"""
ScriptFlow inference service — Hausa ASR + Hausa->English MT on a free HF Space.

Exists so the Render backend stays a thin 512MB web process: Whisper-small plus
an MT model needs ~2GB resident, which the free Render plan cannot hold. A free
Space (2 vCPU / 16GB) can, and both are reached over plain HTTP.

  GET  /          -> {"status", "asr", "mt"}   (never blocks on model load)
  POST /asr       -> {"text", "duration", "words": [{word,start,end,speaker}]}
  POST /translate -> {"translations": [...]}

/asr returns word-level timings in exactly the shape Chirp 3 produced, so the
backend's cue splitter consumes it unchanged.
"""

import os
import time
import threading

import numpy as np
from fastapi import FastAPI, Request, HTTPException, Header
from pydantic import BaseModel

ASR_MODEL = os.environ.get("ASR_MODEL", "NCAIR1/Hausa-ASR")
MT_MODEL = os.environ.get("MT_MODEL", "Helsinki-NLP/opus-mt-ha-en")
# NCAIR1/Hausa-ASR is gated: the token's account must have accepted the licence
# on the model page, or from_pretrained 403s.
HF_TOKEN = os.environ.get("HF_TOKEN") or None
# Free Spaces are world-reachable. A shared secret keeps strangers off the CPU
# budget; unset means open, which is fine for local testing only.
SERVICE_TOKEN = os.environ.get("SERVICE_TOKEN") or None

SAMPLE_RATE = 16000
# Whisper's context is 30s. The pipeline slides this window with overlap so a
# word straddling a boundary is still decoded once, correctly.
CHUNK_LENGTH_S = float(os.environ.get("CHUNK_LENGTH_S", "30"))
STRIDE_LENGTH_S = float(os.environ.get("STRIDE_LENGTH_S", "5"))
MT_BATCH = int(os.environ.get("MT_BATCH", "16"))
# ZeroGPU hands out a GPU only inside @spaces.GPU functions, and refuses to
# start a Space that declares none ("No @spaces.GPU function detected during
# startup"). Importing spaces also tells us we are on that hardware, so the
# device default follows it.
try:
    import spaces
    _GPU = spaces.GPU(duration=int(os.environ.get("GPU_DURATION", "120")))
    _ON_ZEROGPU = True
except Exception:                             # local runs, or a plain CPU host
    _ON_ZEROGPU = False

    def _GPU(fn):
        return fn

DEVICE = os.environ.get("ASR_DEVICE", "cuda" if _ON_ZEROGPU else "cpu")

# Gradio owns the process on a Gradio-SDK Space, but it is a FastAPI app
# underneath — so the REST routes below are the real interface and the little UI
# mounted at /ui is only there to give the Space a face (and to make it obvious
# at a glance whether the models have finished loading).
app = FastAPI()


@app.exception_handler(Exception)
def _unhandled(request, exc):
    """
    Return the traceback in the response body. Reading a Space's logs needs an
    HF token, so an opaque 500 is genuinely hard to diagnose from the client
    that is calling this. Nothing here handles user data worth hiding.
    """
    import traceback
    from fastapi.responses import JSONResponse
    tb = traceback.format_exc()
    print(tb, flush=True)
    return JSONResponse(status_code=500,
                        content={"error": f"{type(exc).__name__}: {exc}",
                                 "traceback": tb.splitlines()[-12:],
                                 "device": DEVICE})

# --------------------------------------------------------------- MODEL LOAD ---
# Loading takes minutes on a cold Space (weights download + CPU init). Doing it
# at import time would make the Space fail its health check and restart-loop, so
# models load lazily behind a lock and / reports progress instead.
_lock = threading.Lock()
_models: dict = {"asr": None, "mt": None}
_errors: dict = {"asr": None, "mt": None}


def _load(kind: str):
    if _models[kind] is not None:
        return _models[kind]
    with _lock:
        if _models[kind] is not None:
            return _models[kind]
        from transformers import pipeline
        t0 = time.time()
        try:
            if kind == "asr":
                obj = pipeline(
                    "automatic-speech-recognition",
                    model=ASR_MODEL,
                    token=HF_TOKEN,
                    chunk_length_s=CHUNK_LENGTH_S,
                    stride_length_s=STRIDE_LENGTH_S,
                    device=DEVICE,
                )
            else:
                # transformers 5 dropped the "translation" pipeline task, so the
                # seq2seq model is driven directly. Beam search, no sampling —
                # the same cue gives the same English on every run, which is the
                # whole reason for moving off Gemini.
                from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
                tok = AutoTokenizer.from_pretrained(MT_MODEL, token=HF_TOKEN)
                mdl = AutoModelForSeq2SeqLM.from_pretrained(MT_MODEL, token=HF_TOKEN)
                obj = (tok, mdl.to(DEVICE).eval())
        except Exception as e:
            _errors[kind] = f"{type(e).__name__}: {e}"
            raise
        print(f"loaded {kind} in {time.time() - t0:.1f}s", flush=True)
        _models[kind] = obj
        return obj


def _auth(token):
    if SERVICE_TOKEN and token != SERVICE_TOKEN:
        raise HTTPException(status_code=401, detail="bad service token")


@app.get("/health")
def health():
    return {
        "status": "ok",
        "asr": {"model": ASR_MODEL, "loaded": _models["asr"] is not None,
                "error": _errors["asr"]},
        "mt": {"model": MT_MODEL, "loaded": _models["mt"] is not None,
               "error": _errors["mt"]},
    }


# ----------------------------------------------------------------------- ASR ---
def _decode(raw: bytes) -> np.ndarray:
    """
    Any container -> float32 mono @16k. The backend already sends 16k mono WAV,
    but decoding defensively costs nothing and keeps the service reusable.
    """
    try:
        from transformers.pipelines.audio_utils import ffmpeg_read
    except ImportError:                       # moved in transformers 5
        from transformers.audio_utils import ffmpeg_read
    return ffmpeg_read(raw, SAMPLE_RATE)


def _words_from_chunks(chunks, audio_secs: float):
    """
    Whisper chunks -> the backend's word contract.

    With return_timestamps="word" each chunk is already one word. With the
    phrase-level fallback a chunk is several words, so its span is divided
    across them by character length — an approximation, but one that keeps cue
    boundaries near the speech instead of collapsing them onto a single instant.
    """
    words = []
    for ch in chunks or []:
        text = (ch.get("text") or "").strip()
        if not text:
            continue
        ts = ch.get("timestamp") or (None, None)
        start = ts[0]
        end = ts[1] if len(ts) > 1 else None
        if start is None:
            continue
        start = float(start)
        if audio_secs:
            start = min(max(start, 0.0), audio_secs)
        if end is not None:
            end = float(end)
            end = min(max(end, start), audio_secs) if audio_secs else max(end, start)

        parts = text.split()
        if len(parts) <= 1:
            # A None end is meaningful downstream: the backend infers a real end
            # from the next word rather than inventing a duration here.
            words.append({"word": text, "start": start, "end": end, "speaker": None})
            continue

        span = (end - start) if end is not None else None
        total = sum(len(p) for p in parts) or 1
        cursor = start
        for p in parts:
            if span is None:
                words.append({"word": p, "start": cursor, "end": None, "speaker": None})
                cursor += 0.3
            else:
                width = span * (len(p) / total)
                words.append({"word": p, "start": cursor, "end": cursor + width,
                              "speaker": None})
                cursor += width
    return words


def _device_fallback() -> bool:
    """
    Drop to CPU and force a reload. ZeroGPU only grants a device inside its
    decorated functions, and the grant can fail for reasons that have nothing
    to do with the model — quota, a worker thread, a cold pool. Falling back is
    slower but keeps the service answering instead of 500ing.
    """
    global DEVICE
    if DEVICE == "cpu":
        return False
    print(f"GPU path failed on {DEVICE} — reloading on CPU", flush=True)
    DEVICE = "cpu"
    _models["asr"] = _models["mt"] = None
    _errors["asr"] = _errors["mt"] = None
    return True


def _asr_impl(audio, gen):
    """Decoding, on the GPU when ZeroGPU has granted one."""
    pipe = _load("asr")
    try:
        return pipe(audio.copy(), return_timestamps="word", generate_kwargs=gen), True
    except Exception as e:
        # Word timings need alignment_heads in the generation config. Fine-tunes
        # frequently drop them; phrase-level timings still make usable cues.
        print(f"word timestamps unavailable ({type(e).__name__}: {e}) — phrase level",
              flush=True)
        return pipe(audio.copy(), return_timestamps=True, generate_kwargs=gen), False


_asr_gpu = _GPU(_asr_impl)


def _asr_run(audio, gen):
    try:
        return _asr_gpu(audio, gen)
    except Exception as e:
        print(f"asr on {DEVICE} failed ({type(e).__name__}: {e})", flush=True)
        if not _device_fallback():
            raise
        return _asr_impl(audio, gen)


@app.post("/asr")
async def asr(request: Request, x_service_token: str = Header(default=None)):
    _auth(x_service_token)
    raw = await request.body()
    if not raw:
        raise HTTPException(status_code=400, detail="empty audio body")

    audio = _decode(raw)
    audio_secs = len(audio) / float(SAMPLE_RATE)

    # language/task are forced because a fine-tune sometimes ships a generation
    # config still defaulting to English transcription, which silently produces
    # garbage on Hausa audio.
    gen = {"language": "ha", "task": "transcribe"}

    t0 = time.time()
    out, word_level = _asr_run(audio, gen)

    words = _words_from_chunks(out.get("chunks"), audio_secs)
    print(f"asr {audio_secs:.1f}s audio -> {len(words)} words in {time.time() - t0:.1f}s "
          f"({'word' if word_level else 'phrase'} timings)", flush=True)

    return {
        "text": (out.get("text") or "").strip(),
        "duration": audio_secs,
        "wordLevel": word_level,
        "words": words,
    }


# --------------------------------------------------------------- TRANSLATION ---
class TranslateIn(BaseModel):
    texts: list[str]


def _mt_impl(texts):
    import torch
    tok, mdl = _load("mt")
    out = []
    for start in range(0, len(texts), MT_BATCH):
        chunk = texts[start:start + MT_BATCH]
        enc = tok(chunk, return_tensors="pt", padding=True, truncation=True,
                  max_length=512).to(DEVICE)
        with torch.inference_mode():
            gen = mdl.generate(**enc, num_beams=4, max_new_tokens=256)
        out.extend(tok.batch_decode(gen, skip_special_tokens=True))
    return out


_mt_gpu = _GPU(_mt_impl)


def _mt_run(texts):
    try:
        return _mt_gpu(texts)
    except Exception as e:
        print(f"mt on {DEVICE} failed ({type(e).__name__}: {e})", flush=True)
        if not _device_fallback():
            raise
        return _mt_impl(texts)


@app.post("/translate")
def translate(body: TranslateIn, x_service_token: str = Header(default=None)):
    _auth(x_service_token)
    texts = body.texts or []
    if not texts:
        return {"translations": []}

    # Blank inputs are held out and reinserted: MarianMT will happily emit a
    # hallucinated sentence for an empty string.
    idx = [i for i, t in enumerate(texts) if (t or "").strip()]
    out = [""] * len(texts)
    if idx:
        res = _mt_run([texts[i] for i in idx])
        for i, r in zip(idx, res):
            out[i] = (r or "").strip()
    return {"translations": out}


# ------------------------------------------------------------------- LAUNCH ---
# A Gradio-SDK Space expects Gradio itself to bind the port. Two earlier shapes
# both failed here: running our own uvicorn raced whatever the runner had
# already put on 7860 ("address already in use"), and defining a bare ASGI app
# with no listener left the container with nothing serving at all.
#
# So Gradio launches, and the REST routes are grafted onto the server it starts.
# They are *prepended*, because Gradio registers catch-all routes for its own
# frontend that would otherwise shadow /asr and /translate.
def _ui_check():
    import json as _json
    return _json.dumps(health(), indent=2)


_UI_TEXT = (
    "## ScriptFlow inference service\n\n"
    "Hausa ASR + Hausa->English MT. The REST API is the real interface:\n\n"
    "- `POST /asr` — raw audio bytes as the body\n"
    "- `POST /translate` — `{\"texts\": [...]}`\n"
    "- `GET /health` — model load state\n\n"
    "Both POSTs take an `X-Service-Token` header. The first call after a cold "
    "start downloads ~1GB of weights and takes several minutes."
)


def _build_ui():
    import gradio as gr
    with gr.Blocks(title="ScriptFlow Hausa inference") as demo:
        gr.Markdown(_UI_TEXT)
        out = gr.Code(label="GET /health", language="json")
        gr.Button("Check status").click(_ui_check, inputs=None, outputs=out)
    return demo


if __name__ == "__main__":
    import gradio as gr

    demo = _build_ui()
    # prevent_thread_lock so this returns and the routes can be attached; the
    # process is held open by block_thread() at the end instead.
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        prevent_thread_lock=True,
        show_error=True,
        # SSR is Gradio 6's default and puts a Node proxy on the public port,
        # serving Python on 7861 behind it. Routes grafted onto the Python app
        # are then unreachable from outside, so the REST API has to have the
        # public port to itself.
        ssr_mode=False,
    )
    demo.app.router.routes[0:0] = app.router.routes
    print("REST routes attached: " +
          ", ".join(sorted(r.path for r in app.router.routes if hasattr(r, "methods"))),
          flush=True)
    # block_thread() returned immediately here and the container exited; an
    # explicit wait keeps the process alive for as long as the server runs.
    try:
        demo.block_thread()
    except KeyboardInterrupt:
        pass
    else:
        threading.Event().wait()