# setup.py
#
# Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
#
# Versão 3.0.0 (Setup Unificado com LTX, SeedVR e VINCIE)
# - Orquestra a instalação de todos os repositórios e modelos para a suíte ADUC-SDR.
# - Garante que todos os downloads usem um cache persistente para inicializações rápidas.

import os
import subprocess
import sys
from pathlib import Path
import yaml
from huggingface_hub import hf_hub_download, snapshot_download

# --- Configuração de Paths e Cache ---
# Assume-se que /data é um volume persistente montado no contêiner.
DEPS_DIR = Path("/data")
CACHE_DIR = DEPS_DIR / ".cache" / "huggingface"

# --- Configuração dos Módulos da Aplicação ---
LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
SEEDVR_MODELS_DIR = DEPS_DIR / "models" / "SeedVR"
VINCIE_REPO_DIR = DEPS_DIR / "VINCIE"
VINCIE_CKPT_DIR = DEPS_DIR / "ckpt" / "VINCIE-3B"

# --- Repositórios Git para Clonar ---
REPOS_TO_CLONE = {
    "LTX-Video": "https://huggingface.co/spaces/Lightricks/ltx-video-distilled",
    "SeedVR": "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler",
    "VINCIE": "https://github.com/ByteDance-Seed/VINCIE",
}

def run_command(command, cwd=None):
    """Executa um comando no terminal de forma segura e com logs claros."""
    print(f"Executando: {' '.join(command)}")
    try:
        # Usamos check=True para lançar uma exceção se o comando falhar.
        subprocess.run(
            command,
            check=True,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
    except subprocess.CalledProcessError as e:
        print(f"ERRO: O comando falhou com o código {e.returncode}")
        print(f"Stderr:\n{e.stderr.strip()}")
        sys.exit(1)
    except FileNotFoundError:
        print(f"ERRO: Comando '{command[0]}' não encontrado. Verifique se o git está instalado no contêiner.")
        sys.exit(1)

def _load_ltx_config():
    """Carrega o arquivo de configuração YAML do LTX-Video."""
    print("--- Carregando Configuração do LTX-Video ---")
    config_file = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled-fp8.yaml"
    
    if not config_file.exists():
        print(f"ERRO: Arquivo de configuração do LTX não encontrado em '{config_file}'")
        return None
        
    print(f"Configuração LTX encontrada: {config_file}")
    with open(config_file, "r") as file:
        return yaml.safe_load(file)

def _download_models_from_hub(repo_id, filenames, cache_dir, local_dir=None):
    """Função auxiliar genérica para baixar arquivos do Hugging Face Hub."""
    for filename in filenames:
        if not filename: continue
        
        print(f"Verificando/Baixando '{filename}' de '{repo_id}'...")
        try:
            hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                cache_dir=str(cache_dir),
                local_dir=str(local_dir) if local_dir else None,
                local_dir_use_symlinks=False, # Salva o arquivo real, não um link simbólico
                token=os.getenv("HF_TOKEN"),
            )
            print(f"-> '{filename}' está disponível.")
        except Exception as e:
            if "404" in str(e) or "Not Found" in str(e):
                 print(f"AVISO: Arquivo '{filename}' não encontrado em '{repo_id}'. Isso pode ser normal se for um ponteiro para um diretório (como em Text Encoders). Continuando.")
            else:
                print(f"ERRO CRÍTICO ao baixar '{filename}': {e}")
                sys.exit(1)

def main():
    """Função principal que orquestra todo o processo de setup."""
    print("--- Iniciando Setup do Ambiente ADUC-SDR (LTX + SeedVR + VINCIE) ---")
    DEPS_DIR.mkdir(exist_ok=True)
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    
    # --- ETAPA 1: Clonar Repositórios ---
    print("\n--- ETAPA 1: Verificando Repositórios Git ---")
    for repo_name, repo_url in REPOS_TO_CLONE.items():
        repo_path = DEPS_DIR / repo_name
        if repo_path.is_dir():
            print(f"Repositório '{repo_name}' já existe em '{repo_path}'. Pulando.")
        else:
            print(f"Clonando '{repo_name}' de {repo_url}...")
            run_command(["git", "clone", "--depth", "1", repo_url, str(repo_path)])
            print(f"-> '{repo_name}' clonado com sucesso.")

    # --- ETAPA 2: Baixar Modelos do LTX-Video e suas Dependências ---
    print("\n--- ETAPA 2: Verificando Modelos LTX-Video e Dependências ---")
    ltx_config = _load_ltx_config()
    if not ltx_config:
        print("ERRO: Não foi possível carregar a configuração do LTX-Video. Abortando.")
        sys.exit(1)
        
    _download_models_from_hub(
        repo_id="Lightricks/LTX-Video",
        filenames=[ltx_config.get("checkpoint_path"), ltx_config.get("spatial_upscaler_model_path")],
        cache_dir=CACHE_DIR
    )
    
    dependency_repos = [
        ltx_config.get("text_encoder_model_name_or_path"),
        ltx_config.get("prompt_enhancer_image_caption_model_name_or_path"),
        ltx_config.get("prompt_enhancer_llm_model_name_or_path"),
    ]
    for repo_id in filter(None, dependency_repos):
         _download_models_from_hub(repo_id=repo_id, filenames=["config.json"], cache_dir=CACHE_DIR)

    # --- ETAPA 3: Baixar Modelos do SeedVR ---
    print("\n--- ETAPA 3: Verificando Modelos SeedVR ---")
    SEEDVR_MODELS_DIR.mkdir(parents=True, exist_ok=True)
    seedvr_files = {
        "seedvr2_ema_7b_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
        "seedvr2_ema_7b_sharp_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
        "ema_vae_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
    }
    for filename, repo_id in seedvr_files.items():
        if not (SEEDVR_MODELS_DIR / filename).is_file():
             _download_models_from_hub(repo_id=repo_id, filenames=[filename], cache_dir=CACHE_DIR, local_dir=SEEDVR_MODELS_DIR)
        else:
            print(f"Arquivo SeedVR '{filename}' já existe. Pulando.")

    # --- ETAPA 4: Baixar Modelos VINCIE ---
    print("\n--- ETAPA 4: Verificando Modelos VINCIE ---")
    VINCIE_CKPT_DIR.mkdir(parents=True, exist_ok=True)
    try:
        print(f"Verificando snapshot do VINCIE em '{VINCIE_CKPT_DIR}'...")
        # snapshot_download é ideal para baixar um repositório inteiro.
        snapshot_download(
            repo_id="ByteDance-Seed/VINCIE-3B",
            local_dir=str(VINCIE_CKPT_DIR),
            cache_dir=str(CACHE_DIR),
            force_download=False, # Não baixa novamente se o diretório já estiver completo
            token=os.getenv("HF_TOKEN")
        )
        print("-> Modelos VINCIE estão disponíveis.")
        
        # Cria o symlink de compatibilidade, se necessário, para que o código do VINCIE encontre os modelos
        repo_ckpt_dir = VINCIE_REPO_DIR / "ckpt"
        repo_ckpt_dir.mkdir(parents=True, exist_ok=True)
        link = repo_ckpt_dir / "VINCIE-3B"
        if not link.exists():
            link.symlink_to(VINCIE_CKPT_DIR.resolve(), target_is_directory=True)
            print(f"-> Symlink de compatibilidade VINCIE criado: '{link}' -> '{VINCIE_CKPT_DIR.resolve()}'")
        else:
            print(f"-> Symlink de compatibilidade VINCIE já existe.")

    except Exception as e:
        print(f"ERRO CRÍTICO ao baixar modelos VINCIE: {e}")
        sys.exit(1)

    print("\n\n--- ✅ Setup Completo do Ambiente ADUC-SDR Concluído com Sucesso! ---")
    print("Todos os repositórios e modelos foram verificados e estão prontos para uso.")

if __name__ == "__main__":
    main()