"""ArdiVoice — Amharic TTS Space (OmniVoice, ZeroGPU).

Exposes one Gradio endpoint (`/tts`) used both by the web UI and by the
ArdiVoiceClient voice agent through `gradio_client`.
"""

import os
import tempfile
import traceback

import gradio as gr
import numpy as np
import soundfile as sf
import spaces  # ZeroGPU allocation
import torch
from omnivoice import OmniVoice, OmniVoiceGenerationConfig

MODEL_ID = os.environ.get("MODEL_ID", "african-low-resource/omnivoice-amharic")
MAX_CHARS = 400  # one sentence at a time; longer prompts degrade prosody

# The model MUST be loaded on CPU here. On ZeroGPU any CUDA initialisation in the
# main process kills the container (with no log output), and from_pretrained also
# places the Higgs audio tokenizer on `model.device` — so device_map="cuda:0"
# would touch CUDA twice before a GPU is even allocated. It is moved to the GPU
# inside the @spaces.GPU function instead, where the allocation is valid.
print(f"Loading {MODEL_ID} on CPU", flush=True)
model = OmniVoice.from_pretrained(MODEL_ID, device_map="cpu", dtype=torch.float32)
SAMPLE_RATE = int(getattr(model, "sampling_rate", 24_000))
print(f"Model loaded (sampling_rate={SAMPLE_RATE})", flush=True)

_state = {"on_gpu": False, "asr_loaded": False}


def _ensure_gpu():
    """Move the model (and its audio tokenizer submodule) onto the allocated GPU."""
    if _state["on_gpu"] or not torch.cuda.is_available():
        return
    model.to("cuda")
    _state["on_gpu"] = True
    print("Model moved to CUDA", flush=True)


def _ensure_asr():
    """Whisper, used only to auto-transcribe a reference clip when ref_text is blank."""
    if _state["asr_loaded"]:
        return
    model.load_asr_model(device="cuda" if _state["on_gpu"] else "cpu")
    _state["asr_loaded"] = True


def _to_wav_array(audio):
    """Normalise whatever generate() returned into a 1-D/2-D float32 numpy array."""
    if isinstance(audio, tuple):
        audio = audio[0]
    if isinstance(audio, torch.Tensor):
        audio = audio.detach().float().cpu().numpy()
    elif isinstance(audio, (list, tuple)):
        audio = np.asarray(audio, dtype=np.float32)
    else:
        audio = np.asarray(audio)

    audio = np.squeeze(audio)
    if audio.ndim == 2 and audio.shape[0] < audio.shape[1]:
        # (channels, frames) -> (frames, channels) for soundfile
        audio = audio.T
    return np.ascontiguousarray(audio, dtype=np.float32)


@spaces.GPU(duration=120)
def tts(text: str, ref_audio, ref_text: str, num_steps: int, guidance_scale: float):
    """Synthesize Amharic speech. Returns (wav_path, status)."""
    if not text or not text.strip():
        return None, "Please enter some Amharic text."

    text = text.strip()
    if len(text) > MAX_CHARS:
        return None, f"Text too long ({len(text)} chars). Send at most {MAX_CHARS}."

    try:
        _ensure_gpu()

        kwargs = {
            "text": text,
            "language": "Amharic",
            "generation_config": OmniVoiceGenerationConfig(
                num_step=int(num_steps),
                guidance_scale=float(guidance_scale),
            ),
        }

        # ref_audio is a filepath from gr.Audio(type="filepath"); None = default voice.
        if ref_audio:
            if not os.path.exists(ref_audio):
                return None, f"Reference audio not found: {ref_audio}"
            ref_text = (ref_text or "").strip()
            if not ref_text:
                # create_voice_clone_prompt(ref_text=None) transcribes the clip, which
                # requires an ASR model to have been loaded first.
                _ensure_asr()
            kwargs["voice_clone_prompt"] = model.create_voice_clone_prompt(
                ref_audio=ref_audio, ref_text=ref_text or None
            )

        audio = model.generate(**kwargs)
        wav = _to_wav_array(audio[0] if len(audio) else audio)

        if wav.size == 0:
            return None, "Model returned empty audio."

        # delete=False: Gradio serves the file after this function returns.
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        tmp.close()
        sf.write(tmp.name, wav, SAMPLE_RATE)

        mode = "voice cloning" if ref_audio else "default voice"
        secs = wav.shape[0] / SAMPLE_RATE
        return tmp.name, f"OK — {mode}, {secs:.1f}s, steps={num_steps}, cfg={guidance_scale}"

    except Exception as exc:  # surfaced to the client instead of a bare 500
        print(traceback.format_exc(), flush=True)
        return None, f"Generation error: {type(exc).__name__}: {exc}"


with gr.Blocks(title="ArdiVoice — Amharic TTS") as demo:
    gr.Markdown(
        "# 🇪🇹 ArdiVoice — Amharic TTS\n"
        "Powered by [omnivoice-amharic](https://huggingface.co/african-low-resource/omnivoice-amharic). "
        "Upload a 5–15s sample to clone a voice, or leave it empty for the default voice."
    )

    text_in = gr.Textbox(
        label="Amharic text",
        value="ሰላም! እንዴት ልረዳዎ እችላለሁ?",
        lines=3,
        max_lines=6,
    )
    ref_in = gr.Audio(label="Reference audio (optional, 5–15s)", type="filepath")
    ref_text_in = gr.Textbox(
        label="Reference transcript (optional)",
        placeholder="Leave empty to transcribe the reference clip automatically",
        lines=1,
    )
    with gr.Row():
        steps_in = gr.Slider(4, 64, value=32, step=1, label="Diffusion steps")
        cfg_in = gr.Slider(1.0, 5.0, value=2.0, step=0.1, label="Guidance scale")

    run_btn = gr.Button("🎙️ Generate", variant="primary")
    audio_out = gr.Audio(label="Synthesized speech", type="filepath")
    status_out = gr.Textbox(label="Status", interactive=False)

    run_btn.click(
        fn=tts,
        inputs=[text_in, ref_in, ref_text_in, steps_in, cfg_in],
        outputs=[audio_out, status_out],
        api_name="tts",
    )

    gr.Examples(
        examples=[
            ["ሰላም፣ እንኳን ደህና መጣችሁ። ይህ የአማርኛ ንግግር ሙከራ ነው።", None, "", 32, 2.0],
            ["ዛሬ ቀን ጥሩ ነው። እንዴት ነህ?", None, "", 32, 2.0],
            ["ኢትዮጵያ ታላቅ አገር ናት።", None, "", 32, 2.0],
        ],
        inputs=[text_in, ref_in, ref_text_in, steps_in, cfg_in],
    )

if __name__ == "__main__":
    demo.queue(max_size=16).launch(show_error=True)