"""Forced aligner sidecar: where does each word of a known text start and end in its audio?

POST /align        multipart: audio (wav/mp3/flac file), text, language, [start, end] seconds
POST /align/batch  multipart: audio, spans (JSON [{start,end,text}]), language -- many sentences, one upload
GET  /health

Times are seconds in the uploaded audio. Words are the whitespace-separated tokens of the text, in order; a
token with no letters (digits, punctuation) cannot be aligned and gets the gap between its neighbours
(aligned=false). Model: torchaudio MMS_FA (weights are fetched once into TORCH_HOME, a volume, never baked in).
No secrets, no auth: the service is reachable only on the internal Docker network.
"""
from __future__ import annotations

import asyncio
import io
import json
import os
import re
import threading
import time
import unicodedata
from typing import Any

import numpy as np
from fastapi import FastAPI, File, Form, HTTPException, UploadFile

SAMPLE_RATE = 16000
MAX_AUDIO_BYTES = int(os.environ.get("ALIGNER_MAX_AUDIO_BYTES", 40 * 1024 * 1024))
MAX_SPAN_SECONDS = float(os.environ.get("ALIGNER_MAX_SPAN_SECONDS", 120))
MAX_AUDIO_SECONDS = float(os.environ.get("ALIGNER_MAX_AUDIO_SECONDS", 900))
MAX_SPANS = int(os.environ.get("ALIGNER_MAX_SPANS", 400))
MAX_TEXT_CHARS = int(os.environ.get("ALIGNER_MAX_TEXT_CHARS", 20000))
REQUEST_TIMEOUT = float(os.environ.get("ALIGNER_REQUEST_TIMEOUT", 110))
QUEUE_WAIT = float(os.environ.get("ALIGNER_QUEUE_WAIT", 30))
DEVICE = os.environ.get("ALIGNER_DEVICE", "cpu")
_LANG_RE = re.compile(r"^[A-Za-z]{2,3}(-[A-Za-z0-9]{2,8})?$")

app = FastAPI(title="aligner", docs_url=None, redoc_url=None, openapi_url=None)
_gate = asyncio.Semaphore(1)            # one alignment at a time: the model is the scarce resource
_state: dict[str, Any] = {"bundle": None, "model": None, "tok": None, "aligner": None, "error": None}
_load_lock = threading.Lock()


# ── text ───────────────────────────────────────────────────────────────────────

def normalize_word(word: str) -> str:
    """Romanize one token to the a-z' alphabet of the MMS_FA model (umlauts and accents folded)."""
    w = unicodedata.normalize("NFKD", word.lower().replace("ß", "ss").replace("æ", "ae").replace("œ", "oe"))
    w = "".join(c for c in w if not unicodedata.combining(c))
    return re.sub(r"[^a-z']", "", w)


def tokenize(text: str) -> list[str]:
    return text.split()


# ── model ──────────────────────────────────────────────────────────────────────

def _load() -> None:
    with _load_lock:
        if _state["model"] is not None:
            return
        import torch                                    # imported lazily: /health works before the model is loaded
        import torchaudio
        torch.set_num_threads(int(os.environ.get("ALIGNER_THREADS", "4")))
        bundle = torchaudio.pipelines.MMS_FA
        model = bundle.get_model(with_star=False).to(DEVICE).eval()
        _state.update(bundle=bundle, model=model, tok=bundle.get_tokenizer(), aligner=bundle.get_aligner(), error=None)


def align_samples(samples: np.ndarray, words: list[str]) -> list[tuple[float, float, bool]]:
    """Per whitespace token (start, end, aligned) in seconds relative to `samples` (16 kHz mono float32)."""
    import torch
    norm = [normalize_word(w) for w in words]
    keep = [i for i, n in enumerate(norm) if n]
    out: list[list[Any]] = [[None, None, False] for _ in words]
    if keep and len(samples) >= 400:
        wav = torch.from_numpy(samples)[None].to(DEVICE)
        with torch.inference_mode():
            emission, _ = _state["model"](wav)
            spans = _state["aligner"](emission[0], _state["tok"]([norm[i] for i in keep]))
        frame = len(samples) / emission.shape[1] / SAMPLE_RATE
        for i, sp in zip(keep, spans):
            out[i] = [sp[0].start * frame, sp[-1].end * frame, True]
    # tokens without letters: the gap between their neighbours
    total = len(samples) / SAMPLE_RATE
    for i, o in enumerate(out):
        if o[2]:
            continue
        prev_end = next((out[j][1] for j in range(i - 1, -1, -1) if out[j][2]), 0.0)
        next_start = next((out[j][0] for j in range(i + 1, len(out)) if out[j][2]), total)
        o[0], o[1] = prev_end, max(prev_end, next_start)
    return [(float(a), float(b), bool(c)) for a, b, c in out]


# ── audio ──────────────────────────────────────────────────────────────────────

def decode(data: bytes) -> np.ndarray:
    """Any libsndfile format (wav, flac, mp3) to 16 kHz mono float32."""
    import soundfile as sf
    try:
        with sf.SoundFile(io.BytesIO(data)) as f:
            if f.frames / max(1, f.samplerate) > MAX_AUDIO_SECONDS:
                raise HTTPException(413, f"audio longer than {MAX_AUDIO_SECONDS:.0f} s")
            x = f.read(dtype="float32", always_2d=True)
            sr = f.samplerate
    except HTTPException:
        raise
    except Exception:
        raise HTTPException(415, "audio is not a readable wav/mp3/flac file")
    x = x.mean(axis=1)
    if sr != SAMPLE_RATE:
        import torch
        import torchaudio.functional as F
        x = F.resample(torch.from_numpy(x), sr, SAMPLE_RATE).numpy()
    return np.ascontiguousarray(x, dtype=np.float32)


def _check_language(language: str) -> None:
    if not _LANG_RE.match(language or ""):
        raise HTTPException(422, "language must be an ISO code such as 'de' or 'en'")


def _check_text(text: str, strict: bool = True) -> list[str]:
    if not isinstance(text, str) or len(text) > MAX_TEXT_CHARS:
        raise HTTPException(422, f"text must be a string of at most {MAX_TEXT_CHARS} characters")
    words = tokenize(text)
    if not words:
        raise HTTPException(422, "text is empty")
    if strict and not any(normalize_word(w) for w in words):
        raise HTTPException(422, "text has no letters the model can align (latin alphabets only)")
    return words


def _slice(x: np.ndarray, start: float, end: float) -> tuple[np.ndarray, float]:
    total = len(x) / SAMPLE_RATE
    if not (0 <= start < end <= total + 0.05):
        raise HTTPException(422, f"span {start}-{end} s is outside the audio (0-{total:.2f} s)")
    if end - start > MAX_SPAN_SECONDS:
        raise HTTPException(422, f"span longer than {MAX_SPAN_SECONDS:.0f} s")
    return x[int(start * SAMPLE_RATE):int(min(end, total) * SAMPLE_RATE)], start


def _run(x: np.ndarray, spans: list[tuple[float, float, str]], deadline: float) -> list[list[dict[str, Any]]]:
    _load()
    result = []
    for start, end, text in spans:
        if time.monotonic() > deadline:
            raise HTTPException(504, "alignment took longer than the request limit")
        words = tokenize(text)
        piece, offset = _slice(x, start, end)
        times = align_samples(piece, words)
        result.append([{"word": w, "start": round(offset + a, 3), "end": round(offset + b, 3), "aligned": ok}
                       for w, (a, b, ok) in zip(words, times)])
    return result


async def _guarded(x: np.ndarray, spans: list[tuple[float, float, str]]) -> list[list[dict[str, Any]]]:
    try:
        await asyncio.wait_for(_gate.acquire(), QUEUE_WAIT)
    except asyncio.TimeoutError:
        raise HTTPException(503, "aligner busy, retry shortly", headers={"Retry-After": "10"})
    try:
        deadline = time.monotonic() + REQUEST_TIMEOUT
        return await asyncio.wait_for(asyncio.to_thread(_run, x, spans, deadline), REQUEST_TIMEOUT + 5)
    except asyncio.TimeoutError:
        raise HTTPException(504, "alignment took longer than the request limit")
    except HTTPException:
        raise
    except Exception as exc:                                    # model download/load failure and the like
        _state["error"] = type(exc).__name__
        raise HTTPException(503, f"aligner model unavailable ({type(exc).__name__})")
    finally:
        _gate.release()


async def _read_audio(audio: UploadFile) -> np.ndarray:
    data = await audio.read(MAX_AUDIO_BYTES + 1)
    if len(data) > MAX_AUDIO_BYTES:
        raise HTTPException(413, "audio file too large")
    return await asyncio.to_thread(decode, data)


@app.on_event("startup")
async def _preload() -> None:
    """Load (and on first run download) the model in the background so the first request is not slow."""
    def work() -> None:
        try:
            _load()
        except Exception as exc:                              # reported by /health, retried on the first request
            _state["error"] = type(exc).__name__
    threading.Thread(target=work, name="preload", daemon=True).start()


# ── routes ─────────────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": _state["model"] is not None, "device": DEVICE,
            "last_error": _state["error"]}


@app.post("/align")
async def align(audio: UploadFile = File(...), text: str = Form(...), language: str = Form("de"),
                start: float | None = Form(None), end: float | None = Form(None)):
    _check_language(language)
    _check_text(text)
    x = await _read_audio(audio)
    s = 0.0 if start is None else start
    e = len(x) / SAMPLE_RATE if end is None else end
    words = (await _guarded(x, [(s, e, text)]))[0]
    return {"language": language, "words": words}


@app.post("/align/batch")
async def align_batch(audio: UploadFile = File(...), spans: str = Form(...), language: str = Form("de")):
    _check_language(language)
    try:
        raw = json.loads(spans)
        parsed = [(float(r["start"]), float(r["end"]), str(r["text"])) for r in raw]
    except (ValueError, KeyError, TypeError):
        raise HTTPException(422, "spans must be a JSON list of {start, end, text}")
    if not parsed or len(parsed) > MAX_SPANS:
        raise HTTPException(422, f"between 1 and {MAX_SPANS} spans are allowed")
    for _, _, t in parsed:
        _check_text(t, strict=False)         # a sentence of digits only just gets its gap back
    if sum(len(t) for _, _, t in parsed) > MAX_TEXT_CHARS * 4:
        raise HTTPException(422, "too much text")
    x = await _read_audio(audio)
    return {"language": language, "spans": await _guarded(x, parsed)}