"""XTTS fine-tuning on HF T4 GPU with auto-upload."""
import os, sys, subprocess

# Monkey-patch TTS to skip CPML license prompt (global solution)
import TTS.utils.manage as _tts_manage
_tts_manage.ModelManager.ask_tos = lambda self, path: True

import torch, gc, glob, csv, time

# Set memory optimization for T4 16GB
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,max_split_size_mb:32"

from TTS.api import TTS

print(f"CUDA: {torch.cuda.is_available()}", flush=True)
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}", flush=True)
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB", flush=True)

gc.collect()
torch.cuda.empty_cache()

tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2")
print("Model loaded!", flush=True)

# Read dataset
def read_csv(path):
    rows = []
    with open(path, newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f, delimiter='|')
        for row in reader:
            rows.append(row)
    return rows

train_rows = read_csv('train.csv')
eval_rows = read_csv('eval.csv')
print(f"Train: {len(train_rows)}, Eval: {len(eval_rows)}", flush=True)

# Copy pre-downloaded XTTS model files from TTS cache to training dir
TTS_CACHE = os.path.expanduser("~/.local/share/tts/tts_models--multilingual--multi-dataset--xtts_v2")
XTTS_DIR = os.path.join("ft_output", "run", "training", "XTTS_v2.0_original_model_files")
os.makedirs(XTTS_DIR, exist_ok=True)

if os.path.exists(TTS_CACHE):
    import shutil
    for f in os.listdir(TTS_CACHE):
        src = os.path.join(TTS_CACHE, f)
        dst = os.path.join(XTTS_DIR, f)
        if not os.path.exists(dst) and os.path.isfile(src):
            print(f"  Copying {f} from cache ({os.path.getsize(src)/1e6:.0f} MB)", flush=True)
            shutil.copy2(src, dst)

# Check what's missing and download with retries
import urllib.request
MISSING = [f for f in ["dvae.pth", "mel_stats.pth", "model.pth", "config.json", "vocab.json"]
           if not os.path.exists(os.path.join(XTTS_DIR, f)) or os.path.getsize(os.path.join(XTTS_DIR, f)) < 1000]

if MISSING:
    BASE_URL = "https://coqui.gateway.scarf.sh/hf-coqui/XTTS-v2/main"
    for fname in MISSING:
        url = f"{BASE_URL}/{fname}"
        dest = os.path.join(XTTS_DIR, fname)
        for attempt in range(5):
            try:
                print(f"  Downloading {fname}... (attempt {attempt+1}/5)", flush=True)
                urllib.request.urlretrieve(url, dest)
                print(f"  {fname} OK ({os.path.getsize(dest)/1e6:.0f} MB)", flush=True)
                break
            except Exception as e:
                print(f"  Retry {attempt+1}: {str(e)[:80]}", flush=True)
                if attempt == 4:
                    print(f"  Gave up on {fname}, will let train_gpt handle", flush=True)
                import time; time.sleep(3)

# Aggressive memory cleanup before training
gc.collect()
torch.cuda.empty_cache()

# Train with reduced memory footprint for T4 16GB
output_path = "ft_output"
os.makedirs(output_path, exist_ok=True)

from TTS.demos.xtts_ft_demo.utils.gpt_train import train_gpt

start = time.time()
train_gpt(
    language="es",
    num_epochs=10,
    batch_size=1,
    grad_acumm=4,
    train_csv="train.csv",
    eval_csv="eval.csv",
    output_path=output_path,
    max_audio_length=255995,
)
elapsed = time.time() - start

# Find best model
best = None
for f in sorted(glob.glob(os.path.join(output_path, "**", "best_model*.pth"), recursive=True)):
    best = f

print(f"\nTraining completed in {elapsed:.0f}s ({elapsed/60:.1f} min)", flush=True)
if best:
    print(f"Best model: {best}", flush=True)
    print(f"Size: {os.path.getsize(best)/1e6:.0f} MB", flush=True)

    # Upload best model to HF Space using huggingface_hub
    print("Uploading model to HF Space...", flush=True)
    
    try:
        from huggingface_hub import HfApi, CommitOperationAdd
        hf_token = os.environ.get("HF_TOKEN", "")
        if not hf_token:
            # Try reading from file
            try:
                with open("/app/.hf_token") as f:
                    hf_token = f.read().strip()
            except:
                pass
        
        if hf_token:
            api = HfApi()
            api.upload_file(
                path_or_fileobj=best,
                path_in_repo="model/best_model.pth",
                repo_id="Converso72/julio-xtts-voice",
                repo_type="model",
                token=hf_token
            )
            print(f"  Model uploaded: model/best_model.pth", flush=True)
        else:
            print("  No HF_TOKEN available, saving model locally", flush=True)
    except Exception as e:
        print(f"  Upload failed: {e}", flush=True)

    print("Model upload complete!", flush=True)

print("DONE!", flush=True)