# -*- coding: utf-8 -*-
"""로컬 상주 faster-whisper 스트리밍 STT 사이드카.

Higgs TTS 서버와 같은 형태다. 모델을 한 번만 올려 두고 loopback WebSocket으로
linear16 PCM을 받아 interim/final 전사를 돌려준다. 오디오는 디스크에 쓰지 않고
발화 단위 메모리 버퍼만 사용하며, 발화가 끝나면 즉시 버린다.

프로토콜(클라이언트 → 서버):
  binary frame  raw linear16 PCM
  {"type": "Finalize"}     현재 발화를 즉시 확정한다
  {"type": "CloseStream"}  남은 발화를 확정하고 닫는다

프로토콜(서버 → 클라이언트):
  {"type": "ready", ...}
  {"type": "transcript", "text": ..., "is_final": bool, "speech_final": bool,
   "confidence": float|null, "words": [...], "duration": float}
  {"type": "error", "detail": "..."}

`--enable` 없이는 뜨지 않고, 기본 바인딩은 loopback뿐이다.
"""

from __future__ import annotations

import argparse
import array
import asyncio
import json
import math
import sys
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Iterable, Protocol
from urllib.parse import parse_qs, urlsplit


DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 9882
DEFAULT_MODEL = "large-v3"
DEFAULT_LANGUAGE = "ko"
DEFAULT_SAMPLE_RATE = 16_000
DEFAULT_CHANNELS = 1
SAMPLE_WIDTH = 2

# 발화 경계 판정
DEFAULT_ENDPOINTING_MS = 300
DEFAULT_UTTERANCE_END_MS = 1_200
DEFAULT_INTERIM_INTERVAL_MS = 700
DEFAULT_SILENCE_RMS = 320
MIN_UTTERANCE_MS = 200
MAX_UTTERANCE_SECONDS = 60.0
MAX_FRAME_BYTES = 1 << 20

ALLOWED_MODELS = (
    "large-v3",
    "large-v3-turbo",
    "medium",
    "small",
    "base",
)


class TranscriptionError(RuntimeError):
    """전사 실패. 원문 오디오는 로그에 남기지 않는다."""


@dataclass(frozen=True, slots=True)
class Word:
    word: str
    start: float
    end: float
    confidence: float | None = None

    def as_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "word": self.word,
            "start": round(self.start, 3),
            "end": round(self.end, 3),
        }
        if self.confidence is not None:
            payload["confidence"] = round(self.confidence, 4)
        return payload


@dataclass(frozen=True, slots=True)
class Transcription:
    text: str
    confidence: float | None = None
    words: tuple[Word, ...] = ()


class Transcriber(Protocol):
    def transcribe(
        self, pcm: bytes, *, sample_rate: int, final: bool
    ) -> Transcription: ...


def frame_rms(pcm: bytes) -> int:
    """linear16 프레임의 RMS. 빈 입력은 무음으로 본다.

    `audioop` 은 Python 3.13에서 제거되므로 표준 정수 연산으로 직접 구한다.
    """

    if not pcm or len(pcm) < SAMPLE_WIDTH:
        return 0
    usable = len(pcm) - (len(pcm) % SAMPLE_WIDTH)
    samples = array.array("h")
    samples.frombytes(pcm[:usable])
    if sys.byteorder != "little":  # pragma: no cover - little-endian 개발 환경
        samples.byteswap()
    if not samples:
        return 0
    total = 0
    for sample in samples:
        total += sample * sample
    return int(math.sqrt(total / len(samples)))


@dataclass
class UtteranceBuffer:
    """한 발화 분량의 PCM과 무음 누적을 추적한다.

    이 클래스는 모델을 전혀 모른다. 덕분에 GPU 없이도 경계 판정을 그대로
    테스트할 수 있다.
    """

    sample_rate: int = DEFAULT_SAMPLE_RATE
    channels: int = DEFAULT_CHANNELS
    endpointing_ms: int = DEFAULT_ENDPOINTING_MS
    utterance_end_ms: int = DEFAULT_UTTERANCE_END_MS
    interim_interval_ms: int = DEFAULT_INTERIM_INTERVAL_MS
    silence_rms: int = DEFAULT_SILENCE_RMS
    _pcm: bytearray = field(default_factory=bytearray, init=False)
    _silence_ms: float = field(default=0.0, init=False)
    _speech_ms: float = field(default=0.0, init=False)
    _since_interim_ms: float = field(default=0.0, init=False)
    _saw_speech: bool = field(default=False, init=False)

    @property
    def bytes_per_ms(self) -> float:
        return self.sample_rate * self.channels * SAMPLE_WIDTH / 1000.0

    def duration_ms(self, pcm_length: int | None = None) -> float:
        length = len(self._pcm) if pcm_length is None else pcm_length
        return length / self.bytes_per_ms if self.bytes_per_ms else 0.0

    @property
    def pcm(self) -> bytes:
        return bytes(self._pcm)

    @property
    def saw_speech(self) -> bool:
        return self._saw_speech

    def append(self, chunk: bytes) -> None:
        if not chunk:
            return
        self._pcm.extend(chunk)
        span_ms = self.duration_ms(len(chunk))
        self._since_interim_ms += span_ms
        if frame_rms(chunk) >= self.silence_rms:
            self._saw_speech = True
            self._speech_ms += span_ms
            self._silence_ms = 0.0
        else:
            self._silence_ms += span_ms

    def should_emit_interim(self) -> bool:
        if not self._saw_speech:
            return False
        if self.duration_ms() < MIN_UTTERANCE_MS:
            return False
        return self._since_interim_ms >= self.interim_interval_ms

    def mark_interim_emitted(self) -> None:
        self._since_interim_ms = 0.0

    def should_finalize(self) -> bool:
        """무음이 endpointing 기준을 넘거나 발화가 최대 길이에 닿으면 확정한다."""

        if not self._saw_speech:
            return False
        if self.duration_ms() >= MAX_UTTERANCE_SECONDS * 1000:
            return True
        return self._silence_ms >= self.endpointing_ms

    def is_stalled(self) -> bool:
        """말이 없는 채로 utterance_end_ms 를 넘긴 상태."""

        return not self._saw_speech and self._silence_ms >= self.utterance_end_ms

    def reset(self) -> None:
        self._pcm.clear()
        self._silence_ms = 0.0
        self._speech_ms = 0.0
        self._since_interim_ms = 0.0
        self._saw_speech = False


def transcript_payload(
    result: Transcription,
    *,
    is_final: bool,
    speech_final: bool,
    duration_seconds: float,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "type": "transcript",
        "text": result.text,
        "is_final": bool(is_final),
        "speech_final": bool(speech_final),
        "duration": round(max(0.0, duration_seconds), 3),
        "words": [word.as_payload() for word in result.words] if is_final else [],
    }
    payload["confidence"] = (
        round(result.confidence, 4) if result.confidence is not None else None
    )
    return payload


def parse_stream_options(path: str) -> dict[str, Any]:
    """접속 쿼리에서 세션 옵션을 읽는다. 알 수 없는 값은 기본값으로 되돌린다."""

    query = parse_qs(urlsplit(path).query)

    def one(name: str) -> str:
        values = query.get(name) or []
        return values[0].strip() if values else ""

    def integer(name: str, default: int, low: int, high: int) -> int:
        raw = one(name)
        if not raw:
            return default
        try:
            return max(low, min(high, int(raw)))
        except ValueError:
            return default

    model = one("model") or DEFAULT_MODEL
    if model not in ALLOWED_MODELS:
        raise ValueError("model_not_allowlisted")
    language = one("language") or DEFAULT_LANGUAGE
    if not language.isalpha() or len(language) > 8:
        raise ValueError("language_invalid")
    return {
        "model": model,
        "language": language,
        "sample_rate": integer("sample_rate", DEFAULT_SAMPLE_RATE, 8_000, 48_000),
        "channels": integer("channels", DEFAULT_CHANNELS, 1, 2),
        "endpointing_ms": integer("endpointing", DEFAULT_ENDPOINTING_MS, 10, 5_000),
        "utterance_end_ms": integer(
            "utterance_end_ms", DEFAULT_UTTERANCE_END_MS, 1_000, 10_000
        ),
        "interim_interval_ms": integer(
            "interim_interval_ms", DEFAULT_INTERIM_INTERVAL_MS, 200, 5_000
        ),
        "silence_rms": integer("silence_rms", DEFAULT_SILENCE_RMS, 1, 10_000),
    }


class StreamSession:
    """한 학습자 발화 스트림. 모델 호출은 주입된 transcriber 로만 한다."""

    def __init__(
        self,
        *,
        options: dict[str, Any],
        transcriber: Transcriber,
        send: Callable[[dict[str, Any]], Any],
        run_blocking: Callable[..., Any] | None = None,
    ) -> None:
        self.options = options
        self._transcriber = transcriber
        self._send = send
        self._run_blocking = run_blocking
        self.buffer = UtteranceBuffer(
            sample_rate=options["sample_rate"],
            channels=options["channels"],
            endpointing_ms=options["endpointing_ms"],
            utterance_end_ms=options["utterance_end_ms"],
            interim_interval_ms=options["interim_interval_ms"],
            silence_rms=options["silence_rms"],
        )
        self.final_segments: list[str] = []

    async def _transcribe(self, pcm: bytes, *, final: bool) -> Transcription:
        def call() -> Transcription:
            return self._transcriber.transcribe(
                pcm, sample_rate=self.options["sample_rate"], final=final
            )

        if self._run_blocking is None:
            return call()
        return await self._run_blocking(call)

    async def feed(self, chunk: bytes) -> None:
        if len(chunk) > MAX_FRAME_BYTES:
            raise TranscriptionError("audio_frame_too_large")
        self.buffer.append(chunk)
        if self.buffer.should_finalize():
            await self.finalize(speech_final=True)
            return
        if self.buffer.should_emit_interim():
            self.buffer.mark_interim_emitted()
            pcm = self.buffer.pcm
            result = await self._transcribe(pcm, final=False)
            if result.text.strip():
                await self._send(
                    transcript_payload(
                        result,
                        is_final=False,
                        speech_final=False,
                        duration_seconds=self.buffer.duration_ms(len(pcm)) / 1000.0,
                    )
                )

    async def finalize(self, *, speech_final: bool) -> None:
        pcm = self.buffer.pcm
        duration_seconds = self.buffer.duration_ms(len(pcm)) / 1000.0
        saw_speech = self.buffer.saw_speech
        self.buffer.reset()
        if not saw_speech or self.buffer.duration_ms(len(pcm)) < MIN_UTTERANCE_MS:
            return
        result = await self._transcribe(pcm, final=True)
        text = result.text.strip()
        if not text:
            return
        self.final_segments.append(text)
        await self._send(
            transcript_payload(
                result,
                is_final=True,
                speech_final=speech_final,
                duration_seconds=duration_seconds,
            )
        )


def warmup_pcm(sample_rate: int, *, milliseconds: int = 500) -> bytes:
    """warmup 용 무음 linear16 버퍼."""

    samples = max(1, int(sample_rate * milliseconds / 1000))
    return b"\x00\x00" * samples


class FasterWhisperTranscriber:
    """상주 faster-whisper 모델 하나. 오디오는 메모리에서만 다룬다."""

    def __init__(self, model: str, *, device: str, compute_type: str) -> None:
        try:
            import numpy as np
            from faster_whisper import WhisperModel
        except ImportError as exc:  # pragma: no cover - 런타임 환경 의존
            raise TranscriptionError("faster_whisper_unavailable") from exc
        self._np = np
        self.model_name = model
        self.device = device
        self.compute_type = compute_type
        self._model = WhisperModel(model, device=device, compute_type=compute_type)
        self._language = DEFAULT_LANGUAGE

    def configure_language(self, language: str) -> None:
        self._language = language

    def warmup(self) -> None:
        """짧은 무음 한 번을 실제로 돌려 디바이스가 살아 있는지 확인한다.

        `get_cuda_device_count() > 0` 만으로는 부족하다. cuDNN 이 없으면 모델은
        올라가지만 첫 추론에서 죽어 스트림이 조용히 끊긴다.
        """

        self.transcribe(
            warmup_pcm(DEFAULT_SAMPLE_RATE),
            sample_rate=DEFAULT_SAMPLE_RATE,
            final=False,
        )

    def transcribe(
        self, pcm: bytes, *, sample_rate: int, final: bool
    ) -> Transcription:
        if not pcm:
            return Transcription(text="")
        audio = (
            self._np.frombuffer(pcm, dtype=self._np.int16).astype(self._np.float32)
            / 32768.0
        )
        segments, _info = self._model.transcribe(
            audio,
            language=self._language,
            beam_size=5 if final else 1,
            word_timestamps=final,
            vad_filter=False,
            condition_on_previous_text=False,
        )
        texts: list[str] = []
        words: list[Word] = []
        confidences: list[float] = []
        for segment in segments:
            piece = (segment.text or "").strip()
            if piece:
                texts.append(piece)
            probability = getattr(segment, "avg_logprob", None)
            if probability is not None:
                confidences.append(min(1.0, max(0.0, float(2 ** probability))))
            for word in getattr(segment, "words", None) or ():
                token = (getattr(word, "word", "") or "").strip()
                start = getattr(word, "start", None)
                end = getattr(word, "end", None)
                if not token or start is None or end is None or end <= start:
                    continue
                probability = getattr(word, "probability", None)
                words.append(
                    Word(
                        word=token,
                        start=float(start),
                        end=float(end),
                        confidence=float(probability)
                        if probability is not None
                        else None,
                    )
                )
        confidence = sum(confidences) / len(confidences) if confidences else None
        return Transcription(
            text=" ".join(texts).strip(),
            confidence=confidence,
            words=tuple(words),
        )


def resolve_device(requested: str, *, cuda_available: bool | None = None) -> tuple[str, str]:
    """요청한 디바이스를 실제 가용 자원으로 해석한다."""

    if requested not in {"auto", "cuda", "cpu"}:
        raise ValueError("device_invalid")
    if cuda_available is None:
        try:  # pragma: no cover - 런타임 환경 의존
            import ctranslate2

            cuda_available = ctranslate2.get_cuda_device_count() > 0
        except Exception:
            cuda_available = False
    if requested == "cuda":
        if not cuda_available:
            raise ValueError("cuda_unavailable")
        return "cuda", "float16"
    if requested == "cpu":
        return "cpu", "int8"
    return ("cuda", "float16") if cuda_available else ("cpu", "int8")


def build_transcriber(
    model: str,
    requested_device: str,
    *,
    factory: Callable[[str, str, str], Any],
    cuda_available: bool | None = None,
) -> Any:
    """디바이스를 해석하고 warmup 까지 통과한 transcriber 만 돌려준다.

    `auto` 는 CUDA 가 보이면 먼저 시도하되, 첫 추론이 실패하면 CPU 로 내려간다.
    명시적 `--device cuda` 는 조용히 강등하지 않고 그대로 실패시킨다.
    """

    device, compute_type = resolve_device(requested_device, cuda_available=cuda_available)
    try:
        transcriber = factory(model, device, compute_type)
        transcriber.warmup()
        return transcriber
    except Exception as exc:
        if requested_device != "auto" or device == "cpu":
            raise TranscriptionError(f"device_unusable:{device}") from exc
        device, compute_type = "cpu", "int8"
        transcriber = factory(model, device, compute_type)
        transcriber.warmup()
        return transcriber


async def serve(args: argparse.Namespace) -> int:  # pragma: no cover - I/O 진입점
    import websockets

    transcriber = build_transcriber(
        args.model,
        args.device,
        factory=lambda model, device, compute_type: FasterWhisperTranscriber(
            model, device=device, compute_type=compute_type
        ),
    )
    loop = asyncio.get_running_loop()

    async def run_blocking(call: Callable[[], Transcription]) -> Transcription:
        return await loop.run_in_executor(None, call)

    async def handler(socket: Any) -> None:
        try:
            options = parse_stream_options(getattr(socket, "path", "") or "")
        except ValueError as exc:
            await socket.send(json.dumps({"type": "error", "detail": str(exc)}))
            await socket.close()
            return
        transcriber.configure_language(options["language"])

        async def send(payload: dict[str, Any]) -> None:
            await socket.send(json.dumps(payload, ensure_ascii=False))

        session = StreamSession(
            options=options,
            transcriber=transcriber,
            send=send,
            run_blocking=run_blocking,
        )
        await send(
            {
                "type": "ready",
                "model": transcriber.model_name,
                "device": transcriber.device,
                "compute_type": transcriber.compute_type,
                "language": options["language"],
                "sample_rate": options["sample_rate"],
            }
        )
        try:
            async for message in socket:
                if isinstance(message, bytes):
                    await session.feed(message)
                    continue
                try:
                    control = json.loads(message)
                except (TypeError, ValueError):
                    continue
                kind = str((control or {}).get("type") or "")
                if kind == "Finalize":
                    await session.finalize(speech_final=True)
                elif kind == "CloseStream":
                    await session.finalize(speech_final=True)
                    break
        except TranscriptionError as exc:
            await send({"type": "error", "detail": exc.args[0] if exc.args else "error"})
        finally:
            with_close = getattr(socket, "close", None)
            if callable(with_close):
                await socket.close()

    print(
        json.dumps(
            {
                "ready": True,
                "host": args.host,
                "port": args.port,
                "model": transcriber.model_name,
                "device": transcriber.device,
                "compute_type": transcriber.compute_type,
            },
            separators=(",", ":"),
        ),
        flush=True,
    )
    async with websockets.serve(
        handler,
        args.host,
        args.port,
        max_size=MAX_FRAME_BYTES,
        max_queue=16,
        ping_interval=20,
        ping_timeout=20,
    ):
        await asyncio.Future()
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--host", default=DEFAULT_HOST, choices=[DEFAULT_HOST, "localhost"])
    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
    parser.add_argument("--model", default=DEFAULT_MODEL, choices=list(ALLOWED_MODELS))
    parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
    parser.add_argument("--enable", action="store_true")
    return parser


def main(argv: Iterable[str] | None = None) -> int:  # pragma: no cover - CLI
    args = build_parser().parse_args(list(argv) if argv is not None else None)
    if not args.enable:
        print("local whisper STT server is disabled; pass --enable", file=sys.stderr)
        return 2
    started = time.monotonic()
    try:
        return asyncio.run(serve(args))
    except KeyboardInterrupt:
        print(
            json.dumps({"stopped": True, "uptime_seconds": round(time.monotonic() - started, 1)}),
            file=sys.stderr,
        )
        return 0


if __name__ == "__main__":  # pragma: no cover - CLI
    raise SystemExit(main())