"""Speech-to-text sidecar — Pipecat local Whisper only (no GapGPT / cloud STT)."""

from __future__ import annotations

import io
import os
import shutil
import subprocess
import tempfile
import wave
from pathlib import Path

import numpy as np
from dotenv import load_dotenv
from fastapi import FastAPI, File, Header, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from pipecat.services.whisper.stt import WhisperSTTService
from pipecat.transcriptions.language import Language

load_dotenv(Path(__file__).resolve().parent / ".env")

MAX_BYTES = 8 * 1024 * 1024
WAV_MAGIC = (b"RIFF", b"WAVE")

app = FastAPI(title="Mento STT", docs_url=None, redoc_url=None)
_stt: WhisperSTTService | None = None


def _env(name: str, default: str = "") -> str:
    return (os.getenv(name) or default).strip()


def _require_token(header_token: str | None) -> None:
    expected = _env("STT_INTERNAL_TOKEN")
    if not expected:
        return
    if (header_token or "").strip() != expected:
        raise HTTPException(status_code=401, detail="توکن سرویس تبدیل صوت نامعتبر است")


def _model_name() -> str:
    """Map env / legacy cloud names to Faster-Whisper local sizes."""
    raw = (_env("STT_MODEL", "small") or "small").lower()
    # Old GapGPT/OpenAI cloud id — not a local Whisper size
    legacy = {
        "whisper-1": "small",
        "whisper": "small",
        "gpt-4o-transcribe": "small",
        "gpt-4o-mini-transcribe": "base",
    }
    if raw in legacy:
        return legacy[raw]
    allowed = {
        "tiny",
        "tiny.en",
        "base",
        "base.en",
        "small",
        "small.en",
        "medium",
        "medium.en",
        "large",
        "large-v1",
        "large-v2",
        "large-v3",
        "distil-large-v2",
        "distil-medium.en",
        "distil-large-v3",
        "large-v3-turbo",
    }
    if raw in allowed:
        return raw
    # unknown → safe default for Persian
    return "small"


def _stt_service() -> WhisperSTTService:
    """Local Faster-Whisper via Pipecat — no OPENAI_API_KEY / GapGPT."""
    global _stt
    if _stt is not None:
        return _stt

    device = _env("STT_DEVICE", "cpu") or "cpu"
    compute = _env("STT_COMPUTE_TYPE", "int8") or "int8"

    try:
        # Settings fields differ by Pipecat version — keep to model + language only
        settings_kwargs: dict = {"model": _model_name()}
        try:
            settings_kwargs["language"] = Language.FA
            settings = WhisperSTTService.Settings(**settings_kwargs)
        except TypeError:
            settings_kwargs.pop("language", None)
            settings = WhisperSTTService.Settings(**settings_kwargs)

        _stt = WhisperSTTService(
            device=device,
            compute_type=compute,
            settings=settings,
        )
    except TypeError:
        # Older constructor: model=/language= kwargs instead of settings=
        try:
            _stt = WhisperSTTService(
                model=_model_name(),
                device=device,
                compute_type=compute,
                language=Language.FA,
            )
        except Exception as exc:
            raise HTTPException(
                status_code=503,
                detail=f"بارگذاری مدل Whisper محلی ناموفق بود: {exc}",
            ) from exc
    except Exception as exc:
        raise HTTPException(
            status_code=503,
            detail=f"بارگذاری مدل Whisper محلی ناموفق بود: {exc}",
        ) from exc

    # force model load at startup path (first request otherwise downloads silently)
    if getattr(_stt, "_model", None) is None and hasattr(_stt, "_load"):
        try:
            _stt._load()
        except Exception as exc:
            raise HTTPException(
                status_code=503,
                detail=(
                    "مدل Whisper دانلود/بارگذاری نشد. "
                    f"اینترنت سرور و فضای دیسک را چک کنید. جزئیات: {exc}"
                ),
            ) from exc

    return _stt


def _is_wav(data: bytes) -> bool:
    return len(data) >= 12 and data[:4] == WAV_MAGIC[0] and data[8:12] == WAV_MAGIC[1]


def _to_wav(data: bytes, filename: str) -> bytes:
    if _is_wav(data):
        return data

    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise HTTPException(
            status_code=400,
            detail="فرمت صوت پشتیبانی نشد. فایل WAV ارسال کنید یا ffmpeg را نصب کنید.",
        )

    suffix = Path(filename or "audio.webm").suffix or ".webm"
    with tempfile.TemporaryDirectory() as tmp:
        src = Path(tmp) / f"in{suffix}"
        dst = Path(tmp) / "out.wav"
        src.write_bytes(data)
        result = subprocess.run(
            [
                ffmpeg,
                "-y",
                "-i",
                str(src),
                "-ac",
                "1",
                "-ar",
                "16000",
                "-f",
                "wav",
                str(dst),
            ],
            capture_output=True,
            check=False,
        )
        if result.returncode != 0 or not dst.exists():
            raise HTTPException(status_code=400, detail="تبدیل فایل صوتی ناموفق بود")
        return dst.read_bytes()


def _wav_to_float32(wav_bytes: bytes) -> tuple[np.ndarray, int]:
    with wave.open(io.BytesIO(wav_bytes), "rb") as wf:
        channels = wf.getnchannels()
        sample_width = wf.getsampwidth()
        rate = wf.getframerate()
        frames = wf.readframes(wf.getnframes())

    if sample_width == 2:
        audio = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0
    elif sample_width == 4:
        audio = np.frombuffer(frames, dtype=np.int32).astype(np.float32) / 2147483648.0
    else:
        raise HTTPException(status_code=400, detail="فرمت نمونهٔ صوتی پشتیبانی نمی‌شود")

    if channels > 1:
        audio = audio.reshape(-1, channels).mean(axis=1)

    return audio, rate


def _transcribe_local(wav_bytes: bytes) -> str:
    service = _stt_service()
    model = getattr(service, "_model", None)
    if model is None:
        raise HTTPException(status_code=503, detail="مدل Whisper محلی آماده نیست")

    audio, rate = _wav_to_float32(wav_bytes)
    # faster-whisper expects 16k mono float32 ideally
    if rate != 16000:
        # simple resample via ffmpeg path already targets 16k for non-wav;
        # for wav at other rates, re-encode quickly
        ffmpeg = shutil.which("ffmpeg")
        if not ffmpeg:
            raise HTTPException(status_code=400, detail="نرخ نمونه‌برداری صوت پشتیبانی نشد")
        with tempfile.TemporaryDirectory() as tmp:
            src = Path(tmp) / "in.wav"
            dst = Path(tmp) / "out.wav"
            src.write_bytes(wav_bytes)
            result = subprocess.run(
                [ffmpeg, "-y", "-i", str(src), "-ac", "1", "-ar", "16000", "-f", "wav", str(dst)],
                capture_output=True,
                check=False,
            )
            if result.returncode != 0 or not dst.exists():
                raise HTTPException(status_code=400, detail="نرمال‌سازی صوت ناموفق بود")
            audio, rate = _wav_to_float32(dst.read_bytes())

    settings = getattr(service, "_settings", None)
    language = getattr(settings, "language", None) or Language.FA
    lang_code = getattr(language, "value", language) if language is not None else "fa"
    if isinstance(lang_code, str) and "-" in lang_code:
        lang_code = lang_code.split("-")[0]
    # Language enum may be "fa" or "FA" or full name — faster-whisper wants ISO code
    if isinstance(lang_code, str):
        lang_code = lang_code.lower()
        if lang_code in ("persian", "farsi"):
            lang_code = "fa"

    initial_prompt = _env(
        "STT_PROMPT",
        "یادداشت و مکالمهٔ فروش به زبان فارسی. اسامی مشتریان و اصطلاحات کسب‌وکار را دقیق بنویس.",
    ) or None

    segments, _info = model.transcribe(
        audio,
        language=lang_code or "fa",
        initial_prompt=initial_prompt,
        vad_filter=True,
    )
    parts = [seg.text.strip() for seg in segments if getattr(seg, "text", None)]
    return " ".join(p for p in parts if p).strip()


def _friendly_stt_error(exc: BaseException) -> tuple[int, str]:
    text = str(exc)
    lowered = text.lower()
    if "timeout" in lowered or "timed out" in lowered:
        return 504, "زمان تبدیل صوت به پایان رسید. دوباره با صدای کوتاه‌تر تلاش کنید."
    if "memory" in lowered or "oom" in lowered:
        return 503, "حافظهٔ سرور برای مدل Whisper کافی نیست. مدل کوچک‌تر (tiny/base) را امتحان کنید."
    if len(text) > 180:
        return 502, "تبدیل صوت به متن محلی ناموفق بود. لاگ mento-stt را بررسی کنید."
    return 502, f"تبدیل صوت به متن ناموفق بود: {text}"


@app.get("/health")
async def health() -> dict[str, object]:
    return {
        "ok": True,
        "engine": "pipecat",
        "service": "WhisperSTTService",
        "mode": "local",
        "model": _env("STT_MODEL", "small"),
        "device": _env("STT_DEVICE", "cpu") or "cpu",
        "cloud": False,
        "gapgpt": False,
    }


@app.post("/transcribe")
async def transcribe(
    audio: UploadFile = File(...),
    x_stt_token: str | None = Header(default=None, alias="X-STT-Token"),
) -> JSONResponse:
    _require_token(x_stt_token)

    data = await audio.read()
    if not data:
        raise HTTPException(status_code=400, detail="فایل صوتی خالی است")
    if len(data) > MAX_BYTES:
        raise HTTPException(status_code=400, detail="حجم فایل صوتی بیش از حد مجاز است")

    wav = _to_wav(data, audio.filename or "audio.wav")

    try:
        text = _transcribe_local(wav)
    except HTTPException:
        raise
    except Exception as exc:
        status, detail = _friendly_stt_error(exc)
        raise HTTPException(status_code=status, detail=detail) from exc

    return JSONResponse({"text": text})