#!/usr/bin/env python3
"""
HuggingFace Space entry point for OmniVoice — batch voice cloning.

One reference audio + one reference text + N target texts -> N generated audios.
"""

import logging
import os
import re
import tempfile
import time
import zipfile
from typing import List, Optional

logging.basicConfig(
    level=logging.WARNING,
    format="%(asctime)s %(name)s %(levelname)s: %(message)s",
)
logging.getLogger("omnivoice").setLevel(logging.DEBUG)

import gradio as gr
import numpy as np
import soundfile as sf
import spaces
import torch
from omnivoice.utils.lang_map import LANG_NAMES, lang_display_name

from omnivoice import OmniVoice, OmniVoiceGenerationConfig

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
CHECKPOINT = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice")

print(f"Loading model from {CHECKPOINT} to cuda ...")
model = OmniVoice.from_pretrained(
    CHECKPOINT,
    device_map="cuda",
    dtype=torch.float16,
    load_asr=True,
)
sampling_rate = model.sampling_rate
print("Model loaded successfully!")

_ALL_LANGUAGES = ["Auto"] + sorted(lang_display_name(n) for n in LANG_NAMES)

# Hard cap on how many lines we accept in one request.
MAX_ITEMS = 512
# Seconds requested per ZeroGPU call. Each sub-batch gets its own call, so this
# bounds a single sub-batch, not the whole job.
GPU_DURATION = 120


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def parse_texts(raw: str) -> List[str]:
    """One target text per line; blank lines ignored."""
    if not raw:
        return []
    return [line.strip() for line in raw.splitlines() if line.strip()]


def slugify(text: str, max_len: int = 32) -> str:
    slug = re.sub(r"[^\w\s-]", "", text, flags=re.UNICODE).strip()
    slug = re.sub(r"[\s_-]+", "-", slug)
    return slug[:max_len].strip("-") or "audio"


# ---------------------------------------------------------------------------
# GPU worker: one sub-batch per call
# ---------------------------------------------------------------------------


@spaces.GPU(duration=GPU_DURATION)
def _generate_chunk(
    texts: List[str],
    ref_audio: str,
    ref_text: Optional[str],
    language: Optional[str],
    gen_config: OmniVoiceGenerationConfig,
    extra: dict,
):
    """Generate one sub-batch. Returns (audios, ref_text_actually_used).

    The voice clone prompt is rebuilt inside every GPU call on purpose: it holds
    CUDA tensors, and under ZeroGPU the CUDA context does not survive between
    calls. Encoding a few seconds of reference audio is cheap. The transcript is
    threaded back out so later chunks skip ASR.
    """
    prompt = model.create_voice_clone_prompt(
        ref_audio=ref_audio,
        ref_text=ref_text,
        preprocess_prompt=gen_config.preprocess_prompt,
    )
    audios = model.generate(
        text=texts,
        language=language,
        voice_clone_prompt=prompt,
        generation_config=gen_config,
        **extra,
    )
    return audios, prompt.ref_text


# ---------------------------------------------------------------------------
# Orchestration (CPU side)
# ---------------------------------------------------------------------------


def batch_generate(
    raw_texts,
    ref_audio,
    ref_text,
    language,
    batch_size,
    num_step,
    guidance_scale,
    denoise,
    speed,
    duration,
    preprocess_prompt,
    postprocess_output,
    progress=gr.Progress(),
):
    texts = parse_texts(raw_texts)
    if not texts:
        return [], "Please enter at least one line of text to synthesize."
    if len(texts) > MAX_ITEMS:
        return [], f"Too many lines ({len(texts)}). Maximum is {MAX_ITEMS}."
    if not ref_audio:
        return [], "Please upload a reference audio."

    gen_config = OmniVoiceGenerationConfig(
        num_step=int(num_step or 32),
        guidance_scale=float(guidance_scale) if guidance_scale is not None else 2.0,
        denoise=bool(denoise) if denoise is not None else True,
        preprocess_prompt=bool(preprocess_prompt),
        postprocess_output=bool(postprocess_output),
    )
    # speed and duration are generate() arguments, not generation-config fields.
    extra = {}
    if speed is not None and float(speed) != 1.0:
        extra["speed"] = float(speed)
    if duration is not None and float(duration) > 0:
        extra["duration"] = float(duration)

    lang = language if (language and language != "Auto") else None
    ref_text = (ref_text or "").strip() or None
    bs = max(1, int(batch_size or 8))

    # Sort by text length so each sub-batch is roughly homogeneous — the model
    # pads to the longest item in a batch, so mixing a 3-word line with a
    # 3-sentence one wastes compute. Original order is restored afterwards.
    order = sorted(range(len(texts)), key=lambda i: len(texts[i]))
    results: List[Optional[np.ndarray]] = [None] * len(texts)

    chunks = [order[i : i + bs] for i in range(0, len(order), bs)]
    start = time.time()
    done = 0
    for chunk in progress.tqdm(chunks, desc="Generating"):
        chunk_texts = [texts[i] for i in chunk]
        try:
            audios, ref_text = _generate_chunk(
                chunk_texts, ref_audio, ref_text, lang, gen_config, extra
            )
        except Exception as e:
            return [], f"Error after {done}/{len(texts)} items: {type(e).__name__}: {e}"
        for idx, audio in zip(chunk, audios):
            results[idx] = audio
        done += len(chunk)

    elapsed = time.time() - start

    # Persist to a fresh temp dir so Gradio can serve/download the files.
    out_dir = tempfile.mkdtemp(prefix="omnivoice_batch_")
    files = []
    total_audio = 0.0
    for i, audio in enumerate(results):
        assert audio is not None
        name = f"{i + 1:03d}_{slugify(texts[i])}.wav"
        path = os.path.join(out_dir, name)
        sf.write(path, audio, sampling_rate)
        total_audio += audio.shape[-1] / sampling_rate
        files.append({"path": path, "text": texts[i], "index": i + 1})

    zip_path = os.path.join(out_dir, "omnivoice_batch.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for f in files:
            zf.write(f["path"], os.path.basename(f["path"]))

    status = (
        f"Done. Generated {len(files)} clips ({total_audio:.1f}s of audio) "
        f"in {elapsed:.1f}s across {len(chunks)} batch(es) of up to {bs}.\n"
        f"Reference text used: {ref_text}"
    )
    return {"files": files, "zip": zip_path}, status


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

CSS = """
.compact-audio { max-height: 220px; }
.result-row { border-bottom: 1px solid var(--border-color-primary); padding: 4px 0; }
"""

with gr.Blocks(title="OmniVoice Batch TTS", css=CSS) as demo:
    gr.Markdown(
        """
# OmniVoice — Batch Voice Clone

Upload **one** reference audio (+ optional transcript), paste **one target text
per line**, and get one generated clip per line.

The reference voice prompt is encoded once per batch and reused for every line,
so this is meaningfully faster than generating the lines one at a time.
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            ref_audio = gr.Audio(
                label="Reference Audio / 参考音频",
                type="filepath",
                elem_classes="compact-audio",
            )
            gr.Markdown(
                "<span style='font-size:0.85em;color:#888;'>"
                "Recommended: 3–10 seconds of clean speech.</span>"
            )
            ref_text = gr.Textbox(
                label="Reference Text (optional) / 参考音频文本（可选）",
                lines=2,
                placeholder="Transcript of the reference audio. Leave empty to "
                "auto-transcribe via ASR.",
            )
            lang = gr.Dropdown(
                label="Language (optional) / 语种 (可选)",
                choices=_ALL_LANGUAGES,
                value="Auto",
                allow_custom_value=False,
                interactive=True,
                info="Applies to every line. Keep as Auto to auto-detect.",
            )
            texts_box = gr.Textbox(
                label="Texts to Synthesize — one per line / 待合成文本（每行一条）",
                lines=12,
                placeholder=(
                    "The first sentence to generate.\n"
                    "The second sentence to generate.\n"
                    "The third sentence to generate."
                ),
            )
            batch_size = gr.Slider(
                1,
                256,
                value=8,
                step=1,
                label="Batch Size",
                info="Lines generated per GPU call. Higher = faster, more VRAM. "
                "Lower this if you hit out-of-memory errors.",
            )
            with gr.Accordion("Generation Settings (optional)", open=False):
                speed = gr.Slider(
                    0.5,
                    1.5,
                    value=1.0,
                    step=0.05,
                    label="Speed",
                    info="1.0 = normal. >1 faster, <1 slower. Ignored if Duration is set.",
                )
                duration = gr.Number(
                    value=None,
                    label="Duration (seconds)",
                    info="Applies to every line — usually leave empty for batch runs.",
                )
                num_step = gr.Slider(
                    4,
                    64,
                    value=32,
                    step=1,
                    label="Inference Steps",
                    info="Default: 32. Lower = faster, higher = better quality.",
                )
                denoise = gr.Checkbox(label="Denoise", value=True)
                guidance_scale = gr.Slider(
                    0.0, 4.0, value=2.0, step=0.1, label="Guidance Scale (CFG)"
                )
                preprocess_prompt = gr.Checkbox(
                    label="Preprocess Prompt",
                    value=True,
                    info="Silence removal / trimming on the reference audio.",
                )
                postprocess_output = gr.Checkbox(
                    label="Postprocess Output",
                    value=True,
                    info="Remove long silences from generated audio.",
                )
            btn = gr.Button("Generate All / 批量生成", variant="primary")

        with gr.Column(scale=1):
            status = gr.Textbox(label="Status / 状态", lines=3)
            zip_out = gr.File(label="Download All (zip)", visible=False)
            results_state = gr.State([])

            @gr.render(inputs=results_state)
            def show_results(payload):
                if not payload:
                    gr.Markdown(
                        "<span style='color:#888;'>Generated clips will appear "
                        "here, one per input line.</span>"
                    )
                    return
                for item in payload["files"]:
                    with gr.Row(elem_classes="result-row"):
                        gr.Audio(
                            value=item["path"],
                            label=f"{item['index']}. {item['text'][:80]}",
                            type="filepath",
                        )

    def _run(*args):
        payload, msg = batch_generate(*args)
        if not payload:
            return [], gr.update(visible=False), msg
        return payload, gr.update(value=payload["zip"], visible=True), msg

    UI_INPUTS = [
        texts_box,
        ref_audio,
        ref_text,
        lang,
        batch_size,
        num_step,
        guidance_scale,
        denoise,
        speed,
        duration,
        preprocess_prompt,
        postprocess_output,
    ]

    btn.click(
        _run,
        inputs=UI_INPUTS,
        outputs=[results_state, zip_out, status],
        api_name=False,  # returns a gr.State; the API uses /batch_generate below
    )

    # -----------------------------------------------------------------
    # Programmatic API endpoint
    # -----------------------------------------------------------------
    # Hidden components, wired to their own event so the public API returns
    # plain files + a status string instead of the UI's gr.State payload.
    # Argument names below are the keyword names gradio_client will expose.
    api_files = gr.File(file_count="multiple", visible=False)
    api_status = gr.Textbox(visible=False)
    api_btn = gr.Button(visible=False)

    def batch_generate_api(
        texts,
        ref_audio,
        ref_text,
        language,
        batch_size,
        num_step,
        guidance_scale,
        denoise,
        speed,
        duration,
        preprocess_prompt,
        postprocess_output,
        progress=gr.Progress(),
    ):
        """Generate one clip per line of `texts`, cloning the voice in `ref_audio`.

        Returns (list of wav filepaths in input-line order, status string).
        On failure the file list is empty and the status describes the error.
        """
        payload, msg = batch_generate(
            texts,
            ref_audio,
            ref_text,
            language,
            batch_size,
            num_step,
            guidance_scale,
            denoise,
            speed,
            duration,
            preprocess_prompt,
            postprocess_output,
            progress=progress,
        )
        if not payload:
            return [], msg
        return [f["path"] for f in payload["files"]], msg

    api_btn.click(
        batch_generate_api,
        inputs=UI_INPUTS,
        outputs=[api_files, api_status],
        api_name="batch_generate",
    )


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