import maya.cmds as cmds
import os

# Define the script content
auto_jump_script = '''# auto_jump.py
import maya.cmds as cmds

# Global variables
auto_jump_enabled = False
script_job_id = -1

def toggle_auto_jump(*args):
    """Toggle the auto jump to selected key function"""
    global auto_jump_enabled, script_job_id
    try:
        # Find the AutoJump button dynamically
        shelf_top_level = cmds.shelfTabLayout('ShelfLayout', query=True, childArray=True)
        toggle_button_name = None
        for shelf in shelf_top_level:
            buttons = cmds.shelfLayout(shelf, query=True, childArray=True) or []
            for button in buttons:
                if cmds.shelfButton(button, query=True, exists=True):
                    if cmds.shelfButton(button, query=True, label=True) == "AutoJump":
                        toggle_button_name = button
                        break
            if toggle_button_name:
                break

        if auto_jump_enabled:
            if cmds.scriptJob(exists=script_job_id):
                cmds.scriptJob(kill=script_job_id)
                print(f"Script job {script_job_id} killed")
            auto_jump_enabled = False
            if toggle_button_name:
                cmds.shelfButton(toggle_button_name, edit=True, backgroundColor=(0, 0, 0))  # Default (OFF)
            print("Auto jump to selected key: OFF")
        else:
            script_job_id = cmds.scriptJob(event=["SelectionChanged", "import auto_jump; auto_jump.auto_jump_to_key()"])
            auto_jump_enabled = True
            if toggle_button_name:
                cmds.shelfButton(toggle_button_name, edit=True, backgroundColor=(0, 0.5, 0))  # Green (ON)
            print(f"Auto jump to selected key: ON - Script job ID: {script_job_id}")
    except Exception as e:
        print(f"Error in toggle_auto_jump: {e}")

def auto_jump_to_key(*args):
    """Jump Graph Editor to selected keyframe"""
    global auto_jump_enabled
    if not auto_jump_enabled:
        print("Auto jump not enabled")
        return
    try:
        selected_keys = cmds.keyframe(query=True, selected=True, timeChange=True)
        if selected_keys:
            target_time = selected_keys[0]
            look_at_range = f"{target_time-0.1}:{target_time+0.1}"
            graph_editor = "graphEditor1GraphEd"
            if cmds.animCurveEditor(graph_editor, exists=True) and cmds.control(graph_editor, query=True, visible=True):
                cmds.animCurveEditor(graph_editor, edit=True, lookAt=look_at_range)
                cmds.currentTime(target_time)
            else:
                print("Graph Editor not found or not visible. Please open the Graph Editor.")
        else:
            print("No keyframes selected in Graph Editor")
    except Exception as e:
        print(f"Error in auto_jump_to_key: {e}")

def add_button_to_shelf():
    """Add Toggle button to the current shelf"""
    try:
        current_shelf = cmds.shelfTabLayout('ShelfLayout', query=True, selectTab=True)
        if not current_shelf:
            print("No active shelf found! Using first available shelf.")
            shelf_top_level = cmds.shelfTabLayout('ShelfLayout', query=True, childArray=True)
            current_shelf = shelf_top_level[0] if shelf_top_level else None
        
        if not current_shelf:
            print("No shelf available to add button!")
            return

        # Remove existing AutoJump button
        buttons = cmds.shelfLayout(current_shelf, query=True, childArray=True) or []
        for button in buttons:
            if cmds.shelfButton(button, query=True, exists=True):
                if cmds.shelfButton(button, query=True, label=True) == "AutoJump":
                    cmds.deleteUI(button)
                    print(f"Removed existing AutoJump button: {button}")

        # Add Toggle button
        toggle_button_name = cmds.shelfButton(
            parent=current_shelf,
            label='AutoJump',
            command='import auto_jump; auto_jump.toggle_auto_jump()',
            sourceType='python',
            annotation='Toggle auto jump to selected key',
            image1='commandButton.png'
        )
        print(f"AutoJump button created: {toggle_button_name}")
    except Exception as e:
        print(f"Error adding button to shelf: {e}")

if __name__ == "__main__":
    add_button_to_shelf()
    print("Script loaded. Use 'AutoJump' to toggle (Green = ON, Default = OFF).")
'''

# Get the Maya scripts directory
scripts_dir = cmds.internalVar(userScriptDir=True)
if not os.path.exists(scripts_dir):
    os.makedirs(scripts_dir)  # Create the scripts folder if it doesn’t exist

# Write the script to Documents/maya/scripts/auto_jump.py
script_path = os.path.join(scripts_dir, "auto_jump.py")
try:
    with open(script_path, "w") as f:
        f.write(auto_jump_script)
    print(f"Successfully wrote auto_jump.py to: {script_path}")
except Exception as e:
    print(f"Error writing auto_jump.py: {e}")
    raise

# Import and run the newly created script
try:
    import auto_jump
    auto_jump.add_button_to_shelf()
except ImportError:
    # If import fails, add the scripts directory to sys.path and try again
    import sys
    if scripts_dir not in sys.path:
        sys.path.append(scripts_dir)
        print(f"Added {scripts_dir} to sys.path")
    import auto_jump
    auto_jump.add_button_to_shelf()