import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox, font
import threading
import socket
import json
import time
import logging
import os
from datetime import datetime
import re

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AdvancedLLMGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("🚀 Distributed LLM System - Advanced")
        self.root.geometry("1000x700")
        
        # Configurazione
        self.coordinator_host = "localhost"
        self.coordinator_port = 8765
        self.socket = None
        self.connected = False
        self.request_counter = 0
        self.chat_history = []
        self.session_file = "chat_session.json"
        
        # Carica sessione precedente
        self._load_session()
        
        # Setup interfaccia
        self._setup_styles()
        self._create_widgets()
        self._connect_to_coordinator()
    
    def _setup_styles(self):
        """Configura stili moderni"""
        self.style = ttk.Style()
        self.style.configure('TFrame', background='#f0f0f0')
        self.style.configure('TLabel', background='#f0f0f0', font=('Arial', 10))
        self.style.configure('TButton', font=('Arial', 10))
        self.style.configure('Title.TLabel', font=('Arial', 16, 'bold'), foreground='#2c3e50')
        
        # Font per codice
        self.code_font = font.Font(family="Consolas", size=10)
        self.normal_font = font.Font(family="Arial", size=10)
    
    def _create_widgets(self):
        """Crea l'interfaccia avanzata"""
        # Frame principale con notebook (tabs)
        self.notebook = ttk.Notebook(self.root)
        self.notebook.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # Tab Chat
        self.chat_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.chat_frame, text="💬 Chat")
        
        # Tab Sistema
        self.system_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.system_frame, text="⚙️ Sistema")
        
        self._create_chat_tab()
        self._create_system_tab()
        
        # Bind per Ctrl+S per salvare
        self.root.bind('<Control-s>', lambda e: self._save_session())
    
    def _create_chat_tab(self):
        """Crea il tab della chat"""
        # Header
        header_frame = ttk.Frame(self.chat_frame)
        header_frame.pack(fill=tk.X, pady=(0, 10))
        
        title_label = ttk.Label(header_frame, text="Distributed LLM Chat", style='Title.TLabel')
        title_label.pack(side=tk.LEFT)
        
        self.status_var = tk.StringVar(value="🔴 Disconnesso")
        status_label = ttk.Label(header_frame, textvariable=self.status_var, font=('Arial', 10, 'bold'))
        status_label.pack(side=tk.RIGHT)
        
        # Area chat con frame scrollabile
        chat_container = ttk.Frame(self.chat_frame)
        chat_container.pack(fill=tk.BOTH, expand=True)
        
        self.chat_area = scrolledtext.ScrolledText(
            chat_container,
            wrap=tk.WORD,
            width=90,
            height=25,
            state=tk.DISABLED,
            font=self.normal_font,
            bg='#ffffff',
            relief=tk.FLAT,
            borderwidth=1
        )
        self.chat_area.pack(fill=tk.BOTH, expand=True)
        
        # Configura tags per formattazione
        self.chat_area.tag_configure("user", foreground="#2c3e50", font=('Arial', 10, 'bold'))
        self.chat_area.tag_configure("ai", foreground="#27ae60", font=('Arial', 10))
        self.chat_area.tag_configure("system", foreground="#7f8c8d", font=('Arial', 9, 'italic'))
        self.chat_area.tag_configure("code", background="#f8f9fa", font=self.code_font, 
                                   relief=tk.SUNKEN, borderwidth=1)
        self.chat_area.tag_configure("warning", foreground="#e74c3c", font=('Arial', 9, 'bold'))
        
        # Controlli input
        input_frame = ttk.Frame(self.chat_frame)
        input_frame.pack(fill=tk.X, pady=(10, 0))
        
        self.input_entry = ttk.Entry(
            input_frame,
            font=('Arial', 12)
        )
        self.input_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 10))
        self.input_entry.bind('<Return>', lambda e: self._send_message())
        self.input_entry.bind('<Control-Return>', lambda e: self._insert_newline())
        
        button_frame = ttk.Frame(input_frame)
        button_frame.pack(side=tk.RIGHT)
        
        self.send_button = ttk.Button(
            button_frame,
            text="Invia",
            command=self._send_message
        )
        self.send_button.pack(side=tk.LEFT, padx=(0, 5))
        
        clear_button = ttk.Button(
            button_frame,
            text="Pulisci",
            command=self._clear_chat
        )
        clear_button.pack(side=tk.LEFT, padx=(0, 5))
        
        save_button = ttk.Button(
            button_frame,
            text="Salva",
            command=self._save_session
        )
        save_button.pack(side=tk.LEFT)
        
        # Carica chat precedente
        self._display_chat_history()
    
    def _create_system_tab(self):
        """Crea il tab del sistema"""
        # Info connessione
        conn_frame = ttk.LabelFrame(self.system_frame, text="Connessione", padding="10")
        conn_frame.pack(fill=tk.X, pady=(0, 10))
        
        self.conn_text = tk.Text(
            conn_frame,
            wrap=tk.WORD,
            height=4,
            state=tk.DISABLED,
            font=self.normal_font
        )
        self.conn_text.pack(fill=tk.X)
        
        # Info worker
        workers_frame = ttk.LabelFrame(self.system_frame, text="Workers Attivi", padding="10")
        workers_frame.pack(fill=tk.BOTH, expand=True)
        
        self.workers_text = tk.Text(
            workers_frame,
            wrap=tk.WORD,
            state=tk.DISABLED,
            font=self.code_font
        )
        self.workers_text.pack(fill=tk.BOTH, expand=True)
        
        # Pulsanti sistema
        button_frame = ttk.Frame(self.system_frame)
        button_frame.pack(fill=tk.X, pady=(10, 0))
        
        refresh_btn = ttk.Button(
            button_frame,
            text="Aggiorna Sistema",
            command=self._refresh_system_info
        )
        refresh_btn.pack(side=tk.LEFT, padx=(0, 10))
        
        export_btn = ttk.Button(
            button_frame,
            text="Esporta Chat",
            command=self._export_chat
        )
        export_btn.pack(side=tk.LEFT)
    
    def _insert_newline(self):
        """Inserisce nuova riga nell'input"""
        current_text = self.input_entry.get()
        cursor_pos = self.input_entry.index(tk.INSERT)
        self.input_entry.insert(cursor_pos, "\n")
    
    def _format_message(self, text):
        """Formatta il testo con evidenziazione codice"""
        # Pattern per rilevare blocchi di codice
        code_pattern = r'```(.*?)```'
        inline_code_pattern = r'`(.*?)`'
        
        # Sostituisce blocchi di codice
        def format_code_block(match):
            code = match.group(1).strip()
            return f"\n[CODE_BLOCK]\n{code}\n[/CODE_BLOCK]\n"
        
        text = re.sub(code_pattern, format_code_block, text, flags=re.DOTALL)
        text = re.sub(inline_code_pattern, r'[INLINE_CODE]\1[/INLINE_CODE]', text)
        
        return text
    
    def _add_to_chat(self, sender, message, raw_message=None):
        """Aggiunge messaggio formattato alla chat"""
        self.chat_area.config(state=tk.NORMAL)
        
        timestamp = datetime.now().strftime("%H:%M:%S")
        
        if sender == "Tu":
            prefix = f"👤 [{timestamp}] Tu:\n"
            tag = "user"
            emoji = "👤"
        elif sender == "AI":
            prefix = f"🤖 [{timestamp}] AI:\n"
            tag = "ai"
            emoji = "🤖"
        else:
            prefix = f"⚙️ [{timestamp}] Sistema:\n"
            tag = "system"
            emoji = "⚙️"
        
        # Aggiungi alla history
        self.chat_history.append({
            "timestamp": timestamp,
            "sender": sender,
            "message": raw_message or message,
            "emoji": emoji
        })
        
        # Inserisci prefix
        self.chat_area.insert(tk.END, prefix, tag)
        
        # Formatta e inserisci il messaggio
        formatted_message = self._format_message(message)
        lines = formatted_message.split('\n')
        
        for line in lines:
            if line.startswith('[CODE_BLOCK]') and line.endswith('[/CODE_BLOCK]'):
                code = line[12:-13].strip()
                if code:
                    self.chat_area.insert(tk.END, "\n```\n", "code")
                    self.chat_area.insert(tk.END, code + "\n", "code")
                    self.chat_area.insert(tk.END, "```\n", "code")
            elif '[INLINE_CODE]' in line:
                # Gestisce codice inline
                parts = line.split('[INLINE_CODE]')
                for part in parts:
                    if '[/INLINE_CODE]' in part:
                        code, rest = part.split('[/INLINE_CODE]', 1)
                        self.chat_area.insert(tk.END, "`" + code + "`", "code")
                        self.chat_area.insert(tk.END, rest)
                    else:
                        self.chat_area.insert(tk.END, part)
                self.chat_area.insert(tk.END, "\n")
            else:
                self.chat_area.insert(tk.END, line + "\n")
        
        self.chat_area.insert(tk.END, "\n")
        self.chat_area.see(tk.END)
        self.chat_area.config(state=tk.DISABLED)
    
    def _display_chat_history(self):
        """Mostra la cronologia della chat"""
        for item in self.chat_history:
            self._add_to_chat(item["sender"], item["message"])
    
    def _save_session(self):
        """Salva la sessione corrente"""
        try:
            session_data = {
                "timestamp": datetime.now().isoformat(),
                "chat_history": self.chat_history
            }
            
            with open(self.session_file, 'w', encoding='utf-8') as f:
                json.dump(session_data, f, ensure_ascii=False, indent=2)
            
            self._add_to_chat("Sistema", f"Chat salvata in {self.session_file}")
            messagebox.showinfo("Successo", f"Sessione salvata in {self.session_file}")
            
        except Exception as e:
            messagebox.showerror("Errore", f"Errore nel salvataggio: {e}")
    
    def _load_session(self):
        """Carica una sessione precedente"""
        if os.path.exists(self.session_file):
            try:
                with open(self.session_file, 'r', encoding='utf-8') as f:
                    session_data = json.load(f)
                
                self.chat_history = session_data.get("chat_history", [])
                
            except Exception as e:
                logger.error(f"Errore caricamento sessione: {e}")
    
    def _export_chat(self):
        """Esporta la chat in formato testo"""
        try:
            filename = f"chat_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
            with open(filename, 'w', encoding='utf-8') as f:
                for item in self.chat_history:
                    f.write(f"[{item['timestamp']}] {item['sender']}: {item['message']}\n\n")
            
            messagebox.showinfo("Successo", f"Chat esportata in {filename}")
        except Exception as e:
            messagebox.showerror("Errore", f"Errore nell'esportazione: {e}")
    
    def _clear_chat(self):
        """Pulisce la chat corrente"""
        if messagebox.askyesno("Conferma", "Vuoi pulire la chat?"):
            self.chat_area.config(state=tk.NORMAL)
            self.chat_area.delete(1.0, tk.END)
            self.chat_area.config(state=tk.DISABLED)
            self.chat_history.clear()
    
    def _refresh_system_info(self):
        """Aggiorna le informazioni di sistema"""
        # Qui potremmo implementare una richiesta di status al coordinator
        self._update_connection_info("Sistema aggiornato")
    
    def _update_connection_info(self, info):
        """Aggiorna informazioni connessione"""
        self.conn_text.config(state=tk.NORMAL)
        self.conn_text.delete(1.0, tk.END)
        self.conn_text.insert(1.0, f"{datetime.now().strftime('%H:%M:%S')} - {info}")
        self.conn_text.config(state=tk.DISABLED)
    
    def _connect_to_coordinator(self):
        """Connette al coordinator"""
        def connect_thread():
            try:
                self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.socket.connect((self.coordinator_host, self.coordinator_port))
                self.connected = True
                
                self.root.after(0, self._on_connected)
                
                # Thread ricezione
                receive_thread = threading.Thread(target=self._receive_loop, daemon=True)
                receive_thread.start()
                
            except Exception as e:
                self.root.after(0, lambda: self._on_connection_error(str(e)))
        
        threading.Thread(target=connect_thread, daemon=True).start()
    
    def _on_connected(self):
        """Callback connessione riuscita"""
        self.status_var.set("🟢 Connesso")
        self._update_connection_info("Connesso al sistema distribuito!")
        self._add_to_chat("Sistema", "✅ Connesso al coordinator! Worker distribuiti pronti.")
    
    def _on_connection_error(self, error):
        """Callback errore connessione"""
        self.status_var.set("🔴 Errore")
        self._update_connection_info(f"Errore: {error}")
        messagebox.showerror("Errore", f"Impossibile connettersi: {error}")
    
    def _send_message(self):
        """Invia messaggio"""
        message = self.input_entry.get().strip()
        if not message:
            return
        
        if not self.connected:
            messagebox.showerror("Errore", "Non connesso al coordinator")
            return
        
        self.input_entry.delete(0, tk.END)
        self._add_to_chat("Tu", message)
        
        self.request_counter += 1
        request_message = {
            'type': 'inference_request',
            'request_id': f"gui_req_{self.request_counter}",
            'prompt': message,
            'timestamp': time.time()
        }
        
        try:
            self.socket.send(json.dumps(request_message).encode('utf-8'))
            self._update_connection_info(f"Richiesta inviata: {message[:30]}...")
        except Exception as e:
            messagebox.showerror("Errore", f"Errore nell'invio: {e}")
            self.connected = False
    
    def _receive_loop(self):
        """Loop ricezione"""
        while self.connected:
            try:
                data = self.socket.recv(16384).decode('utf-8')
                if not data:
                    break
                    
                message = json.loads(data)
                self.root.after(0, lambda: self._handle_message(message))
                    
            except Exception as e:
                if self.connected:
                    logger.error(f"Errore ricezione: {e}")
                break
    
    def _handle_message(self, message):
        """Gestisce messaggi ricevuti"""
        msg_type = message.get('type')
        
        if msg_type == 'inference_response':
            result = message['result']
            worker_id = message.get('worker_id', 'AI')
            
            self._add_to_chat("AI", result, raw_message=result)
            self._update_connection_info(
                f"Risposta da {worker_id} - {len(result)} caratteri"
            )
            
        elif msg_type == 'error':
            error_msg = message.get('error', 'Errore sconosciuto')
            self._add_to_chat("Sistema", f"❌ Errore: {error_msg}")
    
    def on_closing(self):
        """Pulizia alla chiusura"""
        self.connected = False
        if self.socket:
            self.socket.close()
        self._save_session()  # Salva automaticamente alla chiusura
        self.root.destroy()

if __name__ == "__main__":
    root = tk.Tk()
    app = AdvancedLLMGUI(root)
    root.protocol("WM_DELETE_WINDOW", app.on_closing)
    root.mainloop()