import spaces  # Must precede torch: ZeroGPU installs its CUDA hooks here.

import math
import tempfile
import time
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from scipy.signal import resample_poly
from transformers import WhisperForConditionalGeneration, WhisperProcessor

MODEL_ID = "BuzzASR/hausa"
MAX_SECONDS = 30
SAMPLE_RATE = 16000

processor = WhisperProcessor.from_pretrained(MODEL_ID)
model = WhisperForConditionalGeneration.from_pretrained(
    MODEL_ID, dtype=torch.float16, attn_implementation="sdpa",
).to("cuda").eval()


def prepare_audio(audio: tuple[int, np.ndarray] | None) -> tuple[np.ndarray, float]:
    """Validate a short recording and resample it to mono 16 kHz."""
    if audio is None:
        raise gr.Error("Record your voice or upload an audio file first.")
    sample_rate, raw = audio
    samples = np.asarray(raw)
    if sample_rate <= 0 or samples.ndim not in (1, 2) or samples.size == 0:
        raise gr.Error("This recording is empty or invalid. Please record again.")
    duration = samples.shape[0] / sample_rate
    if duration > MAX_SECONDS:
        raise gr.Error("Please use a clip of 30 seconds or less. Longer recordings are not truncated.")
    if np.issubdtype(samples.dtype, np.integer):
        limits = np.iinfo(samples.dtype)
        if limits.min == 0:
            midpoint = (limits.max + 1) / 2
            samples = (samples.astype(np.float32) - midpoint) / midpoint
        else:
            samples = samples.astype(np.float32) / max(abs(limits.min), limits.max)
    else:
        samples = samples.astype(np.float32)
    if samples.ndim == 2:
        samples = samples.mean(axis=1)
    if not np.isfinite(samples).all():
        raise gr.Error("The recording contains invalid audio samples.")
    if sample_rate != SAMPLE_RATE:
        divisor = math.gcd(sample_rate, SAMPLE_RATE)
        samples = resample_poly(samples, SAMPLE_RATE // divisor, sample_rate // divisor)
    return np.ascontiguousarray(samples, dtype=np.float32), duration


@spaces.GPU(duration=60)
def transcribe(audio: tuple[int, np.ndarray] | None) -> tuple[str, str | None, str]:
    """Transcribe a microphone recording or uploaded Hausa clip of up to 30 seconds.

    Returns the unedited transcript, a UTF-8 download, and processing statistics.
    """
    samples, duration = prepare_audio(audio)
    if float(np.sqrt(np.mean(samples ** 2))) < 1e-5:
        return "", None, "No audible speech detected. Check your microphone and try again."
    started = time.perf_counter()
    inputs = processor(
        samples, sampling_rate=SAMPLE_RATE, return_tensors="pt", return_attention_mask=True,
    )
    with torch.inference_mode():
        ids = model.generate(
            inputs.input_features.to(device="cuda", dtype=torch.float16),
            attention_mask=inputs.attention_mask.to("cuda"),
            num_beams=1, no_repeat_ngram_size=3, repetition_penalty=1.2,
        )
    transcript = processor.batch_decode(ids, skip_special_tokens=True)[0].strip()
    elapsed = time.perf_counter() - started
    if not transcript:
        return "", None, "No transcript returned. Try a clearer recording."
    # Gradio copies this output into its managed cache before serving it.
    Path(demo.GRADIO_CACHE).mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", suffix=".txt", dir=demo.GRADIO_CACHE, delete=False) as output:
        output.write(transcript + "\n")
        transcript_path = output.name
    stats = f"**Audio:** {duration:.1f}s · **Processing:** {elapsed:.1f}s · **Words:** {len(transcript.split())}"
    return transcript, transcript_path, stats


CSS = """
.gradio-container {max-width: 1060px !important; margin: auto;}
#hero {padding: 24px 0 12px;}
#hero h1 {font-size: 2.6rem; letter-spacing: -0.04em; margin-bottom: 8px;}
#hero p {font-size: 1.05rem; max-width: 720px;}
#transcript textarea {font-size: 1.15rem; line-height: 1.8;}
"""

with gr.Blocks(title="Hausa Voice Lab", delete_cache=(300, 3600)) as demo:
    gr.Markdown(
        "# Hausa Voice Lab\n"
        "### Ka yi magana. Ka ga rubutun Hausa.\n"
        "Record a short Hausa clip and see what BuzzASR hears. "
        "Listen back, compare names and spelling, and keep the transcript.",
        elem_id="hero",
    )
    with gr.Row():
        with gr.Column(scale=1):
            audio = gr.Audio(
                sources=["microphone", "upload"], type="numpy", label="Your recording",
            )
            gr.Markdown("**1.** Record or upload up to 30 seconds.\n\n**2.** Finish recording, then transcribe.")
            run = gr.Button("Transcribe Hausa", variant="primary", size="lg")
        with gr.Column(scale=1):
            transcript = gr.Textbox(
                label="Hausa transcript", placeholder="Your words will appear here…",
                lines=8, interactive=False, buttons=["copy"], elem_id="transcript",
            )
            stats = gr.Markdown("Ready when you are.")
            download = gr.File(label="Download transcript (.txt)", interactive=False)
    gr.ClearButton([audio, transcript, download, stats], value="Clear recording and transcript")
    with gr.Accordion("Try a short accuracy check", open=False):
        gr.Markdown(
            "Say a sentence containing your name and hometown, then check every word. "
            "For example: **Suna na Auwal. Na zo ne daga garin Ringim.**\n\n"
            "This is a suggested phrase to record, not a prerecorded or verified model result."
        )
    gr.Markdown(
        "Powered by [BuzzASR/hausa](https://huggingface.co/BuzzASR/hausa). "
        "Hausa-focused; English and names may be inaccurate. No automatic name corrections.\n\n"
        "Audio is uploaded to this Space and temporarily cached for processing. "
        "Free GPU access has queues and usage limits."
    )
    run.click(transcribe, inputs=audio, outputs=[transcript, download, stats], api_name="transcribe", concurrency_limit=1)

demo.queue(max_size=12).launch(
    theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="blue"),
    css=CSS, mcp_server=True, max_file_size="20mb",
)