#!/usr/bin/python3

import gi
gi.require_version('Gtk', '4.0')
gi.require_version('Adw', '1')

from gi.repository import Gtk, Adw, Gio, GLib
import sys
import subprocess
import json
import os
import threading
import time
import re

class ComfyUILauncher(Gtk.ApplicationWindow):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
        # Configurazione finestra principale
        self.set_title("ComfyUI Launcher")
        self.set_default_size(900, 700)
        
        # Variabili di stato
        self.comfyui_process = None
        self.is_running = False
        self.selected_conda_env = None
        self.conda_envs = []
        self.status_check_timeout = None
        
        # Crea il layout principale
        self.setup_ui()
        
        # Inizializza il controllo degli ambienti conda
        self.load_conda_environments()
        
        # Avvia il monitoraggio dello stato
        self.start_status_monitoring()
        
    def setup_ui(self):
        # Header bar
        header = Gtk.HeaderBar()
        header.set_title_widget(Gtk.Label(label="ComfyUI Launcher"))
        self.set_titlebar(header)
        
        # Menu button
        menu_button = Gtk.MenuButton()
        menu_button.set_icon_name("open-menu-symbolic")
        header.pack_end(menu_button)
        
        # Main content area
        main_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12)
        main_box.set_margin_top(12)
        main_box.set_margin_bottom(12)
        main_box.set_margin_start(12)
        main_box.set_margin_end(12)
        
        # Title
        title_label = Gtk.Label()
        title_label.set_markup("<span size='x-large'><b>ComfyUI Launcher</b></span>")
        title_label.set_margin_bottom(12)
        main_box.append(title_label)
        
        # Status section
        status_frame = Gtk.Frame()
        status_frame.set_label("Stato ComfyUI")
        status_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        status_box.set_margin_top(12)
        status_box.set_margin_bottom(12)
        status_box.set_margin_start(12)
        status_box.set_margin_end(12)
        
        self.status_indicator = Gtk.Image()
        self.status_indicator.set_from_icon_name("media-playback-stop-symbolic")
        self.status_indicator.set_icon_size(Gtk.IconSize.LARGE)
        status_box.append(self.status_indicator)
        
        self.status_label = Gtk.Label(label="ComfyUI non in esecuzione")
        self.status_label.set_hexpand(True)
        self.status_label.set_halign(Gtk.Align.START)
        status_box.append(self.status_label)
        
        status_frame.set_child(status_box)
        main_box.append(status_frame)
        
        # Environment selection
        env_frame = Gtk.Frame()
        env_frame.set_label("Selezione Ambiente Conda")
        env_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
        env_box.set_margin_top(12)
        env_box.set_margin_bottom(12)
        env_box.set_margin_start(12)
        env_box.set_margin_end(12)
        
        # Dropdown per ambienti conda
        self.env_dropdown = Gtk.DropDown()
        self.env_dropdown.set_enable_search(True)
        self.env_dropdown.connect("notify::selected-item", self.on_env_selected)
        env_box.append(self.env_dropdown)
        
        # Refresh button
        refresh_button = Gtk.Button(label="Aggiorna Lista Ambienti")
        refresh_button.connect("clicked", self.on_refresh_envs)
        env_box.append(refresh_button)
        
        env_frame.set_child(env_box)
        main_box.append(env_frame)
        
        # Control buttons
        control_frame = Gtk.Frame()
        control_frame.set_label("Controlli")
        control_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        control_box.set_margin_top(12)
        control_box.set_margin_bottom(12)
        control_box.set_margin_start(12)
        control_box.set_margin_end(12)
        control_box.set_homogeneous(True)
        
        self.start_button = Gtk.Button(label="Avvia ComfyUI")
        self.start_button.add_css_class("suggested-action")
        self.start_button.connect("clicked", self.on_start_comfyui)
        self.start_button.set_sensitive(False)
        control_box.append(self.start_button)
        
        self.stop_button = Gtk.Button(label="Ferma ComfyUI")
        self.stop_button.add_css_class("destructive-action")
        self.stop_button.connect("clicked", self.on_stop_comfyui)
        self.stop_button.set_sensitive(False)
        control_box.append(self.stop_button)
        
        self.install_button = Gtk.Button(label="Installa ComfyUI")
        self.install_button.connect("clicked", self.on_install_comfyui)
        control_box.append(self.install_button)
        
        control_frame.set_child(control_box)
        main_box.append(control_frame)
        
        # Log section
        log_frame = Gtk.Frame()
        log_frame.set_label("Log")
        log_frame.set_vexpand(True)
        
        scrolled = Gtk.ScrolledWindow()
        scrolled.set_policy(Gtk.PolicyType.AUTOMATIC, Gtk.PolicyType.AUTOMATIC)
        scrolled.set_vexpand(True)
        
        self.log_view = Gtk.TextView()
        self.log_view.set_editable(False)
        self.log_view.set_monospace(True)
        self.log_buffer = self.log_view.get_buffer()
        scrolled.set_child(self.log_view)
        
        log_frame.set_child(scrolled)
        main_box.append(log_frame)
        
        self.set_child(main_box)
        
        # Setup menu
        self.setup_menu(menu_button)
        
    def setup_menu(self, menu_button):
        menu = Gio.Menu()
        menu.append("Apri Cartella ComfyUI", "app.open_folder")
        menu.append("About", "app.about")
        menu.append("Quit", "app.quit")
        menu_button.set_menu_model(menu)
        
    def log_message(self, message):
        """Aggiunge un messaggio al log"""
        def update_log():
            current_text = self.log_buffer.get_text(
                self.log_buffer.get_start_iter(),
                self.log_buffer.get_end_iter(),
                False
            )
            timestamp = time.strftime("%H:%M:%S")
            new_text = f"{current_text}[{timestamp}] {message}\n" if current_text else f"[{timestamp}] {message}\n"
            self.log_buffer.set_text(new_text)
            
            # Scorri alla fine
            mark = self.log_buffer.get_insert()
            self.log_view.scroll_mark_onscreen(mark)
            
        GLib.idle_add(update_log)
        
    def load_conda_environments(self):
        """Carica la lista degli ambienti conda disponibili"""
        def load_envs():
            try:
                # Prova a ottenere la lista degli ambienti conda
                result = subprocess.run(['conda', 'env', 'list', '--json'], 
                                      capture_output=True, text=True, timeout=30)
                if result.returncode == 0:
                    env_data = json.loads(result.stdout)
                    envs = []
                    for env_path in env_data['envs']:
                        env_name = os.path.basename(env_path)
                        if env_name == env_path:  # È il path base, usa 'base'
                            env_name = 'base'
                        envs.append((env_name, env_path))
                    
                    def update_ui():
                        self.conda_envs = envs
                        string_list = Gtk.StringList()
                        for name, path in envs:
                            string_list.append(f"{name} ({path})")
                        self.env_dropdown.set_model(string_list)
                        if envs:
                            self.env_dropdown.set_selected(0)
                        self.log_message(f"Trovati {len(envs)} ambienti conda")
                        
                    GLib.idle_add(update_ui)
                else:
                    GLib.idle_add(lambda: self.log_message(f"Errore nel caricare ambienti conda: {result.stderr}"))
            except subprocess.TimeoutExpired:
                GLib.idle_add(lambda: self.log_message("Timeout nel caricamento ambienti conda"))
            except FileNotFoundError:
                GLib.idle_add(lambda: self.log_message("Conda non trovato. Assicurati che conda sia installato e nel PATH"))
            except Exception as e:
                GLib.idle_add(lambda: self.log_message(f"Errore imprevisto: {str(e)}"))
                
        # Esegui in un thread separato per non bloccare l'UI
        threading.Thread(target=load_envs, daemon=True).start()
        
    def on_env_selected(self, dropdown, param):
        """Gestisce la selezione di un ambiente conda"""
        selected_idx = dropdown.get_selected()
        if selected_idx != Gtk.INVALID_LIST_POSITION and self.conda_envs:
            env_name, env_path = self.conda_envs[selected_idx]
            self.selected_conda_env = (env_name, env_path)
            self.log_message(f"Selezionato ambiente: {env_name}")
            self.update_button_states()
            
    def on_refresh_envs(self, button):
        """Aggiorna la lista degli ambienti conda"""
        self.log_message("Aggiornamento lista ambienti...")
        self.load_conda_environments()
        
    def update_button_states(self):
        """Aggiorna lo stato dei pulsanti in base allo stato corrente"""
        has_env = self.selected_conda_env is not None
        self.start_button.set_sensitive(has_env and not self.is_running)
        self.stop_button.set_sensitive(self.is_running)
        self.install_button.set_sensitive(has_env)
        
    def check_comfyui_running(self):
        """Controlla se ComfyUI è in esecuzione"""
        try:
            # Controlla se il processo ComfyUI è ancora vivo
            if self.comfyui_process and self.comfyui_process.poll() is None:
                return True
                
            # Controlla se c'è un processo ComfyUI in esecuzione (porta 8188)
            try:
                result = subprocess.run(['ss', '-tulpn'], capture_output=True, text=True, timeout=5)
                if ':8188' in result.stdout:
                    return True
            except:
                # Fallback con netstat se ss non è disponibile
                try:
                    result = subprocess.run(['netstat', '-tulpn'], capture_output=True, text=True, timeout=5)
                    if ':8188' in result.stdout:
                        return True
                except:
                    pass
            
            # Controlla processi con nome 'comfy' o 'python' che potrebbero essere ComfyUI
            try:
                result = subprocess.run(['pgrep', '-f', 'comfy.*launch'], capture_output=True, text=True, timeout=5)
                if result.returncode == 0 and result.stdout.strip():
                    return True
            except:
                pass
                
            return False
        except Exception as e:
            self.log_message(f"Errore nel controllo stato: {str(e)}")
            return False
            
    def update_status_ui(self):
        """Aggiorna l'interfaccia utente con lo stato corrente"""
        if self.is_running:
            self.status_indicator.set_from_icon_name("media-playback-start-symbolic")
            self.status_indicator.add_css_class("success")
            self.status_label.set_text("ComfyUI in esecuzione")
        else:
            self.status_indicator.set_from_icon_name("media-playback-stop-symbolic")
            self.status_indicator.remove_css_class("success")
            self.status_label.set_text("ComfyUI non in esecuzione")
            
        self.update_button_states()
        
    def start_status_monitoring(self):
        """Avvia il monitoraggio periodico dello stato"""
        def monitor():
            was_running = self.is_running
            self.is_running = self.check_comfyui_running()
            
            if was_running != self.is_running:
                if self.is_running:
                    self.log_message("ComfyUI rilevato in esecuzione")
                else:
                    self.log_message("ComfyUI non più in esecuzione")
                GLib.idle_add(self.update_status_ui)
            
            return True  # Continua il monitoring
            
        # Controlla ogni 3 secondi
        self.status_check_timeout = GLib.timeout_add(3000, monitor)
        
    def on_start_comfyui(self, button):
        """Avvia ComfyUI"""
        if not self.selected_conda_env:
            self.log_message("Nessun ambiente conda selezionato")
            return
            
        env_name, env_path = self.selected_conda_env
        self.log_message(f"Avvio ComfyUI nell'ambiente {env_name}...")
        
        def start_comfy():
            try:
                # Comando per attivare l'ambiente conda e avviare comfyui
                if env_name == 'base':
                    cmd = ['conda', 'run', '-n', 'base', 'comfy', 'launch']
                else:
                    cmd = ['conda', 'run', '-n', env_name, 'comfy', 'launch']
                
                self.comfyui_process = subprocess.Popen(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    universal_newlines=True
                )
                
                GLib.idle_add(lambda: self.log_message("Processo ComfyUI avviato"))
                
                # Leggi l'output in tempo reale
                for line in iter(self.comfyui_process.stdout.readline, ''):
                    if line:
                        GLib.idle_add(lambda l=line.strip(): self.log_message(f"ComfyUI: {l}"))
                        
            except Exception as e:
                GLib.idle_add(lambda: self.log_message(f"Errore nell'avvio: {str(e)}"))
                
        threading.Thread(target=start_comfy, daemon=True).start()
        
    def on_stop_comfyui(self, button):
        """Ferma ComfyUI"""
        self.log_message("Fermando ComfyUI...")
        
        if self.comfyui_process:
            try:
                self.comfyui_process.terminate()
                # Aspetta un po' per la terminazione graceful
                try:
                    self.comfyui_process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    # Forza la terminazione
                    self.comfyui_process.kill()
                self.comfyui_process = None
                self.log_message("Processo ComfyUI terminato")
            except Exception as e:
                self.log_message(f"Errore nella terminazione: {str(e)}")
        else:
            # Prova a killare qualsiasi processo sulla porta 8188
            try:
                subprocess.run(['pkill', '-f', 'comfy'], timeout=5)
                self.log_message("Terminati processi ComfyUI")
            except:
                self.log_message("Nessun processo ComfyUI trovato da terminare")
                
    def on_install_comfyui(self, button):
        """Installa ComfyUI nell'ambiente selezionato"""
        if not self.selected_conda_env:
            self.log_message("Nessun ambiente conda selezionato")
            return
            
        env_name, env_path = self.selected_conda_env
        self.log_message(f"Installazione ComfyUI nell'ambiente {env_name}...")
        
        def install_comfy():
            try:
                # Prima installa comfy-cli
                cmd1 = ['conda', 'run', '-n', env_name, 'pip', 'install', 'comfy-cli']
                result1 = subprocess.run(cmd1, capture_output=True, text=True, timeout=300)
                
                if result1.returncode == 0:
                    GLib.idle_add(lambda: self.log_message("comfy-cli installato con successo"))
                    
                    # Poi installa ComfyUI
                    cmd2 = ['conda', 'run', '-n', env_name, 'comfy', 'install']
                    result2 = subprocess.run(cmd2, capture_output=True, text=True, timeout=600)
                    
                    if result2.returncode == 0:
                        GLib.idle_add(lambda: self.log_message("ComfyUI installato con successo!"))
                    else:
                        GLib.idle_add(lambda: self.log_message(f"Errore installazione ComfyUI: {result2.stderr}"))
                else:
                    GLib.idle_add(lambda: self.log_message(f"Errore installazione comfy-cli: {result1.stderr}"))
                    
            except subprocess.TimeoutExpired:
                GLib.idle_add(lambda: self.log_message("Timeout durante l'installazione"))
            except Exception as e:
                GLib.idle_add(lambda: self.log_message(f"Errore durante l'installazione: {str(e)}"))
                
        threading.Thread(target=install_comfy, daemon=True).start()

class ComfyUIApp(Adw.Application):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.connect('activate', self.on_activate)
        
        # Configura il style manager per gestire correttamente i temi
        self.setup_style_manager()
        
        # Aggiungi azioni per il menu
        self.setup_actions()
        
    def setup_style_manager(self):
        """Configura il gestore dello stile per evitare warning"""
        style_manager = Adw.StyleManager.get_default()
        style_manager.set_color_scheme(Adw.ColorScheme.DEFAULT)
        
    def setup_actions(self):
        # Azione About
        about_action = Gio.SimpleAction.new("about", None)
        about_action.connect("activate", self.on_about_action)
        self.add_action(about_action)
        
        # Azione Quit
        quit_action = Gio.SimpleAction.new("quit", None)
        quit_action.connect("activate", self.on_quit_action)
        self.add_action(quit_action)
        
        # Azione Open Folder
        open_folder_action = Gio.SimpleAction.new("open_folder", None)
        open_folder_action.connect("activate", self.on_open_folder_action)
        self.add_action(open_folder_action)
        
        # Shortcut per quit
        self.set_accels_for_action("app.quit", ["<primary>q"])

    def on_activate(self, app):
        print("Avvio ComfyUI Launcher...")
        # Crea la finestra principale
        self.win = ComfyUILauncher(application=app)
        print("Finestra launcher creata")
        self.win.present()
        print("Launcher presentato")
        
    def on_about_action(self, action, param):
        # Mostra dialog About
        about = Adw.AboutWindow(
            transient_for=self.win,
            application_name="ComfyUI Launcher",
            application_icon="application-x-executable",
            developer_name="ComfyUI Launcher",
            version="1.0.0",
            developers=["ComfyUI Launcher Team"],
            copyright="© 2025 ComfyUI Launcher",
            comments="Un launcher GTK per ComfyUI con supporto conda"
        )
        about.present()
        
    def on_open_folder_action(self, action, param):
        """Apre la cartella ComfyUI"""
        try:
            # Prova ad aprire la cartella di default di ComfyUI
            home_dir = os.path.expanduser("~")
            comfy_dir = os.path.join(home_dir, "comfy", "ComfyUI")
            if os.path.exists(comfy_dir):
                subprocess.run(['xdg-open', comfy_dir])
            else:
                self.win.log_message("Cartella ComfyUI non trovata")
        except Exception as e:
            self.win.log_message(f"Errore nell'apertura cartella: {str(e)}")
        
    def on_quit_action(self, action, param):
        # Ferma ComfyUI se in esecuzione prima di chiudere
        if hasattr(self, 'win') and self.win.comfyui_process:
            self.win.on_stop_comfyui(None)
        self.quit()

def main():
    print("Avvio del ComfyUI Launcher...")
    try:
        app = ComfyUIApp(application_id="com.example.ComfyUILauncher")
        print("Applicazione launcher creata, avvio in corso...")
        exit_code = app.run(sys.argv)
        print(f"Launcher terminato con codice: {exit_code}")
        return exit_code
    except Exception as e:
        print(f"Errore durante l'esecuzione del launcher: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == '__main__':
    main()
