import sys
import os
import subprocess
try:
    from colorama import init, Fore, Style
    init(autoreset=True)
except ImportError:
    class DummyColor:
        def __getattr__(self, name): return ""
    Fore = DummyColor()
    Style = DummyColor()

ASCII_ART = """
                          @@@@@@@@@@@@@@@@@@@                          
                     @@@@@@@@@@@@@@@@@@@@@@@@@@@@@                     
                 @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                 
               @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@               
            @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@            
          @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@          
         @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@         
       @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@       
      @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@      
     @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@     
    @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@    
   @@@@@@@@@@@@@@@@@@@    @@@@@@@@@@@@@@@@@@%    @@@@@@@@@@@@@@@@@@@   
  @@@@@@@@@@@@@@@@@@         @@@@@@@@@@@@          @@@@@@@@@@@@@@@@@@  
 @@@@@@@@@@@@@@@@@@@                               @@@@@@@@@@@@@@@@@@@ 
 @@@@@@@@@@@@@@@@@@@                               @@@@@@@@@@@@@@@@@@@ 
@@@@@@@@@@@@@@@@@@@@                               @@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@                               @@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@                              @@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@                              .@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@        =@@         @@         @@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@        @@@        =@@@        @@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@                               @@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@                               @@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@              @@@             -@@@@@@@@@@@@@@@@@@@@
 @@@@@@@@@@@@@@@@@@@@@                            @@@@@@@@@@@@@@@@@@@@ 
 @@@@@@@@@@@@@@@@@@@@@@                         @@@@@@@@@@@@@@@@@@@@@@ 
  @@@@@@@@@@@@@@@@@@@@@@@@                   @@@@@@@@@@@@@@@@@@@@@@@@  
   @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@   
    @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@    
     @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@     
      @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@      
       @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@       
         @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@         
          @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@          
            @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@            
               @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@               
                 @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                 
                     @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                    
                         @@@@@@@@@@@@@@@@@@@@@                         
"""

def print_info():
    print(Fore.CYAN + ASCII_ART)
    print(Fore.YELLOW + "===================================================")
    print(Fore.GREEN + "  KomaruScriptV2 [ Beta Full ]")
    print(Fore.WHITE + "  Version: 2.20.1")
    print(Fore.WHITE + "  Authors:")
    print(Fore.MAGENTA + "    Telegram Channel: @KomaruCharge")
    print(Fore.MAGENTA + "    Telegram Author: @KomaruEat")
    print(Fore.YELLOW + "===================================================")

def print_menu():
    print(Fore.CYAN + "+-------------------------------------------------+")
    print(Fore.CYAN + "|                 " + Fore.GREEN + "KOMARU MENU" + Fore.CYAN + "                     |")
    print(Fore.CYAN + "+-------------------------------------------------+")
    print(Fore.CYAN + "| " + Fore.WHITE + "1. Run a .ks file" + Fore.CYAN + "                               |")
    print(Fore.CYAN + "| " + Fore.WHITE + "2. Compile .ks to .exe (PyInstaller required)" + Fore.CYAN + "   |")
    print(Fore.CYAN + "| " + Fore.WHITE + "3. Install pip packages" + Fore.CYAN + "                         |")
    print(Fore.CYAN + "| " + Fore.WHITE + "4. Create new .ks project" + Fore.CYAN + "                       |")
    print(Fore.CYAN + "| " + Fore.WHITE + "5. System Info (neofetch style)" + Fore.CYAN + "                 |")
    print(Fore.CYAN + "| " + Fore.WHITE + "6. Fix file icon association for Windows (.ks)" + Fore.CYAN + "  |")
    print(Fore.CYAN + "| " + Fore.WHITE + "0. Exit" + Fore.CYAN + "                                         |")
    print(Fore.CYAN + "+-------------------------------------------------+")

def handle_menu():
    while True:
        print_menu()
        choice = input(Fore.YELLOW + "Select an option > " + Fore.WHITE)
        if choice == '1':
            file_path = input(Fore.YELLOW + "Enter path to .ks file > " + Fore.WHITE)
            run_file(file_path)
        elif choice == '2':
            file_path = input(Fore.YELLOW + "Enter path to .ks file to compile > " + Fore.WHITE)
            compile_to_exe(file_path)
        elif choice == '3':
            package = input(Fore.YELLOW + "Enter pip package name to install > " + Fore.WHITE)
            subprocess.run([sys.executable, "-m", "pip", "install", package])
        elif choice == '4':
            filename = input(Fore.YELLOW + "Enter new project name (without extension) > " + Fore.WHITE)
            create_project(filename)
        elif choice == '5':
            print_info()
        elif choice == '6':
            associate_icon()
        elif choice == '0':
            print(Fore.GREEN + "Goodbye!")
            break
        else:
            print(Fore.RED + "Invalid choice. Please try again.")

def create_project(filename):
    if not filename.endswith('.ks'):
        filename += '.ks'
    with open(filename, 'w', encoding='utf-8') as f:
        f.write("# Welcome to KomaruScriptV2!\n")
        f.write("# Pre-installed: flask, pygame, pyTelegramBotAPI, pyautogui, sherlock, and more!\n\n")
        f.write("print('Hello from KomaruScriptV2!')\n")
    print(Fore.GREEN + f"[*] Created {filename} successfully. You can run it via 'komaru {filename}'")

def run_file(filepath):
    if not os.path.exists(filepath):
        print(Fore.RED + f"[!] File not found: {filepath}")
        return
    with open(filepath, 'r', encoding='utf-8') as f:
        code = f.read()
    
    print(Fore.CYAN + f"[*] Running {filepath}...\n")
    try:
        exec(code, {'__name__': '__main__', '__file__': filepath})
    except Exception as e:
        print(Fore.RED + f"[!] Error executing file:\n{e}")

def compile_to_exe(filepath):
    if not os.path.exists(filepath):
        print(Fore.RED + f"[!] File not found: {filepath}")
        return
    print(Fore.CYAN + "[*] Compiling to EXE... (requires PyInstaller)")
    try:
        import PyInstaller.__main__
        PyInstaller.__main__.run([
            filepath,
            '--onefile',
            '--console'
        ])
        print(Fore.GREEN + "[*] Success! Check the 'dist' folder.")
    except ImportError:
        print(Fore.RED + "[!] PyInstaller is not installed. Installing it now...")
        subprocess.run([sys.executable, "-m", "pip", "install", "pyinstaller"])
        print(Fore.GREEN + "[*] PyInstaller installed. Try compiling again.")

def associate_icon():
    print(Fore.CYAN + "[*] Setting up file association for .ks files on Windows...")
    if os.name != "nt":
        print(Fore.RED + "[!] This feature is only available on Windows.")
        return
    try:
        import winreg
        from PIL import Image
        
        script_dir = os.path.dirname(os.path.abspath(__file__))
        png_path = os.path.join(script_dir, "logo.png")
        ico_path = os.path.join(script_dir, "komaru.ico")
        
        if os.path.exists(png_path) and not os.path.exists(ico_path):
            print(Fore.CYAN + "[*] Generating .ico file from logo...")
            img = Image.open(png_path)
            img.save(ico_path, format="ICO", sizes=[(256, 256)])
        
        if not os.path.exists(ico_path):
            print(Fore.RED + "[!] Warning: Icon file 'logo.png' or 'komaru.ico' not found in package directory.")
            # We'll set the association anyway, but without the custom icon if it fails
        
        # Windows Registry paths
        ext_key_path = r"Software\Classes\.ks"
        winreg.CreateKey(winreg.HKEY_CURRENT_USER, ext_key_path)
        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, ext_key_path, 0, winreg.KEY_WRITE)
        winreg.SetValue(key, "", winreg.REG_SZ, "KomaruScriptV2File")
        winreg.CloseKey(key)

        class_key_path = r"Software\Classes\KomaruScriptV2File"
        winreg.CreateKey(winreg.HKEY_CURRENT_USER, class_key_path)
        class_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, class_key_path, 0, winreg.KEY_WRITE)
        winreg.SetValue(class_key, "", winreg.REG_SZ, "KomaruScriptV2 File")

        # Icon
        default_icon_path = r"Software\Classes\KomaruScriptV2File\DefaultIcon"
        winreg.CreateKey(winreg.HKEY_CURRENT_USER, default_icon_path)
        icon_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, default_icon_path, 0, winreg.KEY_WRITE)
        if os.path.exists(ico_path):
            winreg.SetValue(icon_key, "", winreg.REG_SZ, ico_path)
        else:
            winreg.SetValue(icon_key, "", winreg.REG_SZ, sys.executable + ",0")
        
        # Command (Run via python -m komaru <file>)
        cmd_path = r"Software\Classes\KomaruScriptV2File\shell\open\command"
        winreg.CreateKey(winreg.HKEY_CURRENT_USER, cmd_path)
        cmd_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, cmd_path, 0, winreg.KEY_WRITE)
        
        command_string = f'"{sys.executable}" -m komaru.cli "%1"'
        winreg.SetValue(cmd_key, "", winreg.REG_SZ, command_string)
        
        print(Fore.GREEN + "[*] Successfully linked .ks files to KomaruScriptV2!")
        print(Fore.YELLOW + "[!] You may need to restart your terminal or file explorer to see the new icon.")
    except Exception as e:
        print(Fore.RED + f"[!] Failed to set registry keys: {e}")

def run_direct(args):
    filepath = args[0]
    run_file(filepath)

def main():
    if len(sys.argv) < 2:
        print(Fore.RED + "KomaruScriptV2 [ Full ]")
        print(Fore.WHITE + "Usage: komaru <command> [args]")
        print("Commands:")
        print("  menu    - Open interactive menu")
        print("  info    - Show system info")
        print("  <file>  - Run a .ks file directly")
        return

    command = sys.argv[1].lower()

    if command == "info":
        print_info()
    elif command == "menu":
        handle_menu()
    else:
        run_direct(sys.argv[1:])

if __name__ == "__main__":
    main()
