#!/usr/bin/env python3
"""
ComfyUI Hugging Face Space
Versão otimizada para execução gratuita no Hugging Face Spaces
"""

import subprocess
import os
import time
import socket
import threading
import logging
import json
import gradio as gr
from pathlib import Path
from typing import Dict, List, Optional, Any
import psutil
import signal
import sys
import requests
from datetime import datetime

# Configuração de logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('/tmp/comfyui.log')
    ]
)
logger = logging.getLogger(__name__)

class ComfyUISpaceManager:
    """Gerenciador principal do ComfyUI no Hugging Face Space"""
    
    def __init__(self):
        self.app_dir = Path("/tmp/ComfyUI")
        self.models_dir = Path("/data/models")  # Persistente no HF Spaces
        self.port = 7860  # Porta padrão do Gradio
        self.comfy_port = 8188  # Porta do ComfyUI
        self.process = None
        self.setup_complete = False
        self.setup_logs = []
        
        # Cria diretórios necessários
        self.app_dir.mkdir(parents=True, exist_ok=True)
        self.models_dir.mkdir(parents=True, exist_ok=True)
        
    def log_setup(self, message: str, level: str = "info"):
        """Adiciona mensagem aos logs de setup"""
        timestamp = datetime.now().strftime("%H:%M:%S")
        log_entry = f"[{timestamp}] {message}"
        self.setup_logs.append(log_entry)
        
        if level == "info":
            logger.info(message)
        elif level == "warning":
            logger.warning(message)
        elif level == "error":
            logger.error(message)
            
        # Mantém apenas últimas 100 entradas
        if len(self.setup_logs) > 100:
            self.setup_logs = self.setup_logs[-100:]
    
    def get_hf_token(self) -> str:
        """Obtém token do Hugging Face"""
        token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_TOKEN")
        if not token:
            self.log_setup("⚠️ Token HF não encontrado, alguns modelos podem não baixar", "warning")
            return ""
        self.log_setup(f"✅ Token HF configurado: ...{token[-8:]}")
        return token
    
    def setup_comfyui(self) -> bool:
        """Setup completo do ComfyUI"""
        try:
            self.log_setup("🚀 Iniciando setup do ComfyUI...")
            
            # 1. Clona repositório
            if not self._clone_repository():
                return False
            
            # 2. Instala dependências
            if not self._install_dependencies():
                return False
            
            # 3. Setup de modelos
            if not self._setup_models():
                return False
            
            # 4. Configura custom nodes
            if not self._setup_custom_nodes():
                return False
            
            self.log_setup("✅ Setup completo! ComfyUI pronto para usar")
            self.setup_complete = True
            return True
            
        except Exception as e:
            self.log_setup(f"❌ Erro no setup: {e}", "error")
            return False
    
    def _clone_repository(self) -> bool:
        """Clona repositório do ComfyUI"""
        if (self.app_dir / ".git").exists():
            self.log_setup("📁 Repositório já existe, atualizando...")
            try:
                os.chdir(self.app_dir)
                subprocess.run(["git", "pull"], check=True, capture_output=True)
                return True
            except subprocess.CalledProcessError as e:
                self.log_setup(f"⚠️ Erro ao atualizar repositório: {e}", "warning")
                # Continua mesmo se falhar
        
        self.log_setup("📥 Clonando ComfyUI...")
        try:
            subprocess.run([
                "git", "clone", "https://github.com/comfyanonymous/ComfyUI.git",
                str(self.app_dir)
            ], check=True, capture_output=True)
            self.log_setup("✅ ComfyUI clonado com sucesso")
            return True
        except subprocess.CalledProcessError as e:
            self.log_setup(f"❌ Erro ao clonar ComfyUI: {e}", "error")
            return False
    
    def _install_dependencies(self) -> bool:
        """Instala dependências do ComfyUI"""
        self.log_setup("📦 Instalando dependências...")
        
        # Lista de dependências essenciais
        dependencies = [
            "torch>=2.0.0",
            "torchvision",
            "torchaudio",
            "numpy<2.0",
            "pillow",
            "transformers>=4.28.1",
            "accelerate",
            "diffusers",
            "huggingface-hub[hf_transfer]",
            "safetensors>=0.4.2",
            "aiohttp",
            "psutil",
            "opencv-python-headless"
        ]
        
        for dep in dependencies:
            try:
                self.log_setup(f"📦 Instalando {dep}...")
                subprocess.run([
                    "pip", "install", "--upgrade", dep
                ], check=True, capture_output=True)
            except subprocess.CalledProcessError as e:
                self.log_setup(f"⚠️ Falha ao instalar {dep}: {e}", "warning")
                # Continua mesmo se falhar
        
        self.log_setup("✅ Dependências instaladas")
        return True
    
    def _setup_models(self) -> bool:
        """Setup de modelos com cache persistente"""
        self.log_setup("🤖 Configurando modelos...")
        
        # Verifica se modelos já existem
        models_exist = self._check_existing_models()
        
        if models_exist:
            self.log_setup("✅ Modelos já existem no cache, vinculando...")
            self._link_models()
            return True
        
        # Download de modelos essenciais
        self.log_setup("📥 Primeira execução - baixando modelos essenciais...")
        self._download_essential_models()
        
        return True
    
    def _check_existing_models(self) -> bool:
        """Verifica se modelos já existem no storage persistente"""
        essential_models = [
            "checkpoints/flux1-schnell.safetensors",
            "vae/ae.safetensors",
            "clip/clip_l.safetensors"
        ]
        
        for model_path in essential_models:
            full_path = self.models_dir / model_path
            if not full_path.exists() or full_path.stat().st_size < 1024*1024:  # < 1MB
                return False
        
        return True
    
    def _link_models(self):
        """Vincula modelos do storage persistente ao ComfyUI"""
        comfy_models = self.app_dir / "models"
        
        # Remove diretório de modelos se existir
        if comfy_models.exists():
            import shutil
            shutil.rmtree(comfy_models)
        
        # Cria symlink para storage persistente
        comfy_models.symlink_to(self.models_dir)
        self.log_setup("🔗 Modelos vinculados ao storage persistente")
    
    def _download_essential_models(self):
        """Download de modelos essenciais"""
        token = self.get_hf_token()
        
        # Modelos essenciais (versões menores para HF Spaces)
        essential_models = [
            {
                "repo": "black-forest-labs/FLUX.1-schnell",
                "file": "flux1-schnell.safetensors",
                "dir": "checkpoints"
            },
            {
                "repo": "black-forest-labs/FLUX.1-schnell", 
                "file": "ae.safetensors",
                "dir": "vae"
            },
            {
                "repo": "openai/clip-vit-large-patch14",
                "file": "pytorch_model.bin",
                "dir": "clip",
                "rename": "clip_l.safetensors"
            }
        ]
        
        for model in essential_models:
            self._download_hf_model(
                model["repo"], 
                model["file"], 
                model["dir"],
                token,
                model.get("rename")
            )
        
        # LoRAs populares (menores)
        self._download_popular_loras()
        
        # Vincula modelos após download
        self._link_models()
    
    def _download_hf_model(self, repo: str, filename: str, target_dir: str, token: str, rename: str = None):
        """Download de modelo individual do HuggingFace"""
        target_path = self.models_dir / target_dir
        target_path.mkdir(parents=True, exist_ok=True)
        
        final_name = rename if rename else filename
        file_path = target_path / final_name
        
        if file_path.exists():
            self.log_setup(f"✅ {final_name} já existe")
            return
        
        self.log_setup(f"📥 Baixando {final_name} de {repo}...")
        
        try:
            from huggingface_hub import hf_hub_download
            
            downloaded_path = hf_hub_download(
                repo_id=repo,
                filename=filename,
                token=token if token else None,
                cache_dir="/tmp/hf_cache"
            )
            
            # Copia para storage persistente
            import shutil
            shutil.copy2(downloaded_path, file_path)
            
            size_mb = file_path.stat().st_size / (1024*1024)
            self.log_setup(f"✅ {final_name} baixado ({size_mb:.1f}MB)")
            
        except Exception as e:
            self.log_setup(f"⚠️ Erro ao baixar {final_name}: {e}", "warning")
    
    def _download_popular_loras(self):
        """Download de LoRAs populares e pequenos"""
        loras_dir = self.models_dir / "loras"
        loras_dir.mkdir(exist_ok=True)
        
        # LoRAs pequenos e populares
        popular_loras = [
            "https://huggingface.co/alvdansen/pony-realism/resolve/main/ponyRealismV21MainVAE.safetensors",
            "https://huggingface.co/Hyper-SD/Hyper-FLUX.1-dev-8steps-lora/resolve/main/Hyper-FLUX.1-dev-8steps-lora.safetensors"
        ]
        
        for url in popular_loras:
            try:
                filename = url.split("/")[-1]
                file_path = loras_dir / filename
                
                if file_path.exists():
                    continue
                
                self.log_setup(f"📥 Baixando LoRA {filename}...")
                subprocess.run([
                    "wget", "-q", "--timeout=60", url, "-O", str(file_path)
                ], check=True)
                
                if file_path.stat().st_size > 1024:  # > 1KB
                    self.log_setup(f"✅ LoRA {filename} baixado")
                else:
                    file_path.unlink()
                    
            except Exception as e:
                self.log_setup(f"⚠️ Erro ao baixar LoRA: {e}", "warning")
    
    def _setup_custom_nodes(self) -> bool:
        """Setup de custom nodes essenciais"""
        self.log_setup("🔧 Configurando custom nodes...")
        
        custom_nodes_dir = self.app_dir / "custom_nodes"
        custom_nodes_dir.mkdir(exist_ok=True)
        
        # Custom nodes essenciais e leves
        essential_nodes = [
            "https://github.com/ltdrdata/ComfyUI-Manager.git",
            "https://github.com/WASasquatch/was-node-suite-comfyui.git"
        ]
        
        for repo_url in essential_nodes:
            try:
                repo_name = repo_url.split("/")[-1].replace(".git", "")
                node_dir = custom_nodes_dir / repo_name
                
                if node_dir.exists():
                    self.log_setup(f"✅ {repo_name} já existe")
                    continue
                
                self.log_setup(f"📥 Clonando {repo_name}...")
                subprocess.run([
                    "git", "clone", "--depth", "1", repo_url, str(node_dir)
                ], check=True, capture_output=True, timeout=120)
                
                # Instala requirements se existir
                req_file = node_dir / "requirements.txt"
                if req_file.exists():
                    subprocess.run([
                        "pip", "install", "-r", str(req_file)
                    ], check=False, capture_output=True, timeout=120)
                
                self.log_setup(f"✅ {repo_name} instalado")
                
            except Exception as e:
                self.log_setup(f"⚠️ Erro ao instalar {repo_name}: {e}", "warning")
        
        self.log_setup("✅ Custom nodes configurados")
        return True
    
    def start_comfyui(self) -> bool:
        """Inicia o servidor ComfyUI"""
        if self.process and self.process.poll() is None:
            self.log_setup("✅ ComfyUI já está rodando")
            return True
        
        self.log_setup("🚀 Iniciando servidor ComfyUI...")
        
        try:
            os.chdir(self.app_dir)
            
            # Inicia ComfyUI em background
            self.process = subprocess.Popen([
                "python", "main.py",
                "--listen", "0.0.0.0",
                "--port", str(self.comfy_port),
                "--preview-method", "auto",
                "--dont-print-server"
            ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            
            # Aguarda servidor ficar pronto
            for attempt in range(60):  # 60 tentativas = 2 minutos
                if self._test_comfyui_connection():
                    self.log_setup("✅ ComfyUI está rodando!")
                    return True
                
                if self.process.poll() is not None:
                    stdout, stderr = self.process.communicate()
                    self.log_setup(f"❌ ComfyUI parou: {stderr}", "error")
                    return False
                
                time.sleep(2)
            
            self.log_setup("❌ Timeout aguardando ComfyUI", "error")
            return False
            
        except Exception as e:
            self.log_setup(f"❌ Erro ao iniciar ComfyUI: {e}", "error")
            return False
    
    def _test_comfyui_connection(self) -> bool:
        """Testa conexão com ComfyUI"""
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
                sock.settimeout(1)
                return sock.connect_ex(('127.0.0.1', self.comfy_port)) == 0
        except:
            return False
    
    def get_status(self) -> Dict[str, Any]:
        """Retorna status do sistema"""
        return {
            "setup_complete": self.setup_complete,
            "comfyui_running": self.process and self.process.poll() is None,
            "comfyui_url": f"http://127.0.0.1:{self.comfy_port}" if self.setup_complete else None,
            "logs": self.setup_logs[-20:],  # Últimas 20 linhas
            "memory_usage": psutil.virtual_memory().percent,
            "disk_usage": psutil.disk_usage('/').percent if os.path.exists('/') else 0,
            "models_dir_size": self._get_dir_size(self.models_dir)
        }
    
    def _get_dir_size(self, path: Path) -> str:
        """Retorna tamanho do diretório formatado"""
        try:
            total_size = sum(f.stat().st_size for f in path.glob('**/*') if f.is_file())
            size_gb = total_size / (1024**3)
            return f"{size_gb:.2f} GB"
        except:
            return "N/A"
    
    def stop_comfyui(self):
        """Para o ComfyUI"""
        if self.process:
            self.log_setup("🛑 Parando ComfyUI...")
            self.process.terminate()
            try:
                self.process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                self.process.kill()
            self.process = None

# Instância global
comfy_manager = ComfyUISpaceManager()

def setup_comfyui():
    """Interface para setup do ComfyUI"""
    if comfy_manager.setup_complete:
        return "✅ Setup já completo!", comfy_manager.get_status()
    
    # Executa setup em thread separada
    def run_setup():
        comfy_manager.setup_comfyui()
        comfy_manager.start_comfyui()
    
    setup_thread = threading.Thread(target=run_setup, daemon=True)
    setup_thread.start()
    
    return "🚀 Setup iniciado! Acompanhe o progresso abaixo...", comfy_manager.get_status()

def get_status_update():
    """Atualização de status para interface"""
    status = comfy_manager.get_status()
    
    # Formata logs
    logs_text = "\n".join(status["logs"]) if status["logs"] else "Aguardando logs..."
    
    # Status geral
    if status["setup_complete"] and status["comfyui_running"]:
        main_status = "🟢 ComfyUI Rodando - Pronto para usar!"
        comfyui_link = f'<a href="http://127.0.0.1:{comfy_manager.comfy_port}" target="_blank">🎯 Abrir ComfyUI</a>'
    elif status["setup_complete"]:
        main_status = "🟡 Setup completo - Iniciando ComfyUI..."
        comfyui_link = "⏳ Aguardando..."
    else:
        main_status = "🟡 Setup em progresso..."
        comfyui_link = "⏳ Aguardando setup..."
    
    # Informações do sistema
    system_info = f"""
    💾 Memória: {status['memory_usage']:.1f}%
    💿 Disco: {status['disk_usage']:.1f}%
    📁 Modelos: {status['models_dir_size']}
    """
    
    return main_status, logs_text, system_info, comfyui_link

def restart_comfyui():
    """Reinicia ComfyUI"""
    comfy_manager.stop_comfyui()
    time.sleep(2)
    success = comfy_manager.start_comfyui()
    
    if success:
        return "✅ ComfyUI reiniciado com sucesso!"
    else:
        return "❌ Erro ao reiniciar ComfyUI"

# Interface Gradio
def create_interface():
    """Cria interface Gradio"""
    
    with gr.Blocks(
        title="ComfyUI Hugging Face Space",
        theme=gr.themes.Soft(),
        css="""
        .status-box { padding: 20px; border-radius: 10px; margin: 10px 0; }
        .running { background-color: #d4edda; border: 1px solid #c3e6cb; }
        .setup { background-color: #fff3cd; border: 1px solid #ffeaa7; }
        .error { background-color: #f8d7da; border: 1px solid #f5c6cb; }
        """
    ) as demo:
        
        gr.Markdown("""
        # 🚀 ComfyUI Hugging Face Space
        
        **ComfyUI executando gratuitamente no Hugging Face Spaces!**
        
        ✅ GPU Tesla T4 gratuita  
        ✅ Modelos persistentes  
        ✅ Interface web completa  
        ✅ Custom nodes incluídos  
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                # Status principal
                status_display = gr.Markdown("🟡 Aguardando inicialização...", elem_classes=["status-box", "setup"])
                
                # Botões de controle
                with gr.Row():
                    setup_btn = gr.Button("🚀 Iniciar Setup", variant="primary", scale=2)
                    restart_btn = gr.Button("🔄 Reiniciar", scale=1)
                
                # Link para ComfyUI
                comfyui_link = gr.HTML("⏳ Aguardando setup...")
                
                # Informações do sistema
                system_info = gr.Textbox(
                    label="📊 Informações do Sistema",
                    interactive=False,
                    lines=4
                )
            
            with gr.Column(scale=3):
                # Logs em tempo real
                logs_display = gr.Textbox(
                    label="📋 Logs do Setup",
                    lines=20,
                    max_lines=30,
                    interactive=False,
                    autoscroll=True
                )
        
        # Informações de ajuda
        with gr.Accordion("📚 Como Usar", open=False):
            gr.Markdown("""
            ### 🎯 Passos para usar:
            
            1. **Clique em "Iniciar Setup"** - O sistema irá baixar e configurar tudo automaticamente
            2. **Aguarde ~10-15 minutos** na primeira execução (downloads dos modelos)
            3. **Clique em "Abrir ComfyUI"** quando aparecer o link
            4. **Use normalmente** - Todos os dados ficam salvos automaticamente!
            
            ### 🔧 Recursos incluídos:
            - **FLUX.1-schnell** - Modelo de geração rápida
            - **VAE e CLIP** - Encoders necessários
            - **LoRAs populares** - Para estilos específicos
            - **ComfyUI Manager** - Para instalar novos nodes
            - **WAS Node Suite** - Nodes úteis extras
            
            ### 💡 Dicas:
            - ✅ Seus modelos ficam salvos entre sessões
            - ✅ Pode fechar e abrir o Space sem perder dados
            - ✅ Use "Reiniciar" se algo der errado
            - ✅ Logs mostram o progresso detalhado
            """)
        
        # Event handlers
        setup_btn.click(
            fn=setup_comfyui,
            outputs=[status_display, logs_display]
        )
        
        restart_btn.click(
            fn=restart_comfyui,
            outputs=[status_display]
        )
        
        # Auto-refresh de status a cada 3 segundos
        demo.load(
            fn=get_status_update,
            outputs=[status_display, logs_display, system_info, comfyui_link],
            every=3
        )
    
    return demo

# Tratamento de sinais para shutdown gracioso
def signal_handler(signum, frame):
    """Handler para shutdown gracioso"""
    logger.info(f"Recebido sinal {signum}, parando ComfyUI...")
    comfy_manager.stop_comfyui()
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

# Inicia interface
if __name__ == "__main__":
    # Setup inicial em background se não foi feito
    if not comfy_manager.setup_complete:
        setup_thread = threading.Thread(
            target=lambda: comfy_manager.setup_comfyui() and comfy_manager.start_comfyui(),
            daemon=True
        )
        setup_thread.start()
    
    # Cria e lança interface
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,  # HF Spaces já é público
        inbrowser=False
    )