import os
from huggingface_hub import HfApi, HfFolder
from TTS.train import train  # API oficial de entrenamiento en v0.22+
import json

# === Configuración ===
HF_TOKEN = os.environ.get("HF_TOKEN")  # definir en los Secrets del Space
if not HF_TOKEN:
    raise RuntimeError("Falta HF_TOKEN en los Secrets del Space.")

HF_REPO_ID = "sob111/xttsv2-es-finetuned"  # tu repo de destino
CONFIG_PATH = os.path.join(os.getcwd(), "config.json")
OUTPUT_PATH = "/tmp/output_model"

# Dataset local
DATA_DIR = "/tmp/voxpopuli_es_500"
META_TRAIN = os.path.join(DATA_DIR, "metadata.json")
AUDIO_DIR = os.path.join(DATA_DIR, "wav_data")

# === Validar dataset ===
print("=== Comprobando dataset local ===")
#if not os.path.isfile(META_TRAIN):
#    raise FileNotFoundError(f"No se encontró {META_TRAIN}.")
#if not os.path.isdir(AUDIO_DIR):
#    raise FileNotFoundError(f"No se encontró el directorio {AUDIO_DIR}.")

num_lines = sum(1 for _ in open(META_TRAIN, "r", encoding="utf-8"))
print(f"metadata.json: {num_lines} líneas")

# === Guardar token de Hugging Face ===
print("=== Guardando token de Hugging Face ===")
HfFolder.save_token(HF_TOKEN)

# === Entrenamiento XTTSv2 ===
print("=== Iniciando entrenamiento XTTSv2 ===")
try:
    train(config_path=CONFIG_PATH)
except Exception as e:
    raise RuntimeError("❌ El entrenamiento XTTSv2 falló. Revisa los logs anteriores.") from e

print("=== Entrenamiento finalizado ===")

# === Subir modelo fine-tune a Hugging Face ===
print("=== Subiendo modelo fine-tune a Hugging Face ===")
api = HfApi()
api.create_repo(repo_id=HF_REPO_ID, exist_ok=True, token=HF_TOKEN)
api.upload_folder(
    folder_path=OUTPUT_PATH,
    repo_id=HF_REPO_ID,
    repo_type="model",
    token=HF_TOKEN
)

print(f"✅ Fine-tuning completado y subido a {HF_REPO_ID}")