"""
Companion process for the proof-of-personhood talk.

What it does:
  * Captures mic audio at 16 kHz, mono.
  * Uses webrtcvad to detect utterance boundaries (start of speech →
    sustained silence → close).
  * Transcribes each completed utterance with faster-whisper.
  * Broadcasts results over a local WebSocket (ws://localhost:8787) so
    the Slidev panel can render them live.

Two guarantees the architecture preserves:
  1. Audio is ALWAYS recorded to disk (WAV) regardless of STT state. If
     the transcription pipeline fails mid-talk, the audio is still safe
     and can be re-transcribed + signed post-hoc. Live failure becomes
     a UI glitch, not lost content.
  2. JSONL of every final transcript is also written, with wall-clock
     timestamps. That's the source of truth for the receipt later.

Usage:
  python main.py --list-mics
  python main.py --model small               # pick model size
  python main.py --device 2                  # pick mic
  python main.py --replay recordings/foo.wav # replay a WAV instead of mic
"""

from __future__ import annotations

import argparse
import asyncio
import concurrent.futures
import json
import queue
import signal
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import sounddevice as sd
import soundfile as sf
import webrtcvad
import websockets
from faster_whisper import WhisperModel


# ---- Audio constants ------------------------------------------------------

SAMPLE_RATE = 16_000               # whisper wants 16 kHz mono
FRAME_MS = 30                       # webrtcvad accepts 10 / 20 / 30 ms
FRAME_SAMPLES = SAMPLE_RATE * FRAME_MS // 1000  # = 480

# ---- VAD chunking thresholds ---------------------------------------------

# A short burst of speech frames to begin an utterance; prevents single
# noise frames from triggering a transcription run.
START_THRESHOLD_FRAMES = 3          # ~90 ms

# How much silence closes an utterance. Tunable: shorter = lower latency
# but choppier; longer = smoother but slower.
END_THRESHOLD_FRAMES = 17           # ~510 ms

# Hard cap on utterance length so a long monologue without pauses still
# gets transcribed in chunks.
MAX_UTTERANCE_FRAMES = SAMPLE_RATE * 12 // FRAME_SAMPLES  # ~12 s

# Frames of audio kept before an utterance starts, so the first word
# isn't clipped if VAD trips a frame late.
PREROLL_FRAMES = 7                  # ~210 ms

# How often we flush the WAV file to disk while recording.
WAV_FLUSH_EVERY_FRAMES = 100        # ~3 s


# ---- Helpers --------------------------------------------------------------


def list_mics() -> None:
    """Print available input devices."""
    print("Input devices:")
    for i, d in enumerate(sd.query_devices()):
        if d["max_input_channels"] > 0:
            print(
                f"  [{i:>2}] {d['name']}  "
                f"({int(d['default_samplerate'])} Hz, {d['max_input_channels']} ch)"
            )


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")


# ---- Recorder -------------------------------------------------------------


class Recorder:
    """Always-on WAV + JSONL writer. Independent of STT."""

    def __init__(self, out_dir: Path) -> None:
        out_dir.mkdir(parents=True, exist_ok=True)
        stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        self.wav_path = out_dir / f"talk-{stamp}.wav"
        self.jsonl_path = out_dir / f"talk-{stamp}.jsonl"
        self.wav = sf.SoundFile(
            str(self.wav_path),
            mode="w",
            samplerate=SAMPLE_RATE,
            channels=1,
            subtype="PCM_16",
        )
        self.jsonl = open(self.jsonl_path, "a", buffering=1, encoding="utf-8")
        self._frames_since_flush = 0

    def write_audio(self, samples_int16: np.ndarray) -> None:
        self.wav.write(samples_int16)
        self._frames_since_flush += 1
        if self._frames_since_flush >= WAV_FLUSH_EVERY_FRAMES:
            self.wav.flush()
            self._frames_since_flush = 0

    def write_event(self, event: dict) -> None:
        self.jsonl.write(json.dumps(event) + "\n")

    def close(self) -> None:
        try:
            self.wav.close()
        except Exception:
            pass
        try:
            self.jsonl.close()
        except Exception:
            pass


# ---- WebSocket broadcaster ------------------------------------------------


class Broadcaster:
    """Tracks connected clients and broadcasts JSON events to all of them."""

    def __init__(self, history_limit: int = 50) -> None:
        self.clients: set = set()
        self.history: list[dict] = []
        self.history_limit = history_limit

    async def register(self, ws) -> None:
        self.clients.add(ws)
        # Late joiners get the last 10 events so the panel isn't empty.
        for event in self.history[-10:]:
            try:
                await ws.send(json.dumps(event))
            except Exception:
                pass

    async def unregister(self, ws) -> None:
        self.clients.discard(ws)

    async def broadcast(self, event: dict) -> None:
        self.history.append(event)
        if len(self.history) > self.history_limit:
            self.history = self.history[-self.history_limit :]
        message = json.dumps(event)
        dead = []
        for ws in list(self.clients):
            try:
                await ws.send(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.clients.discard(ws)


# ---- VAD chunker (runs in a background thread) ----------------------------


def vad_chunker(
    audio_q: "queue.Queue",
    vad: webrtcvad.Vad,
    on_utterance,
) -> None:
    """
    Pulls int16 frames off `audio_q` and calls `on_utterance(samples, t_start, t_end)`
    whenever VAD detects a complete utterance. Runs forever; exit on a None
    sentinel in the queue.
    """
    preroll: list[np.ndarray] = []
    speech_streak = 0
    silence_streak = 0
    in_utterance = False
    utterance_buf: list[np.ndarray] = []
    t_start_frame = 0
    frame_idx = 0

    while True:
        try:
            frame = audio_q.get(timeout=0.5)
        except queue.Empty:
            continue
        if frame is None:
            return

        frame_idx += 1
        try:
            is_speech = vad.is_speech(frame.tobytes(), SAMPLE_RATE)
        except Exception:
            # webrtcvad raises if the frame size is wrong. Skip and continue.
            continue

        if not in_utterance:
            preroll.append(frame)
            if len(preroll) > PREROLL_FRAMES:
                preroll.pop(0)

            if is_speech:
                speech_streak += 1
                if speech_streak >= START_THRESHOLD_FRAMES:
                    in_utterance = True
                    t_start_frame = frame_idx - len(preroll)
                    utterance_buf = list(preroll)
                    preroll = []
                    silence_streak = 0
            else:
                speech_streak = 0
            continue

        # in_utterance
        utterance_buf.append(frame)
        if is_speech:
            silence_streak = 0
        else:
            silence_streak += 1

        should_close = (
            silence_streak >= END_THRESHOLD_FRAMES
            or len(utterance_buf) >= MAX_UTTERANCE_FRAMES
        )
        if should_close:
            samples = np.concatenate(utterance_buf)
            t_start = t_start_frame * FRAME_MS / 1000.0
            t_end = frame_idx * FRAME_MS / 1000.0
            try:
                on_utterance(samples, t_start, t_end)
            except Exception as e:
                print(f"[vad] on_utterance error: {e}", file=sys.stderr)
            in_utterance = False
            utterance_buf = []
            preroll = []
            speech_streak = 0
            silence_streak = 0


# ---- Transcription --------------------------------------------------------


def transcribe(model: WhisperModel, samples_int16: np.ndarray, language: str) -> str:
    """Run faster-whisper on a chunk of audio. Returns plain text."""
    audio_float = samples_int16.astype(np.float32) / 32768.0
    segments, _ = model.transcribe(
        audio_float,
        language=language,
        vad_filter=False,            # we already VAD'd
        beam_size=1,                 # greedy: lower quality, much faster
        condition_on_previous_text=False,
        no_speech_threshold=0.5,
    )
    return "".join(seg.text for seg in segments).strip()


# ---- Main -----------------------------------------------------------------


async def amain() -> None:
    parser = argparse.ArgumentParser(description="Speech-to-text companion for the talk.")
    parser.add_argument("--list-mics", action="store_true", help="List input devices and exit.")
    parser.add_argument("--device", type=int, default=None, help="Input device index.")
    parser.add_argument("--model", default="small",
                        help="Model size: tiny | base | small | medium | large-v3.")
    parser.add_argument("--compute-type", default="int8",
                        help="faster-whisper compute_type (int8 | int8_float16 | float16 | float32).")
    parser.add_argument("--device-type", default="cpu",
                        help="Compute device (cpu | cuda).")
    parser.add_argument("--language", default="en", help="Language hint.")
    parser.add_argument("--out", default="recordings", help="Output directory.")
    parser.add_argument("--ws-port", type=int, default=8787, help="WebSocket port.")
    parser.add_argument("--vad", type=int, default=2, choices=[0, 1, 2, 3],
                        help="webrtcvad aggressiveness (0=lax .. 3=strict).")
    parser.add_argument("--replay", default=None,
                        help="Replay a WAV file at real-time speed instead of mic capture.")
    args = parser.parse_args()

    if args.list_mics:
        list_mics()
        return

    # ---- Boot the model ----
    print(f"[boot] loading whisper model='{args.model}' compute_type='{args.compute_type}' "
          f"device='{args.device_type}'", flush=True)
    t0 = time.time()
    model = WhisperModel(args.model, device=args.device_type, compute_type=args.compute_type)
    print(f"[boot] model loaded in {time.time() - t0:.1f}s", flush=True)

    print("[boot] warming up...", flush=True)
    warmup = np.zeros(SAMPLE_RATE, dtype=np.float32)  # 1 s of silence
    list(model.transcribe(warmup, language=args.language))
    print("[boot] warmed.", flush=True)

    # ---- Infra ----
    recorder = Recorder(Path(args.out))
    broadcaster = Broadcaster()
    vad = webrtcvad.Vad(args.vad)
    audio_q: "queue.Queue" = queue.Queue()
    loop = asyncio.get_running_loop()

    # Single-worker executor for transcription so slow utterances serialize
    # instead of fighting for CPU. The default executor is shared with WAV
    # writes etc.; keeping STT isolated makes latency predictable.
    transcribe_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1,
                                                            thread_name_prefix="stt")

    # ---- Audio source ----
    if args.replay:
        async def replay_loop() -> None:
            data, sr = sf.read(args.replay, dtype="int16")
            if data.ndim > 1:
                data = data[:, 0]
            if sr != SAMPLE_RATE:
                print(f"[replay] WAV is {sr} Hz; resampling not implemented. "
                      f"Convert with: ffmpeg -i in.wav -ar {SAMPLE_RATE} -ac 1 out.wav",
                      file=sys.stderr)
                return
            print(f"[replay] {args.replay}: {len(data)/SAMPLE_RATE:.1f} s")
            for i in range(0, len(data), FRAME_SAMPLES):
                frame = data[i : i + FRAME_SAMPLES]
                if len(frame) < FRAME_SAMPLES:
                    break
                audio_q.put(frame)
                # No recorder.write_audio() in replay — the WAV is the source.
                await asyncio.sleep(FRAME_MS / 1000.0)
            print("[replay] done.")

        asyncio.create_task(replay_loop())
        stream = None
    else:
        def audio_callback(indata, frames, time_info, status) -> None:
            if status:
                print(f"[mic] status: {status}", file=sys.stderr)
            # float32 [-1,1] → int16 PCM
            samples_int16 = (indata[:, 0] * 32767.0).astype(np.int16)
            audio_q.put(samples_int16)
            recorder.write_audio(samples_int16)

        stream = sd.InputStream(
            samplerate=SAMPLE_RATE,
            channels=1,
            dtype="float32",
            callback=audio_callback,
            blocksize=FRAME_SAMPLES,
            device=args.device,
        )
        stream.start()
        dev_name = "default" if args.device is None else f"device #{args.device}"
        print(f"[mic] capturing from {dev_name} at {SAMPLE_RATE} Hz, frame={FRAME_MS} ms",
              flush=True)

    # ---- Utterance handling ----
    def on_utterance(samples_int16: np.ndarray, t_start: float, t_end: float) -> None:
        # Called from the VAD thread. Schedule async work on the main loop.
        fut = asyncio.run_coroutine_threadsafe(
            handle_utterance(samples_int16, t_start, t_end), loop
        )
        # Drop the future (errors will be reported by handle_utterance itself).
        del fut

    async def handle_utterance(samples_int16, t_start, t_end) -> None:
        try:
            t_inf_start = time.time()
            text = await loop.run_in_executor(
                transcribe_pool, transcribe, model, samples_int16, args.language
            )
            t_inf = time.time() - t_inf_start
        except Exception as e:
            print(f"[stt] error: {e}", file=sys.stderr)
            return

        if not text:
            return
        dur = t_end - t_start
        event = {
            "type": "final",
            "text": text,
            "t_start": round(t_start, 3),
            "t_end": round(t_end, 3),
            "duration": round(dur, 3),
            "inference_ms": int(t_inf * 1000),
            "wall_ts": utc_now(),
        }
        print(f"  ⇢ [{t_start:6.1f}s  +{dur:.1f}s  stt={int(t_inf*1000):>4}ms]  {text}",
              flush=True)
        recorder.write_event(event)
        await broadcaster.broadcast(event)

    # ---- VAD chunker thread ----
    loop.run_in_executor(None, vad_chunker, audio_q, vad, on_utterance)

    # ---- WebSocket server ----
    async def ws_handler(ws) -> None:
        peer = getattr(ws, "remote_address", "?")
        print(f"[ws] client connected: {peer}", flush=True)
        await broadcaster.register(ws)
        try:
            async for _msg in ws:
                pass  # ignore inbound messages
        except Exception:
            pass
        finally:
            await broadcaster.unregister(ws)
            print(f"[ws] client disconnected: {peer}", flush=True)

    server = await websockets.serve(ws_handler, "localhost", args.ws_port)
    print(f"[ws] listening on ws://localhost:{args.ws_port}", flush=True)
    print(f"[out] {recorder.wav_path}")
    print(f"[out] {recorder.jsonl_path}")
    print("[ready] speak into the mic — Ctrl+C to stop.\n", flush=True)

    # ---- Shutdown plumbing ----
    stop_event = asyncio.Event()

    def request_stop() -> None:
        if not stop_event.is_set():
            print("\n[shutdown] stopping...", flush=True)
            stop_event.set()

    try:
        loop.add_signal_handler(signal.SIGINT, request_stop)
        loop.add_signal_handler(signal.SIGTERM, request_stop)
    except NotImplementedError:
        # Windows: fall back to default KeyboardInterrupt handling.
        pass

    await stop_event.wait()

    # ---- Cleanup ----
    audio_q.put(None)                       # poison the VAD thread
    server.close()
    await server.wait_closed()
    if stream is not None:
        stream.stop()
        stream.close()
    transcribe_pool.shutdown(wait=False, cancel_futures=True)
    recorder.close()
    print("[shutdown] done.")


def main() -> None:
    try:
        asyncio.run(amain())
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()
