pimport sys
import os
import winreg as reg

# --- Configuration ---
# The text that will appear in the right-click menu
MENU_ITEM_NAME = "Open AI Assistant Here"

# --- Main Logic ---
def install():
    try:
        # Get the absolute path to the python executable and the main script
        python_exe = sys.executable
        script_path = os.path.abspath("local_agent_ui.py")

        # The command that will be executed when the menu item is clicked
        # %V is a placeholder that Windows replaces with the directory you right-clicked in
        command = f'"{python_exe}" "{script_path}" "%V"'

        # Registry path for the context menu on the background of a folder
        key_path = r'Directory\\Background\\shell'

        # Create the main key for our menu item
        with reg.CreateKey(reg.HKEY_CLASSES_ROOT, f'{key_path}\\{MENU_ITEM_NAME}') as key:
            # Create the 'command' subkey and set its value to our command
            with reg.CreateKey(key, 'command') as command_key:
                reg.SetValue(command_key, None, reg.REG_SZ, command)

        print(f"Successfully installed '{MENU_ITEM_NAME}' context menu item.")
        print("You can now right-click inside any folder to launch the assistant with that folder as context.")

    except Exception as e:
        print(f"Error: {e}")
        print("Please try running this script as an administrator.")

def uninstall():
    try:
        key_path = r'Directory\\Background\\shell'
        reg.DeleteKey(reg.HKEY_CLASSES_ROOT, f'{key_path}\\{MENU_ITEM_NAME}\\command')
        reg.DeleteKey(reg.HKEY_CLASSES_ROOT, f'{key_path}\\{MENU_ITEM_NAME}')
        print(f"Successfully uninstalled '{MENU_ITEM_NAME}'.")
    except FileNotFoundError:
        print("Menu item not found. Nothing to uninstall.")
    except Exception as e:
        print(f"Error: {e}")
        print("Please try running this script as an administrator.")


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == 'uninstall':
        uninstall()
    else:
        install()