import gc
import threading
import traceback

import gradio as gr
import numpy as np
import torch
from transformers import pipeline

MODEL_ID = "fishaudio/s2-pro"

_pipe = None
_pipe_error = None
_pipe_lock = threading.Lock()
_gen_lock = threading.Lock()


def load_pipeline():
    global _pipe, _pipe_error
    if _pipe is not None:
        return _pipe
    if _pipe_error is not None:
        raise RuntimeError(_pipe_error)

    with _pipe_lock:
        if _pipe is not None:
            return _pipe
        if _pipe_error is not None:
            raise RuntimeError(_pipe_error)
        try:
            _pipe = pipeline(
                task="text-to-audio",
                model=MODEL_ID,
                device=-1,
                trust_remote_code=True,
            )
            return _pipe
        except Exception as e:
            _pipe_error = f"Failed to load {MODEL_ID}: {e}"
            raise RuntimeError(_pipe_error) from e


def synthesize(text: str):
    text = (text or "").strip()
    if not text:
        raise gr.Error("Please enter some text.")

    pipe = load_pipeline()

    with _gen_lock:
        try:
            result = pipe(text)
        except Exception as e:
            tb = traceback.format_exc(limit=2)
            raise gr.Error(f"Generation failed: {e}\n\n{tb}") from e
        finally:
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    if isinstance(result, dict):
        audio = result.get("audio")
        sr = result.get("sampling_rate") or result.get("sample_rate") or 24000
    elif isinstance(result, tuple) and len(result) == 2:
        sr, audio = result
    else:
        raise gr.Error(f"Unexpected model output type: {type(result)}")

    if audio is None:
        raise gr.Error("Model returned no audio.")

    audio = np.asarray(audio)
    if audio.ndim > 1:
        audio = audio.squeeze()

    return (int(sr), audio), f"Done. Model: {MODEL_ID}"


def app_info():
    return (
        "This app tries to run fishaudio/s2-pro directly from Hugging Face using Transformers and trust_remote_code=True. "
        "If the upstream repo requires a different runtime than standard Transformers, you may need a Docker Space or GPU."
    )


with gr.Blocks() as demo:
    gr.Markdown("# Fish Audio S2 Pro Text to Speech")
    gr.Markdown(app_info())
    with gr.Row():
        text = gr.Textbox(
            label="Text",
            lines=6,
            placeholder="Type text to convert to speech...",
        )
    btn = gr.Button("Generate Speech")
    audio = gr.Audio(label="Audio", type="numpy")
    status = gr.Textbox(label="Status", interactive=False)

    btn.click(synthesize, inputs=text, outputs=[audio, status], api_name="tts")


if __name__ == "__main__":
    demo.launch()