"""STT backends, transcription endpoints."""
from __future__ import annotations

import asyncio
import json
import re
import struct
import time
import uuid
from pathlib import Path

import requests
from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile

from core.config import _load_settings
from core.constants import (
    _WHISPER_DEFAULT, _FASTER_WHISPER_DEFAULT, _WHISPER_CPP_DEFAULT,
    _GROQ_STT_ENDPOINT, _NVIDIA_ASR_DEFAULT, _NVIDIA_ROUTER_DEFAULT,
    _MAX_UPLOAD_BYTES, _STT_REQUEST_TIMEOUT,
)
from core.registry import _registry_get, TEMP_DIR
from core.validation import _validate_http_url, _copy_limited
from core.audio import _to_wav_16k
from core.voice import _AUDIO_EXTS, _UPLOAD_EXTS

router = APIRouter()

# ── STT backend definitions ───────────────────────────────────────────────────

_STT_BACKEND_ALIASES = {
    "": "configured",
    "default": "configured",
    "whisper": "configured",
    "configured_whisper": "configured",
    "configured_stt": "configured",
    "parakeet": "nvidia_parakeet",
    "nvidia": "nvidia_parakeet",
    "nvidia_asr": "nvidia_parakeet",
    "nvidia_parakeet_asr": "nvidia_parakeet",
    "router": "nvidia_router",
    "speech_router": "nvidia_router",
    "nvidia_speech_router": "nvidia_router",
    "faster_whisper_server": "faster_whisper",
    "faster-whisper": "faster_whisper",
    "ctranslate2": "faster_whisper",
    "faster_w": "faster_whisper",
    "whisper-cpp": "whisper_cpp",
    "whisper_cpp_server": "whisper_cpp",
    "cpp": "whisper_cpp",
    "groq": "groq_whisper",
    "groq_stt": "groq_whisper",
    "groq-whisper": "groq_whisper",
}

_STT_BACKEND_METRICS: dict[str, dict] = {
    "configured":     {"speed": "GPU / CPU",        "latency": "1–5 s",       "quality": "large-v3",       "ram": "3 GB VRAM"},
    "faster_whisper": {"speed": "~70× RT · GPU",    "latency": "0.5–2 s",     "quality": "large-v3",       "ram": "1.5 GB VRAM"},
    "whisper_cpp":    {"speed": "~8–15× RT · CPU",  "latency": "1–5 s",       "quality": "large-v3 Q5",    "ram": "~1 GB RAM"},
    "groq_whisper":   {"speed": "fastest cloud",    "latency": "0.5–1 s",     "quality": "Whisper Turbo",  "ram": "cloud · 0"},
    "nvidia_parakeet":{"speed": "~200× RT · GPU",   "latency": "<0.3 s",      "quality": "Parakeet-TDT",   "ram": "~3 GB"},
    "nvidia_router":  {"speed": "GPU routed",       "latency": "~0.5 s",      "quality": "varies",         "ram": "~11 GB"},
}

_STT_VALID_BACKENDS = {"configured", "nvidia_parakeet", "nvidia_router", "faster_whisper", "whisper_cpp", "groq_whisper"}


def _clean_stt_backend(value: str) -> str:
    import re
    original = str(value or "").strip()
    if original.startswith("custom:"):
        return original
    key = re.sub(r"[^a-z0-9]+", "_", original.lower()).strip("_")
    key = _STT_BACKEND_ALIASES.get(key, key)
    return key if key in _STT_VALID_BACKENDS else "configured"


def _stt_backend_url(settings: dict, backend: str) -> str:
    backend = _clean_stt_backend(backend)
    if backend.startswith("custom:"):
        custom_id = backend[len("custom:"):]
        for card in settings.get("custom_engine_cards", []):
            if str(card.get("id", card.get("name", ""))) == custom_id:
                return card.get("url", "")
        return ""
    if backend == "nvidia_parakeet":
        return settings.get("nvidia_asr_url") or _NVIDIA_ASR_DEFAULT
    if backend == "nvidia_router":
        return settings.get("nvidia_router_url") or _NVIDIA_ROUTER_DEFAULT
    if backend == "faster_whisper":
        return settings.get("faster_whisper_url") or _FASTER_WHISPER_DEFAULT
    if backend == "whisper_cpp":
        return settings.get("whisper_cpp_url") or _WHISPER_CPP_DEFAULT
    if backend == "groq_whisper":
        return _GROQ_STT_ENDPOINT
    return settings.get("whisper_url") or _WHISPER_DEFAULT


def _stt_backend_model(backend: str) -> str:
    backend = _clean_stt_backend(backend)
    if backend in {"nvidia_parakeet", "nvidia_router"}:
        return "whisper-1"
    if backend == "whisper_cpp":
        return "whisper-1"
    if backend == "groq_whisper":
        return "whisper-large-v3-turbo"
    return "large-v3"


def _backend_port_label(url: str) -> str:
    from urllib.parse import urlsplit
    try:
        parts = urlsplit(url)
        if parts.port:
            return str(parts.port)
    except Exception:
        pass
    return ""


def _stt_backend_label(backend: str, url: str) -> str:
    labels = {
        "configured":     "Configured Whisper/STT",
        "nvidia_parakeet": "NVIDIA Parakeet ASR",
        "nvidia_router":  "NVIDIA Speech Router",
        "faster_whisper": "faster-whisper (CTranslate2 GPU)",
        "whisper_cpp":    "whisper.cpp (CPU/CUDA)",
        "groq_whisper":   "Groq Whisper (cloud · free)",
    }
    port = _backend_port_label(url)
    label = labels.get(backend, backend)
    return f"{port} {label}" if port else label


def _stt_backend_api_key(settings: dict, backend: str) -> str:
    if backend == "groq_whisper":
        return settings.get("groq_api_key", "").strip()
    return settings.get("whisper_api_key", "").strip()


def _make_minimal_wav(duration_ms: int = 500, sample_rate: int = 16000) -> bytes:
    """Minimal WAV: mono 16-bit silence of given duration at given sample rate."""
    num_frames = sample_rate * duration_ms // 1000
    data = b"\x00\x00" * num_frames
    header = struct.pack(
        "<4sI4s4sIHHIIHH4sI",
        b"RIFF", 36 + len(data), b"WAVE",
        b"fmt ", 16, 1, 1, sample_rate, sample_rate * 2, 2, 16,
        b"data", len(data),
    )
    return header + data


def _stt_backend_health(url: str) -> tuple[bool, list[str]]:
    base = _validate_http_url(url, allow_private=True).rstrip("/")
    models: list[str] = []
    ok = False
    try:
        r = requests.get(f"{base}/health", timeout=2)
        ok = r.status_code == 200
    except Exception:
        pass
    try:
        r = requests.get(f"{base}/v1/models", timeout=3)
        if r.status_code == 200:
            ok = True
            payload = r.json()
            data = payload.get("data", []) if isinstance(payload, dict) else []
            for item in data:
                if isinstance(item, dict) and item.get("id"):
                    models.append(str(item["id"]))
                elif isinstance(item, str):
                    models.append(item)
    except Exception:
        pass
    if not ok:
        return False, models
    try:
        wav = _make_minimal_wav()
        endpoint_found = False
        timed_out = False
        for path in ("/v1/audio/transcriptions", "/transcribe"):
            try:
                r = requests.post(
                    f"{base}{path}",
                    files={"file": ("probe.wav", wav, "audio/wav")},
                    data={"model": "whisper-1", "response_format": "text"},
                    timeout=8,
                )
                if r.status_code in {404, 405}:
                    continue
                if r.status_code == 500:
                    try:
                        detail = r.json().get("detail", "")
                        audio_quality_words = ("too short", "no speech", "audio", "empty",
                                               "duration", "length", "silence")
                        ok = bool(detail) and any(w in detail.lower() for w in audio_quality_words)
                    except Exception:
                        ok = False
                else:
                    ok = True
                endpoint_found = True
                break
            except requests.exceptions.ConnectionError:
                ok = False
                endpoint_found = True
                break
            except Exception:
                timed_out = True
        if not endpoint_found and not timed_out:
            ok = False
    except Exception:
        pass
    return ok, models


def _transcription_text_from_response(resp: requests.Response) -> str:
    try:
        payload = resp.json()
        if isinstance(payload, str):
            return payload.strip()
        if isinstance(payload, dict):
            for key in ("text", "transcript", "transcription"):
                if payload.get(key) is not None:
                    return str(payload[key]).strip()
            segments = payload.get("segments")
            if isinstance(segments, list):
                return " ".join(s.get("text", "").strip() for s in segments if s.get("text")).strip()
    except Exception:
        pass
    return resp.text.strip()


def _transcribe_audio(src: Path, settings: dict, backend: str = "configured") -> tuple[str, str]:
    backend = _clean_stt_backend(backend)
    stt_url = _validate_http_url(_stt_backend_url(settings, backend), allow_private=True).rstrip("/")
    stt_key = _stt_backend_api_key(settings, backend)
    hdrs = {"Authorization": f"Bearer {stt_key}"} if stt_key else {}
    model = _stt_backend_model(backend)
    paths = ["/v1/audio/transcriptions", "/transcribe"] if backend.startswith("custom:") else ["/v1/audio/transcriptions"]
    resp = None
    for path in paths:
        with src.open("rb") as f:
            resp = requests.post(
                f"{stt_url}{path}",
                files={"file": ("audio.wav", f, "audio/wav")},
                data={"model": model, "response_format": "text"},
                headers=hdrs,
                timeout=_STT_REQUEST_TIMEOUT,
            )
        if resp.status_code == 404 and len(paths) > 1:
            continue
        break
    if resp.status_code in {400, 404, 422, 500} and model != "whisper-1":
        path = paths[-1]
        with src.open("rb") as f:
            resp = requests.post(
                f"{stt_url}{path}",
                files={"file": ("audio.wav", f, "audio/wav")},
                data={"model": "whisper-1", "response_format": "text"},
                headers=hdrs,
                timeout=60,
            )
    if not resp.ok:
        try:
            body = resp.json()
            detail = body.get("detail") or body.get("error") or body.get("message") or str(body)
        except Exception:
            detail = resp.text[:300].strip()
        if not detail or detail.lower() in {"internal server error", "unknown error"}:
            detail = f"HTTP {resp.status_code} — backend may be misconfigured or missing CUDA support"
        if "'NoneType'" in detail and "'to'" in detail:
            detail = ("Speaker diarization failed — pyannote/speaker-diarization-3.1 requires "
                      "a HuggingFace token. Get one at hf.co/settings/tokens and accept the "
                      "model license at hf.co/pyannote/speaker-diarization-3.1, then add the "
                      "token to the whisperx-gpu container env as HF_TOKEN.")
        raise RuntimeError(f"STT ({stt_url}): {detail}")
    return _transcription_text_from_response(resp), backend


def _word_tokens(text: str) -> list[str]:
    return re.findall(r"[\w']+", (text or "").lower(), flags=re.UNICODE)


def _word_accuracy(reference: str, hypothesis: str) -> float | None:
    ref = _word_tokens(reference)
    hyp = _word_tokens(hypothesis)
    if not ref:
        return None
    prev = list(range(len(hyp) + 1))
    for i, rw in enumerate(ref, 1):
        cur = [i]
        for j, hw in enumerate(hyp, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (0 if rw == hw else 1)))
        prev = cur
    return max(0.0, (1.0 - prev[-1] / max(1, len(ref))) * 100.0)


def _stt_device_hint(engine: dict, settings: dict) -> str:
    text = " ".join(str(engine.get(k, "")) for k in ("id", "backend", "label", "url", "model", "source")).lower()
    if any(token in text for token in ("nvidia", "parakeet", "nemotron", "whisperx", "faster", "cuda", "gpu", "blackwell")):
        return "CUDA/GPU"
    if "cpu" in text or "whisper_cpp" in text or "whisper.cpp" in text:
        return "CPU"
    backend = _clean_stt_backend(str(engine.get("backend") or engine.get("id") or ""))
    metrics = _STT_BACKEND_METRICS.get(backend, {})
    speed = str(metrics.get("speed", "")).lower()
    if "gpu" in speed or "cuda" in speed:
        return "CUDA/GPU"
    if "cpu" in speed:
        return "CPU"
    return "Unknown"


def _transcribe_url_for_benchmark(src: Path, url: str, model: str, api_key: str = "") -> str:
    base = _validate_http_url(url, allow_private=True).rstrip("/")
    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
    last_resp: requests.Response | None = None
    last_error = ""
    for path in ("/v1/audio/transcriptions", "/transcribe", "/audio/transcriptions"):
        try:
            with src.open("rb") as f:
                resp = requests.post(
                    f"{base}{path}",
                    files={"file": ("benchmark.wav", f, "audio/wav")},
                    data={"model": model or "whisper-1", "response_format": "text"},
                    headers=headers,
                    timeout=_STT_REQUEST_TIMEOUT,
                )
            last_resp = resp
            if resp.status_code in {404, 405}:
                continue
            if not resp.ok and (model or "whisper-1") != "whisper-1":
                with src.open("rb") as f:
                    retry = requests.post(
                        f"{base}{path}",
                        files={"file": ("benchmark.wav", f, "audio/wav")},
                        data={"model": "whisper-1", "response_format": "text"},
                        headers=headers,
                        timeout=_STT_REQUEST_TIMEOUT,
                    )
                if retry.ok:
                    return _transcription_text_from_response(retry)
                last_resp = retry
            if resp.ok:
                return _transcription_text_from_response(resp)
            try:
                body = resp.json()
                last_error = body.get("detail") or body.get("error") or body.get("message") or str(body)
            except Exception:
                last_error = resp.text[:400].strip()
        except Exception as e:
            last_error = str(e)
    status = f"HTTP {last_resp.status_code}" if last_resp is not None else "STT request failed"
    raise RuntimeError(f"{status}: {last_error or 'No compatible transcription endpoint'}")


def _benchmark_stt_engine(src: Path, engine: dict, settings: dict, reference: str) -> dict:
    label = str(engine.get("label") or engine.get("id") or engine.get("backend") or "STT")
    backend = str(engine.get("backend") or engine.get("id") or "").strip()
    url = str(engine.get("url") or "").strip()
    model = str(engine.get("model") or "").strip() or (backend and _stt_backend_model(backend)) or "whisper-1"
    device = _stt_device_hint(engine, settings)
    t0 = time.perf_counter()
    try:
        if url:
            api_key = _stt_backend_api_key(settings, _clean_stt_backend(backend)) if backend else settings.get("whisper_api_key", "")
            text = _transcribe_url_for_benchmark(src, url, model, api_key)
            used_url = _validate_http_url(url, allow_private=True).rstrip("/")
        else:
            text, used_backend = _transcribe_audio(src, settings, backend or "configured")
            used_url = _stt_backend_url(settings, used_backend)
            model = model or _stt_backend_model(used_backend)
        elapsed = time.perf_counter() - t0
        accuracy = _word_accuracy(reference, text)
        return {
            "ok": True,
            "engine": label,
            "backend": backend,
            "url": used_url,
            "model": model,
            "device": device,
            "time_sec": elapsed,
            "accuracy": accuracy,
            "output": text,
        }
    except Exception as e:
        return {
            "ok": False,
            "engine": label,
            "backend": backend,
            "url": url,
            "model": model,
            "device": device,
            "time_sec": time.perf_counter() - t0,
            "accuracy": None,
            "output": "",
            "error": str(e),
        }


async def _default_stt_benchmark_engines() -> list[dict]:
    payload = await stt_backends()
    engines = []
    for item in payload.get("backends", []):
        if item.get("available"):
            engines.append({
                "id": item.get("id"),
                "backend": item.get("id"),
                "label": item.get("label"),
                "url": item.get("url"),
                "model": item.get("model") or (item.get("models") or ["whisper-1"])[0],
            })
    return engines


# ── STT routes ────────────────────────────────────────────────────────────────

@router.get("/api/stt-backends")
async def stt_backends():
    settings = _load_settings()
    items = []
    seen_urls: set[tuple[str, str]] = set()
    ordered = ("configured", "faster_whisper", "whisper_cpp", "groq_whisper", "nvidia_parakeet", "nvidia_router")
    for backend in ordered:
        raw_url = _stt_backend_url(settings, backend)
        url = _validate_http_url(raw_url, allow_private=True).rstrip("/")
        key = (backend, url)
        if key in seen_urls:
            continue
        seen_urls.add(key)
        api_key = _stt_backend_api_key(settings, backend)
        if backend == "groq_whisper":
            ok = bool(api_key)
            models: list[str] = ["whisper-large-v3-turbo", "whisper-large-v3", "distil-whisper-large-v3-en"]
        else:
            ok, models = _stt_backend_health(url)
        items.append({
            "id": backend,
            "label": _stt_backend_label(backend, url),
            "url": url,
            "port": _backend_port_label(url),
            "available": ok,
            "model": _stt_backend_model(backend),
            "models": models,
            "metrics": _STT_BACKEND_METRICS.get(backend, {}),
        })
    for card in settings.get("custom_engine_cards", []):
        if card.get("role") not in ("stt", "stt+tts"):
            continue
        raw_url = card.get("url", "").strip()
        url = _validate_http_url(raw_url, allow_private=True).rstrip("/")
        if not url:
            continue
        card_id = "custom:" + str(card.get("id", card.get("name", "")))
        if (card_id, url) in seen_urls:
            continue
        seen_urls.add((card_id, url))
        ok, models = _stt_backend_health(url)
        port = _backend_port_label(url)
        label = card.get("label") or card.get("name") or "Custom STT"
        if port:
            label = f"{port} {label}"
        items.append({
            "id": card_id,
            "label": label,
            "url": url,
            "port": port,
            "available": ok,
            "model": "whisper-1",
            "models": models,
            "metrics": {},
        })
    return {"backends": items}


@router.post("/api/transcribe")
async def transcribe(request: Request):
    data = await request.json()
    fid: str = data["id"]
    src = _registry_get(fid)
    if src is None or not src.exists():
        raise HTTPException(404, "Audio not found")
    settings = _load_settings()
    backend = _clean_stt_backend(str(data.get("backend") or data.get("stt_backend") or "configured"))
    try:
        text, used_backend = await asyncio.to_thread(_transcribe_audio, src, settings, backend)
        return {"text": text, "backend": used_backend}
    except Exception as e:
        raise HTTPException(502, f"STT error ({backend}): {e}")


@router.post("/api/transcribe-bytes")
async def transcribe_bytes(
    file: UploadFile = File(...),
    backend: str = Form("configured"),
):
    """Accept raw audio upload and return transcription directly (used by hotkey daemon)."""
    suffix = Path(file.filename or "audio.wav").suffix.lower() or ".wav"
    if suffix not in _AUDIO_EXTS:
        raise HTTPException(400, "Unsupported audio type")
    tmp = TEMP_DIR / f"{uuid.uuid4().hex}_daemon{suffix}"
    wav_tmp = tmp
    try:
        with tmp.open("wb") as f:
            _copy_limited(file.file, f, _MAX_UPLOAD_BYTES)
        if suffix != ".wav":
            wav_tmp = _to_wav_16k(tmp)
        settings = _load_settings()
        stt_backend = _clean_stt_backend(backend)
        text, used_backend = await asyncio.to_thread(_transcribe_audio, wav_tmp, settings, stt_backend)
        return {"text": text, "backend": used_backend}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(502, f"STT error: {e}")
    finally:
        for p in {tmp, wav_tmp}:
            try:
                p.unlink(missing_ok=True)
            except Exception:
                pass

@router.post("/api/stt-benchmark")
async def stt_benchmark(
    audio: UploadFile = File(...),
    reference_text: str = Form(""),
    reference_file: UploadFile | None = File(None),
    engines_json: str = Form(""),
):
    suffix = Path(audio.filename or "audio.wav").suffix.lower() or ".wav"
    if suffix not in (_AUDIO_EXTS | _UPLOAD_EXTS):
        raise HTTPException(400, "Unsupported audio type")
    tmp = TEMP_DIR / f"{uuid.uuid4().hex}_stt_bench{suffix}"
    wav_tmp = tmp
    try:
        with tmp.open("wb") as f:
            _copy_limited(audio.file, f, _MAX_UPLOAD_BYTES)
        if suffix != ".wav":
            wav_tmp = _to_wav_16k(tmp)
        reference = (reference_text or "").strip()
        if reference_file is not None:
            raw = await reference_file.read()
            if raw:
                reference = raw[:2_000_000].decode("utf-8", errors="replace").strip()
        if not reference:
            raise HTTPException(400, "Reference text is required")
        try:
            parsed = json.loads(engines_json) if engines_json else []
            engines = parsed if isinstance(parsed, list) else []
        except Exception:
            engines = []
        if not engines:
            engines = await _default_stt_benchmark_engines()
        if not engines:
            raise HTTPException(400, "No STT engines selected")
        settings = _load_settings()
        results = []
        for engine in engines:
            if not isinstance(engine, dict):
                continue
            results.append(await asyncio.to_thread(_benchmark_stt_engine, wav_tmp, engine, settings, reference))
        ok_rows = [r for r in results if r.get("ok")]
        best_time = min((r.get("time_sec") for r in ok_rows if isinstance(r.get("time_sec"), (int, float))), default=None)
        best_accuracy = max((r.get("accuracy") for r in ok_rows if isinstance(r.get("accuracy"), (int, float))), default=None)
        return {
            "ok": True,
            "reference_words": len(_word_tokens(reference)),
            "results": results,
            "best_time_sec": best_time,
            "best_accuracy": best_accuracy,
        }
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, f"STT benchmark failed: {e}")
    finally:
        for p in {tmp, wav_tmp}:
            try:
                p.unlink(missing_ok=True)
            except Exception:
                pass