import os

os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")

import spaces  # noqa: E402  (must precede torch)

import json
import subprocess
import sys
import sysconfig
import time
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from huggingface_hub import snapshot_download

MODEL_ID = "LocalDoc/speakwrite-azerbaijani"
HERE = Path(__file__).resolve().parent
BALBAL = HERE / "balbal"  # balbal @ 1be9f88b77c63fbafa7713c6e46e79847f6c390e (Apache-2.0), refiner.json from the model repo
NATIVE = BALBAL / "inference" / "native"
MAX_SECONDS = 300
CHUNK_SECONDS = 30.0

# ---- build balbal's native C++/pybind11 KenLM decoders once (flashlight-text 0.0.7 ABI) ----
if not (NATIVE / ("ctc_fast_decoder" + sysconfig.get_config_var("EXT_SUFFIX"))).exists():
    print("Building balbal native decoders ...", flush=True)
    subprocess.run([sys.executable, str(BALBAL / "inference" / "build_native.py")], check=True,
                   env=dict(os.environ, CXXFLAGS="-O3"))

sys.path[:0] = [str(BALBAL / "inference"), str(BALBAL)]
from balbal import models as M  # noqa: E402
from balbal.audio import read_audio  # noqa: E402
from decoder import KenLMDecoder  # noqa: E402

# ---- weights ----
model_dir = Path(snapshot_download(MODEL_ID))
settings = dict(M.decoder_config("balbal-v2-manas"), **json.loads((model_dir / "decoder.json").read_text()))

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
net = M.build_model("balbal-v2-manas", model_dir / "tokenizer", model_dir / "model.pt", device="cpu")
net = net.to(DEVICE).eval()
torch.set_float32_matmul_precision("highest")

lm_decoder = KenLMDecoder(model_dir / "lm", settings, fast=True)
sp = lm_decoder.sp
VOCAB = lm_decoder.vocab_size
print(f"Model + KenLM loaded (device={DEVICE}).", flush=True)


def _forward(wave: np.ndarray) -> np.ndarray:
    """Mono float32 16 kHz waveform (<= 30 s) -> [frames, V + 1] natural-log CTC probabilities (blank last)."""
    wave = np.ascontiguousarray(wave, dtype=np.float32)
    if len(wave) < 4000:  # the conv subsampler needs >= 0.25 s
        wave = np.pad(wave, (0, 4000 - len(wave)))
    with torch.inference_mode():
        lp, n, _ = net(input_signal=torch.from_numpy(wave)[None].to(DEVICE),
                       input_signal_length=torch.tensor([len(wave)], device=DEVICE))
    return lp[0, : int(n[0])].float().cpu().numpy()


def _greedy(lp: np.ndarray) -> str:
    ids = lp.argmax(axis=1)
    out, prev = [], -1
    for t in ids:
        t = int(t)
        if t != prev and t < VOCAB:
            out.append(t)
        prev = t
    return sp.decode(out)


@spaces.GPU(duration=20)
def transcribe(audio_path: str, use_lm: bool = True, hotwords: str = ""):
    """Transcribe Azerbaijani speech to text with speakwrite-azerbaijani (FastConformer-CTC + 6-gram KenLM).

    Args:
        audio_path: path to an audio file (any format ffmpeg can read); resampled to 16 kHz mono. Max 5 minutes.
        use_lm: decode with the 6-gram token KenLM beam search (lower WER). If False, greedy CTC argmax.
        hotwords: optional comma-separated names/words expected in the audio; they get a small bonus during LM decoding.

    Returns:
        The transcript (cased, punctuated, numbers as words), word timestamps, and timing info.
    """
    if not audio_path:
        raise gr.Error("Please upload or record an audio clip.")
    wave = read_audio(audio_path)
    duration = len(wave) / 16000
    if duration > MAX_SECONDS:
        raise gr.Error(f"Audio is {duration:.0f} s long; please keep it under {MAX_SECONDS // 60} minutes.")
    hot = [w.strip() for w in (hotwords or "").split(",") if w.strip()] or None
    t0 = time.perf_counter()
    step = int(CHUNK_SECONDS * 16000)
    texts, words = [], []
    for i in range(0, len(wave), step):
        lp = _forward(wave[i : i + step])
        if use_lm:
            out = lm_decoder.decode_words(lp, hot, i / 16000)
            words += [[w["word"], round(w["start"], 2), round(w["end"], 2)] for w in out["words"]]
            text = out["text"]
        else:
            text = _greedy(lp)
        if text:
            texts.append(text)
    elapsed = time.perf_counter() - t0
    info = (f"Audio: {duration:.1f} s · inference: {elapsed:.2f} s "
            f"({duration / elapsed if elapsed else 0:.0f}× real time) · decoder: {'KenLM beam search' if use_lm else 'greedy CTC'}")
    return " ".join(texts), words, info


CSS = """
#col-container { max-width: 1000px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎙️ speakwrite-azerbaijani\n"
            "Azerbaijani speech recognition: FastConformer-CTC with a 3-pass loop refiner, product-key acoustic memory "
            "and a 6-gram token KenLM. Output is cased and punctuated, with numbers written as words.  \n"
            "[Model card](https://huggingface.co/LocalDoc/speakwrite-azerbaijani) · "
            "[balbal](https://github.com/erayyap/balbal) · "
            "*speakwrite-azerbaijani by LocalDoc* — CC BY-NC 4.0 (non-commercial use only)."
        )
        with gr.Row():
            with gr.Column():
                audio = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Azerbaijani speech")
                run = gr.Button("Transcribe", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    use_lm = gr.Checkbox(value=True, label="Use KenLM beam search (uncheck for greedy CTC)")
                    hotwords = gr.Textbox(value="", label="Hotwords (comma-separated names/words to favour; LM only)",
                                          placeholder="e.g. Bakı, Nizami")
            with gr.Column():
                text = gr.Textbox(label="Transcript", lines=6, buttons=["copy"])
                info = gr.Markdown()
                words = gr.Dataframe(headers=["word", "start (s)", "end (s)"], label="Word timestamps", interactive=False)

        run.click(transcribe, inputs=[audio, use_lm, hotwords], outputs=[text, words, info])

        gr.Examples(
            examples=[["examples/suni_zeka.wav"], ["examples/honkonq.wav"],
                      ["examples/medeni_sok.wav"], ["examples/quslar.wav"]],
            inputs=[audio],
            outputs=[text, words, info],
            fn=transcribe,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples (FLEURS Azerbaijani test set, CC BY 4.0)",
        )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)