import base64
import json
import os
import shutil
import tempfile
import threading
import time
from pathlib import Path

import gradio as gr
import spaces

VOICE_API_TOKEN = os.getenv("VOICE_API_TOKEN", "").strip()
MODEL_NAME = os.getenv("F5_MODEL", "F5TTS_v1_Base")

started_at = time.time()
model_lock = threading.Lock()
f5_model = None


def check_key(api_key: str) -> None:
    if VOICE_API_TOKEN and api_key != VOICE_API_TOKEN:
        raise gr.Error("invalid api_key")


def decode_audio_b64(audio_base64: str, filename: str) -> Path:
    if "," in audio_base64[:128]:
        audio_base64 = audio_base64.split(",", 1)[1]
    suffix = Path(filename or "reference.wav").suffix or ".wav"
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    tmp_path = Path(tmp.name)
    try:
        tmp.write(base64.b64decode(audio_base64))
        tmp.close()
        return tmp_path
    except Exception:
        tmp.close()
        tmp_path.unlink(missing_ok=True)
        raise gr.Error("invalid base64 audio")


def get_model():
    global f5_model
    with model_lock:
        if f5_model is not None:
            return f5_model
        from f5_tts.api import F5TTS

        f5_model = F5TTS(model=MODEL_NAME, device="cpu")
        return f5_model


def health(api_key: str = "") -> str:
    check_key(api_key)
    return json.dumps(
        {
            "ok": True,
            "model": MODEL_NAME,
            "model_loaded": f5_model is not None,
            "uptime_seconds": round(time.time() - started_at, 3),
        },
        ensure_ascii=False,
    )


@spaces.GPU(duration=10)
def zerogpu_probe() -> str:
    return "ok"


def clone_b64(
    api_key: str,
    ref_audio_base64: str,
    ref_filename: str,
    ref_text: str,
    gen_text: str,
    speed: float = 1.0,
    nfe_step: int = 32,
) -> str:
    check_key(api_key)
    speed = float(speed)
    nfe_step = int(nfe_step)
    if not ref_text.strip():
        raise gr.Error("ref_text is required")
    if not gen_text.strip():
        raise gr.Error("gen_text is required")
    if len(gen_text) > 1000:
        raise gr.Error("gen_text is too long for free CPU; max 1000 characters")
    if speed < 0.5 or speed > 2.0:
        raise gr.Error("speed must be between 0.5 and 2.0")
    if nfe_step < 8 or nfe_step > 64:
        raise gr.Error("nfe_step must be between 8 and 64")

    ref_path = decode_audio_b64(ref_audio_base64, ref_filename or "reference.wav")
    output_dir = Path(tempfile.mkdtemp(prefix="f5_clone_"))
    output_wav = output_dir / "output.wav"
    output_spec = output_dir / "output.png"

    try:
        model = get_model()
        try:
            model.infer(
                ref_file=str(ref_path),
                ref_text=ref_text,
                gen_text=gen_text,
                file_wave=str(output_wav),
                file_spec=str(output_spec),
                speed=speed,
                nfe_step=nfe_step,
                seed=None,
            )
        except TypeError:
            model.infer(
                ref_file=str(ref_path),
                ref_text=ref_text,
                gen_text=gen_text,
                file_wave=str(output_wav),
                file_spec=str(output_spec),
                speed=speed,
                seed=None,
            )
        if not output_wav.exists():
            raise gr.Error("voice clone produced no output file")
        audio_base64 = base64.b64encode(output_wav.read_bytes()).decode("ascii")
        return json.dumps(
            {
                "audio_base64": audio_base64,
                "mime_type": "audio/wav",
                "model": MODEL_NAME,
            },
            ensure_ascii=False,
        )
    finally:
        ref_path.unlink(missing_ok=True)
        shutil.rmtree(output_dir, ignore_errors=True)


with gr.Blocks(title="F5 Voice Clone API") as demo:
    gr.Markdown("# F5 Voice Clone API")
    gr.Markdown("Use `clone_b64` from n8n. Reference audio and output audio are base64 for simple JSON automation.")
    api_key = gr.Textbox(label="api_key", type="password")
    with gr.Tab("Health"):
        health_out = gr.Textbox(label="result")
        gr.Button("Check").click(health, inputs=[api_key], outputs=[health_out], api_name="health")
        gr.Button("ZeroGPU probe", visible=False).click(
            zerogpu_probe,
            inputs=[],
            outputs=[health_out],
            api_name="zerogpu_probe",
        )
    with gr.Tab("Clone"):
        ref_audio_base64 = gr.Textbox(label="ref_audio_base64", lines=5)
        ref_filename = gr.Textbox(label="ref_filename", value="reference.wav")
        ref_text = gr.Textbox(label="ref_text", lines=3)
        gen_text = gr.Textbox(label="gen_text", lines=4)
        speed = gr.Number(label="speed", value=1.0)
        nfe_step = gr.Number(label="nfe_step", value=32, precision=0)
        result = gr.Textbox(label="result", lines=8)
        gr.Button("Clone").click(
            clone_b64,
            inputs=[api_key, ref_audio_base64, ref_filename, ref_text, gen_text, speed, nfe_step],
            outputs=[result],
            api_name="clone_b64",
        )

demo.queue(default_concurrency_limit=1).launch()