#!/usr/bin/env python3
"""
Passo 0: Setup do Ambiente
Instala todas as dependências para:
- Geração de dataset (Soprano TTS, Whisper, SNAC, NeMo NFA)
- Treinamento do modelo Speech-to-Speech

PyTorch 2.7.0 com suporte oficial a Blackwell (sm_120)

Suporta:
- Blackwell: RTX 5090, 5080, 5070, B100, B200 (sm_120) - CUDA 12.8
- Hopper: H100, H200 (sm_90) - CUDA 12.8
- Ada: RTX 4090, 4080, L40S (sm_89) - CUDA 12.4
- Ampere: A100, RTX 3090 (sm_80/86) - CUDA 12.4

Usage:
    python passo0_setup.py [--skip_test]
"""

import os
import sys
import subprocess
import shutil

HF_TOKEN = os.environ.get("HF_TOKEN", "")


def log(msg):
    print(f"[SETUP] {msg}")
    sys.stdout.flush()


def run(cmd, check=True):
    """Execute command."""
    log(f"$ {cmd}")
    result = subprocess.run(cmd, shell=True)
    if check and result.returncode != 0:
        log(f"Command failed with code {result.returncode}")
        return False
    return True


def get_gpu_info():
    """Get GPU info and architecture."""
    result = subprocess.run(['nvidia-smi', '--query-gpu=name,compute_cap', '--format=csv,noheader'],
                          capture_output=True, text=True)
    if result.returncode != 0:
        return "", ""
    lines = result.stdout.strip().split('\n')
    if lines:
        parts = lines[0].split(',')
        name = parts[0].strip()
        compute = parts[1].strip() if len(parts) > 1 else ""
        return name, compute
    return "", ""


def needs_cuda_128(gpu_name, compute_cap):
    """Check if GPU needs CUDA 12.8+ (Hopper/Blackwell architecture)."""
    # H100, H200 = Hopper (sm_90)
    # RTX 50xx, B100, B200 = Blackwell (sm_120) - requires PyTorch 2.7+
    hopper_blackwell = ["H100", "H200", "RTX 50", "5090", "5080", "5070", "B100", "B200"]
    if any(arch in gpu_name for arch in hopper_blackwell):
        return True
    # Compute capability 9.0+ (Hopper) or 12.0+ (Blackwell) needs CUDA 12.8
    try:
        cap = float(compute_cap)
        if cap >= 9.0:  # sm_90 (Hopper) or sm_120 (Blackwell)
            return True
    except:
        pass
    return False


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--skip_test", action="store_true")
    args = parser.parse_args()

    log("="*60)
    log("SETUP DO AMBIENTE - H200/Hopper Compatible")
    log("="*60)

    gpu_name, compute_cap = get_gpu_info()
    log(f"GPU: {gpu_name} (compute {compute_cap})")
    use_cuda_128 = needs_cuda_128(gpu_name, compute_cap)
    log(f"Using CUDA 12.8: {use_cuda_128}")

    # 1. System dependencies
    log("\n[1/7] System dependencies...")
    if shutil.which('apt-get'):
        run("apt-get update -qq && apt-get install -y -qq espeak-ng espeak libsndfile1 ffmpeg git wget curl build-essential sox libsox-fmt-all", check=False)

    # 2. PyTorch (version depends on GPU architecture)
    # PyTorch 2.7.0 has official Blackwell (sm_120) support
    log("\n[2/7] PyTorch 2.7.0 (Blackwell compatible)...")
    if use_cuda_128:
        # Hopper (H100/H200) and Blackwell (RTX 50xx) need CUDA 12.8
        log("  Installing PyTorch 2.7.0 with CUDA 12.8 (Blackwell/Hopper support)")
        run("pip install torch==2.7.0 torchaudio==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu128 -q")
    else:
        # Ampere (A100, RTX 30xx) and Ada (RTX 40xx) can use CUDA 12.4
        log("  Installing PyTorch 2.7.0 with CUDA 12.4 (Ampere/Ada)")
        run("pip install torch==2.7.0 torchaudio==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu124 -q")

    # 3. Core packages
    log("\n[3/7] Core packages...")
    packages = [
        "'numpy>=2.0.2,<2.1.0'",
        "scipy", "soundfile", "librosa",
        "'transformers>=4.52.0'", "accelerate", "peft", "safetensors",
        "huggingface-hub", "snac", "omegaconf", "tqdm", "requests",
        "'pandas>=2.2.3'",
        "tensorboard", "psutil", "groq",
        "Cython",  # Required for NeMo
    ]
    run(f"pip install {' '.join(packages)} -q")

    # 4. NeMo for Forced Alignment (GPU-accelerated)
    log("\n[4/7] NeMo Forced Aligner (GPU-accelerated)...")
    # NeMo requires specific versions
    run("pip install 'nemo_toolkit[asr]>=2.0.0' -q")

    # 5. Whisper (native transformers - H200 compatible, no ctranslate2)
    log("\n[5/7] Whisper (transformers - H200 compatible)...")
    # We use native Whisper from transformers (already installed above)
    # No whisperx/ctranslate2 needed - they don't support H200 (sm_90)

    # 6. Soprano TTS 80M (ultra-lightweight, 2000x realtime)
    # See: https://github.com/ekwek1/soprano
    log("\n[6/8] Soprano TTS 80M...")

    # Install soprano-tts with lmdeploy for fastest inference
    # lmdeploy provides 2000x realtime speed vs ~10x for transformers backend
    run("pip install soprano-tts -q")

    # lmdeploy doesn't support Blackwell (RTX 50xx) yet, so only install on supported GPUs
    blackwell_gpus = ["RTX 50", "5090", "5080", "5070", "B100", "B200"]
    is_blackwell = any(arch in gpu_name for arch in blackwell_gpus)

    if is_blackwell:
        log("  Blackwell GPU detected - skipping lmdeploy (not supported yet)")
        log("  Soprano will use transformers backend (slower but compatible)")
    else:
        log("  Installing lmdeploy for 2000x realtime speed...")
        run("pip install lmdeploy -q")

    # 7. Create directories
    log("\n[7/8] Creating directories...")
    for d in ["./data", "./data/raw", "./data/processed", "./checkpoints", "./logs"]:
        os.makedirs(d, exist_ok=True)

    # 8. Test Soprano with lmdeploy
    if not args.skip_test:
        log("\n[8/8] Testing components...")
        log("\n" + "="*60)
        log("TESTING SOPRANO TTS (lmdeploy backend)")
        log("="*60)
        run("""python3 -c "
import torch
_load = torch.load
torch.load = lambda *a, **k: _load(*a, **{**k, 'weights_only': False})

from soprano import SopranoTTS
import time

print('Loading Soprano TTS with lmdeploy backend...')
try:
    tts = SopranoTTS(backend='lmdeploy', device='cuda', cache_size_mb=2000, decoder_batch_size=8)
    print('Using lmdeploy backend (fastest)')
except Exception as e:
    print(f'lmdeploy failed ({e}), falling back to transformers')
    tts = SopranoTTS(backend='transformers', device='cuda')

# Warmup
for _ in range(3):
    tts.infer('warmup')

# Speed test
t = time.time()
for _ in range(10):
    tts.infer('Hello, this is a test.')
print(f'Speed: {10/(time.time()-t):.1f} calls/s')
print('Soprano TTS OK!')
"
""")

    # Test NeMo NFA (if not skipped)
    if not args.skip_test:
        log("\n" + "="*60)
        log("TESTING NEMO FORCED ALIGNER")
        log("="*60)
        run("""python3 -c "
import torch
_load = torch.load
torch.load = lambda *a, **k: _load(*a, **{**k, 'weights_only': False})

try:
    import nemo.collections.asr as nemo_asr
    print('NeMo ASR import OK')
    # Quick model check (downloads small model)
    print('NeMo Forced Aligner ready!')
except Exception as e:
    print(f'NeMo warning (may still work): {e}')
"
""")

    log("\n" + "="*60)
    log("SETUP COMPLETO!")
    log("="*60)
    log("""
Para gerar dataset:
    cd datasets && python create_dataset.py --count 1000 --output ../data/dataset.pt --gpus 1

Para treinar:
    python passo2_finetune_stage1.py --data ./data/dataset.pt
    python passo3_finetune_stage2.py --data ./data/dataset.pt --stage1_ckpt ./checkpoints/stage1_best.pt
""")


if __name__ == "__main__":
    main()