#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Key Offset Tool
Made by Martin Chang, 2021
Modified for multi-version Maya compatibility, 2025
"""

import maya.cmds as cmds

# Global variables
bakeLayer = True
keyOffsetField = None
startOffsetField = None
endOffsetField = None

def offsetKeys(forward, *args):
    """Offset keys for selected objects"""
    global keyOffsetField
    
    objects = cmds.ls(orderedSelection=True, long=False)
    
    if not keyOffsetField or not cmds.floatField(keyOffsetField, exists=True):
        cmds.warning("Key Offset Tool window not found. Please open the tool first.")
        return
    
    try:
        keyOffset = cmds.floatField(keyOffsetField, query=True, value=True)
    except RuntimeError:
        cmds.warning("Failed to get keyframe offset value.")
        return
    
    if not objects:
        cmds.confirmDialog(title='Error', message='Select one or more objects to convert to offset keys')
        return
    
    if forward:
        for i in range(len(objects)):
            try:
                cmds.keyframe(objects[i], edit=True, relative=True, timeChange=keyOffset*(i+1))
            except RuntimeError as e:
                cmds.warning("Failed to offset keys for {}: {}".format(objects[i], str(e)))
    else:
        for i in range(len(objects)):
            try:
                startTime = cmds.playbackOptions(q=True, minTime=True)
                firstKey = cmds.findKeyframe(objects[i], which='first')
                cmds.keyframe(objects[i], edit=True, relative=True, timeChange=-(firstKey-startTime))
            except RuntimeError as e:
                cmds.warning("Failed to revert offset for {}: {}".format(objects[i], str(e)))
        
def scaleToStart(*args):
    """Scale keys to start of time range"""
    global startOffsetField
    
    objects = cmds.ls(orderedSelection=True, long=False)
    
    if not startOffsetField or not cmds.floatField(startOffsetField, exists=True):
        cmds.warning("Key Offset Tool window not found. Please open the tool first.")
        return
    
    try:
        startOffset = cmds.floatField(startOffsetField, query=True, value=True)
        startTime = cmds.playbackOptions(q=True, minTime=True)
    except RuntimeError:
        cmds.warning("Failed to get offset or time range values.")
        return

    if not objects:
        cmds.confirmDialog(title='Error', message='Select one or more objects to scale keys to start')
        return
    
    for i in range(len(objects)):
        try:
            firstKey = cmds.findKeyframe(objects[i], which='first')
            if startOffset == 0:
                cmds.warning("Offset from Start cannot be zero for {}".format(objects[i]))
                continue
            cmds.selectKey(objects[i], time=(firstKey, firstKey+startOffset))
            cmds.scaleKey(timePivot=firstKey+startOffset, timeScale=(firstKey-startTime)/startOffset + 1)
            if firstKey+startOffset <= startTime:
                cmds.warning('The desired keyframe to scale from is in front of the beginning of the time range, which will cause scaling issues')
        except (RuntimeError, ZeroDivisionError) as e:
            cmds.warning("Failed to scale keys for {}: {}".format(objects[i], str(e)))
        
def scaleToEnd(*args):
    """Scale keys to end of time range"""
    global endOffsetField
    
    objects = cmds.ls(orderedSelection=True, long=False)
    
    if not endOffsetField or not cmds.floatField(endOffsetField, exists=True):
        cmds.warning("Key Offset Tool window not found. Please open the tool first.")
        return
    
    try:
        endOffset = cmds.floatField(endOffsetField, query=True, value=True)
        endTime = cmds.playbackOptions(q=True, maxTime=True)
    except RuntimeError:
        cmds.warning("Failed to get offset or time range values.")
        return

    if not objects:
        cmds.confirmDialog(title='Error', message='Select one or more objects to scale keys to end')
        return
    
    for i in range(len(objects)):
        try:
            lastKey = cmds.findKeyframe(objects[i], which='last')
            if endOffset == 0:
                cmds.warning("Offset from End cannot be zero for {}".format(objects[i]))
                continue
            cmds.selectKey(objects[i], time=(lastKey, lastKey-endOffset))
            cmds.scaleKey(timePivot=lastKey-endOffset, timeScale=1-((lastKey-endTime)/endOffset))
            if lastKey-endOffset >= endTime:
                cmds.warning('The desired keyframe to scale from is behind the end of the time range, which will cause scaling issues')
        except (RuntimeError, ZeroDivisionError) as e:
            cmds.warning("Failed to scale keys for {}: {}".format(objects[i], str(e)))


def show():
    """Show the Key Offset Tool window"""
    global keyOffsetField, startOffsetField, endOffsetField
    
    # Delete existing window if it exists
    if cmds.window("keyWin", exists=True):
        cmds.deleteUI("keyWin")
    
    # Create window
    physWindow = cmds.window(
        "keyWin",
        title="Key Offset Tool",
        resizeToFitChildren=True,
        sizeable=False,
        minimizeButton=False,
        maximizeButton=False,
        menuBar=False
    )
    cmds.window(physWindow, edit=True, h=100, w=100)

    cmds.columnLayout(columnAttach=('both', 0), columnWidth=190)

    cmds.text(label='Offset Keys:', height=30)

    cmds.setParent('..')
    cmds.rowLayout(numberOfColumns=2)
    cmds.text(label='Keyframe Offset:', width=120)
    keyOffsetField = cmds.floatField(value=0.5, precision=2)

    cmds.setParent('..')
    cmds.rowLayout(numberOfColumns=2, height=30)
    cmds.button(label='Offset', width=75, command=lambda *args: offsetKeys(1))
    cmds.button(label='RevertOffset', width=75, command=lambda *args: offsetKeys(0))

    cmds.setParent('..')
    cmds.text(label='Scale Keys to Time Range:', height=30)

    cmds.setParent('..')
    cmds.rowLayout(numberOfColumns=2)
    cmds.text(label='Offset from Start:', width=120)
    startOffsetField = cmds.floatField(value=10, precision=1)

    cmds.setParent('..')
    cmds.button(label='Scale To Start', width=75, command=scaleToStart)

    cmds.setParent('..')
    cmds.rowLayout(numberOfColumns=2)
    cmds.text(label='Offset from End:', width=120)
    endOffsetField = cmds.floatField(value=10, precision=1)

    cmds.setParent('..')
    cmds.button(label='Scale To End', width=75, command=scaleToEnd)

    cmds.showWindow(physWindow)


if __name__ == '__main__':
    show()