import os
os.environ["COQUI_TOS_AGREED"] = "1"                # Coqui TOS auto-accept
os.environ.setdefault("COQUI_CACHE_DIR", "/home/user/.coqui")

import glob
import gradio as gr
import numpy as np
import soundfile as sf
from huggingface_hub import snapshot_download
from TTS.api import TTS

HF_REPO = "suhaibrashid17/XTTS-v2-Urdu-FT"
SAMPLE_RATE = 24000

def _find_file(patterns, root):
    for pat in patterns:
        hits = glob.glob(os.path.join(root, pat), recursive=True)
        if hits:
            return hits[0]
    return None

def load_finetuned_xtts():
    # Download to local cache (no symlinks to avoid weird paths)
    repo_dir = snapshot_download(
        repo_id=HF_REPO,
        local_dir_use_symlinks=False,
        allow_patterns=["**/*"]  # pull everything once; avoids missing tokenizer/vocab
    )

    # Locate config + weight file
    cfg = _find_file(["config.json", "**/config.json"], repo_dir)
    mdl = _find_file(
        ["model.pth", "pytorch_model.bin", "**/model.pth", "**/pytorch_model.bin"],
        repo_dir,
    )
    if not (cfg and mdl):
        raise RuntimeError(
            "Could not find model/config. Need config.json and model.pth (or pytorch_model.bin)."
        )

    # >>> IMPORTANT FIX <<<
    # Coqui XTTS expects model_path = DIRECTORY that contains model.pth
    model_dir = os.path.dirname(mdl)

    # Load by giving directory (not the file) to model_path
    tts = TTS(model_path=model_dir, config_path=cfg, gpu=False, progress_bar=False)
    return tts

# ---- load once on startup
TTS_ENGINE = load_finetuned_xtts()

# ---------- helpers ----------
def _ensure_wav_path(audio):
    """
    Gradio v5: Audio can be filepath or dict/tuple/ndarray.
    Always return a real wav path.
    """
    if audio is None:
        return None
    if isinstance(audio, str):
        return audio
    if isinstance(audio, dict) and "data" in audio:
        data = audio["data"]
        sr = int(audio.get("sampling_rate", SAMPLE_RATE))
        path = "ref.wav"
        sf.write(path, data, sr)
        return path
    if isinstance(audio, (tuple, list)):
        # (sr, data) or (data, sr)
        if isinstance(audio[0], (int, float)):
            sr, data = audio
        else:
            data, sr = audio
        path = "ref.wav"
        sf.write(path, data, int(sr))
        return path
    if isinstance(audio, np.ndarray):
        path = "ref.wav"
        sf.write(path, audio, SAMPLE_RATE)
        return path
    raise ValueError("Unsupported audio input format")

def clean_urdu_text(txt: str) -> str:
    if not txt:
        return ""
    t = " ".join(txt.split())
    if t and t[-1] not in "۔!?؟":
        t += "۔"
    return t

# ---------- inference ----------
def clone_urdu(reference_audio, text):
    ref_path = _ensure_wav_path(reference_audio)
    if not ref_path:
        return gr.Error("Reference audio required (5–30 seconds).")

    txt = clean_urdu_text(text)
    try:
        out_path = "output_xtts_ur.wav"
        # XTTS fine-tuned for Urdu
        TTS_ENGINE.tts_to_file(
            text=txt,
            file_path=out_path,
            speaker_wav=ref_path,
            language="ur",
        )
        return out_path
    except Exception as e:
        return gr.Error(f"XTTS-Urdu error: {e}")

# ---------- UI ----------
with gr.Blocks(title="XTTS-v2 Urdu — Voice Clone (CPU)") as demo:
    gr.Markdown("## XTTS-v2 Urdu — Voice Clone")

    ref = gr.Audio(
        label="Reference voice (WAV/MP3, ~5–30 sec)",
        sources=["upload", "microphone"],
        type="filepath",
        waveform_options={"show_controls": True},
    )

    txt = gr.Textbox(
        label="Urdu text",
        lines=4,
        placeholder="یہاں اردو متن لکھیں… مختصر جملے بہتر نتیجہ دیتے ہیں۔"
    )

    btn = gr.Button("Generate", variant="primary")
    out = gr.Audio(label="Output", type="filepath")

    btn.click(clone_urdu, inputs=[ref, txt], outputs=out)

if __name__ == "__main__":
    demo.queue(max_size=8).launch(server_name="0.0.0.0", server_port=7860)