# setup.py

import os
import subprocess
import sys
from pathlib import Path
import yaml

# Adiciona a importação do hf_hub_download
from huggingface_hub import hf_hub_download

# --- Configuração ---
DEPS_DIR = Path("/data")
MODELS_ROOT = DEPS_DIR / "ltx_models" # Diretório centralizado para modelos

# URLs explícitas para os repositórios
REPOS_TO_CLONE = {
    "LTX-Video": "https://huggingface.co/spaces/Lightricks/ltx-video-distilled",
    "SeedVR_Space": "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B",
    "MMAudio": "https://github.com/hkchengrex/MMAudio.git"
}

def run_command(command, cwd=None, is_pip=False):
    """Executa um comando no terminal e exibe a saída em tempo real."""
    if is_pip:
        command = [sys.executable, "-m", "pip"] + command
        
    print(f"Executando: {' '.join(command)}")
    try:
        with subprocess.Popen(
            command, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            text=True, encoding='utf-8', bufsize=1
        ) as process:
            for line in iter(process.stdout.readline, ''):
                print(line, end='')
            
            return_code = process.wait()
            if return_code != 0:
                raise subprocess.CalledProcessError(return_code, command)
    except subprocess.CalledProcessError as e:
        print(f"ERRO: O comando falhou com o código de saída {e.returncode}")
        sys.exit(1)
    except FileNotFoundError:
        print(f"ERRO: O comando '{command[0]}' não foi encontrado.")
        sys.exit(1)

def clone_repos():
    """ETAPA 1: Garante que todos os repositórios de código necessários existam."""
    print("--- ETAPA 1: Verificando e clonando repositórios de código ---")
    for repo_name, repo_url in REPOS_TO_CLONE.items():
        repo_path = DEPS_DIR / repo_name
        if repo_path.exists() and (repo_path / ".git").exists():
            print(f"Repositório '{repo_name}' já existe. Pulando.")
        else:
            print(f"Clonando '{repo_name}' de {repo_url}...")
            run_command(["git", "clone", "--depth", "1", "-q", repo_url, str(repo_path)])
            print(f"'{repo_name}' clonado com sucesso.")

def install_ltx_dependencies():
    """ETAPA 2: Instala as dependências Python do LTX-Video."""
    print("\n--- ETAPA 2: Instalando dependências Python do LTX-Video ---")
    ltx_repo_path = DEPS_DIR / "LTX-Video"
    if ltx_repo_path.exists():
        run_command(["install", "-e", ".[inference]"], cwd=str(ltx_repo_path), is_pip=True)
        print("Dependências do LTX-Video instaladas com sucesso.")
    else:
        print("ERRO CRÍTICO: Diretório LTX-Video não encontrado. Não foi possível instalar dependências.")
        sys.exit(1)

def download_ltx_models():
    """
    ETAPA 3: Verifica e baixa os modelos LTX necessários, lendo os configs
    do repositório já clonado.
    """
    print("\n--- ETAPA 3: Verificando e baixando modelos LTX ---")
    LTX_REPO_ID = "Lightricks/LTX-Video" # Repositório dos modelos no Hub
    ltx_configs_dir = DEPS_DIR / "LTX-Video" / "configs"
    
    if not ltx_configs_dir.is_dir():
        print(f"ERRO CRÍTICO: Diretório de configs '{ltx_configs_dir}' não encontrado. O clone do LTX-Video pode ter falhado.")
        sys.exit(1)

    required_model_files = set()
    for cfg_path in ltx_configs_dir.glob("*.yaml"):
        try:
            with open(cfg_path, 'r') as f:
                config = yaml.safe_load(f)
                if config and config.get("checkpoint_path"):
                    required_model_files.add(config["checkpoint_path"])
                if config and config.get("spatial_upscaler_model_path"):
                    required_model_files.add(config["spatial_upscaler_model_path"])
        except Exception as e:
            print(f"AVISO: Falha ao ler o arquivo de configuração {cfg_path}: {e}")

    if not required_model_files:
        print("AVISO: Nenhum arquivo de modelo foi encontrado nos arquivos de configuração. Verifique os configs.")
        return

    print(f"Modelos necessários encontrados: {required_model_files}")
    MODELS_ROOT.mkdir(exist_ok=True)
    
    for model_filename in required_model_files:
        local_path = MODELS_ROOT / model_filename
        if local_path.exists():
            print(f"Modelo '{model_filename}' já existe localmente. Pulando.")
        else:
            print(f"Baixando modelo '{model_filename}'...")
            try:
                hf_hub_download(
                    repo_id=LTX_REPO_ID,
                    filename=model_filename,
                    local_dir=str(MODELS_ROOT),
                    local_dir_use_symlinks=False,
                    token=os.getenv("HF_TOKEN")
                )
                print(f"'{model_filename}' baixado com sucesso.")
            except Exception as e:
                print(f"ERRO: Falha ao baixar '{model_filename}'. Verifique o nome do arquivo e a conexão. Erro: {e}")

def main():
    print("--- Iniciando Setup Completo do Ambiente ADUC-SDR ---")
    DEPS_DIR.mkdir(exist_ok=True)
    
    # Executa as etapas na ordem correta
    clone_repos()
    install_ltx_dependencies()
    download_ltx_models()

    print("\n--- ✅ Setup do Ambiente Concluído com Sucesso! ---")
    print("Todos os repositórios, dependências e modelos estão prontos.")

if __name__ == "__main__":
    main()