"""Gradio Space: forced alignment with Qwen3-ForcedAligner-0.6B.

Serves both:
  1. A usable web UI for interactive alignment.
  2. A programmatic API endpoint (`/align`), auto-exposed by Gradio, that
     other services can call via `gradio_client.Client` — see
     `examples/client_example.py`.
"""

from __future__ import annotations

# `spaces` must be imported before `torch` (imported transitively via
# `aligner`) so it can patch CUDA initialization for ZeroGPU Spaces. On
# non-ZeroGPU hardware (or local dev), `spaces.GPU` is a harmless no-op.
import spaces

import gradio as gr

from aligner import MAX_AUDIO_SECONDS, SUPPORTED_LANGUAGES, align

DATAFRAME_HEADERS = ["#", "Text", "Start (s)", "End (s)"]

EXAMPLE_AUDIO_URL = (
    "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-ASR-Repo/asr_zh.wav"
)
EXAMPLE_TEXT = "甚至出现交易几乎停滞的情况。"


@spaces.GPU(duration=120)
def run_alignment(audio, text: str, language: str):
    """Event handler for the UI button, and the function exposed as the API.

    Returns (dataframe_rows, raw_json) so the UI shows both a readable
    table and the full structured payload; API callers get the same tuple.

    `@spaces.GPU` is required on ZeroGPU Spaces: it's how the ZeroGPU
    scheduler detects which functions need a GPU attached and allocates one
    for the duration of the call. Model loading happens lazily inside
    `aligner.get_model()`, on first call to `align()` below, so it also
    runs inside this GPU-attached window (required — CUDA calls made
    outside a `@spaces.GPU` call fail on ZeroGPU hardware).
    """
    try:
        spans = align(audio=audio, text=text, language=language)
    except Exception as exc:  # surface a clean error instead of a traceback
        raise gr.Error(str(exc)) from exc

    rows = [[s.index, s.text, round(s.start_time, 3), round(s.end_time, 3)] for s in spans]
    raw = [
        {"index": s.index, "text": s.text, "start_time": s.start_time, "end_time": s.end_time}
        for s in spans
    ]
    return rows, raw


with gr.Blocks(title="Qwen3 Forced Aligner") as demo:
    gr.Markdown(
        """
        # Qwen3 Forced Aligner
        Align a text transcript to its audio and get per-unit start/end timestamps,
        powered by [Qwen/Qwen3-ForcedAligner-0.6B](https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B).

        Upload or record audio, paste the matching transcript, pick the language, and align.
        This UI also doubles as an API — see the "Use via API" link at the bottom of the page.
        """
    )

    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label="Audio",
            )
            text_input = gr.Textbox(
                label="Text to align",
                placeholder="Enter the transcript that matches the audio...",
                lines=4,
            )
            language_input = gr.Dropdown(
                choices=SUPPORTED_LANGUAGES,
                value="English",
                label="Language",
            )
            gr.Markdown(
                f"_Max audio length: ~{MAX_AUDIO_SECONDS // 60} minutes per request._"
            )
            align_button = gr.Button("Align", variant="primary")

        with gr.Column():
            output_table = gr.Dataframe(
                headers=DATAFRAME_HEADERS,
                datatype=["number", "str", "number", "number"],
                label="Aligned spans",
                wrap=True,
            )
            with gr.Accordion("Raw JSON output", open=False):
                output_json = gr.JSON(label="Raw result")

    gr.Examples(
        examples=[[EXAMPLE_AUDIO_URL, EXAMPLE_TEXT, "Chinese"]],
        inputs=[audio_input, text_input, language_input],
    )

    align_button.click(
        fn=run_alignment,
        inputs=[audio_input, text_input, language_input],
        outputs=[output_table, output_json],
        api_name="align",
        concurrency_limit=1,  # one model instance on one GPU
    )

demo.queue()

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