#!/usr/bin/env python3
"""
meetrec.py — record an in-person meeting and transcribe it with Whisper.

Records from any PipeWire/ALSA input (microphone, USB conference mic,
USB audio interface, or a Pulse/PipeWire sink "monitor") while optionally
printing a live rolling transcript. On Ctrl+C it stops cleanly and runs a
final, higher-quality full transcription.

Outputs (with -o meeting):
    meeting.wav      raw recording, 16 kHz mono WAV
    meeting.live.srt live rolling transcript (only when --live > 0)
    meeting.txt      final transcript
    meeting.srt      final transcript with timestamps
    meeting.json     final transcript, structured

Setup on Arch:
    sudo pacman -S python pipewire pipewire-pulse
    python -m venv .venv && source .venv/bin/activate
    pip install numpy sounddevice faster-whisper

Usage:
    python meetrec.py --list-sources
    python meetrec.py -s "Conference Mic" --model small --live 8 -o meeting
    python meetrec.py -s monitor --model base --language en   # record system audio

    Ctrl+C to stop -> final transcripts are written.
"""

import argparse
import json
import queue
import shutil
import subprocess
import sys
import threading
import time
import wave

import numpy as np
import sounddevice as sd
from faster_whisper import WhisperModel

SR = 16000          # sample rate — what Whisper wants
LIVE_WINDOW = 60    # seconds of rolling audio kept for live transcription


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

def fmt_hms(t: float) -> str:
    h, rem = divmod(int(t), 3600)
    m, s = divmod(rem, 60)
    return f"{h:02d}:{m:02d}:{s:02d}"


def fmt_srt(t: float) -> str:
    ms = int(round((t - int(t)) * 1000))
    h, rem = divmod(int(t) + (1 if ms >= 1000 else 0), 3600)
    m, s = divmod(rem, 60)
    if ms >= 1000:
        ms -= 1000
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def write_srt(path: str, items) -> None:
    with open(path, "w") as f:
        for i, (a, b, text) in enumerate(items, 1):
            f.write(f"{i}\n{fmt_srt(a)} --> {fmt_srt(b)}\n{text}\n\n")


def list_sources() -> None:
    print("PortAudio (PipeWire/Pulse/ALSA) input devices:\n")
    for i, d in enumerate(sd.query_devices()):
        if d["max_input_channels"] > 0:
            print(f"  {i:>3}  {d['name']}\n"
                  f"        api={d['hostapi']}  in={d['max_input_channels']}ch  "
                  f"rate={d['default_samplerate']:.0f}")
    pactl = shutil.which("pactl")
    if pactl:
        print("\nPulseAudio/PipeWire sources (pactl):\n")
        try:
            out = subprocess.run(
                [pactl, "list", "short", "sources"],
                capture_output=True, text=True, timeout=5).stdout.strip()
            for line in out.splitlines():
                print("  " + line)
        except Exception:
            pass
    print("\nTip: sink monitors (record what the computer plays) appear as "
          "\"<sink name>.monitor\". Match them with -s <part of name>.")


def resolve_source(name: str):
    """Return a sounddevice device index, or None for the default device."""
    if name in ("default", None, ""):
        return None
    inputs = [(i, d) for i, d in enumerate(sd.query_devices())
              if d["max_input_channels"] > 0]
    matches = [(i, d) for i, d in inputs if name.lower() in d["name"].lower()]
    if not matches:
        print(f"error: no input device matching {name!r}. "
              f"Run with --list-sources.", file=sys.stderr)
        sys.exit(1)
    if len(matches) > 1:
        print(f"note: multiple devices match {name!r}, using: "
              f"{matches[0][1]['name']}", file=sys.stderr)
    idx, dev = matches[0]
    print(f"recording from: [{idx}] {dev['name']} ({dev['hostapi']})")
    return idx


# --------------------------------------------------------------------------
# Recording
# --------------------------------------------------------------------------

class Recorder:
    """Records mono 16 kHz float32 to a WAV file and keeps a rolling buffer."""

    def __init__(self, device, wav_path: str):
        self.device = device
        self.wav_path = wav_path
        self.q: "queue.Queue" = queue.Queue()
        self.lock = threading.Lock()
        self.buf = np.zeros(0, dtype=np.float32)   # rolling window
        self.base = 0                              # samples recorded before buf
        self.total = 0                             # total samples recorded
        self.peak = 0.0                            # peak amplitude of last block
        self.stop_event = threading.Event()
        self.stream = None
        self._writer = threading.Thread(target=self._write, daemon=True)
        self._writer.start()

    def _callback(self, indata, frames, time_info, status):
        if status:
            print(f"[rec] status: {status}", file=sys.stderr)
        self.peak = float(np.max(np.abs(indata[:, 0]))) if frames else 0.0
        self.q.put(indata.copy())
        with self.lock:
            block = indata[:, 0].astype(np.float32, copy=False)
            self.buf = np.append(self.buf, block)
            self.total += len(block)
            limit = LIVE_WINDOW * SR
            if len(self.buf) > limit:
                drop = len(self.buf) - limit
                self.buf = self.buf[drop:]
                self.base += drop

    def duration(self) -> float:
        with self.lock:
            return self.total / SR

    def _write(self):
        with wave.open(self.wav_path, "wb") as w:
            w.setnchannels(1)
            w.setsampwidth(2)
            w.setframerate(SR)
            while True:
                block = self.q.get()
                if block is None:
                    break
                pcm = (np.clip(block[:, 0], -1.0, 1.0) * 32767).astype(np.int16)
                w.writeframes(pcm.tobytes())

    def snapshot(self):
        """Return (rolling audio, sample offset of its start)."""
        with self.lock:
            return self.buf.copy(), self.base

    def start(self):
        self.stream = sd.InputStream(
            samplerate=SR, channels=1, dtype="float32",
            device=self.device, callback=self._callback)
        self.stream.start()

    def stop(self):
        self.stop_event.set()
        self.q.put(None)
        if self.stream:
            self.stream.stop()
            self.stream.close()
        # The writer must drain the queue and close the WAV before we return,
        # otherwise the file can be truncated. Join in slices instead of one
        # long timeout so a slow disk still gets the full flush; only warn
        # after a total of 60 s (effectively never in practice).
        waited = 0.0
        while self._writer.is_alive() and waited < 60:
            self._writer.join(timeout=2.0)
            waited += 2.0
        if self._writer.is_alive():
            print("warning: WAV writer still running after 60 s; "
                  "the file may be truncated", file=sys.stderr)


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

def live_worker(model, rec: Recorder, args, live_srt_path: str,
                stop: threading.Event, on_line=None):
    """Transcribes the rolling window every `args.live` seconds and appends
    newly-finalized segments to the live .srt. If `on_line` is given, new
    lines are delivered to it (a, text) instead of being printed."""
    def emit(a, text):
        if on_line is not None:
            on_line(a, text)
        else:
            print(f"[{fmt_hms(a)}] {text}", flush=True)

    items = []          # (abs_start, abs_end, text) for the live .srt
    last_end = 0.0
    if on_line is None:
        print("live transcription on (Ctrl+C to stop)", flush=True)

    while not stop.is_set():
        stop.wait(args.live)
        if stop.is_set():
            break
        buf, base = rec.snapshot()
        if len(buf) < 2 * SR:      # need at least 2 s of audio
            continue
        try:
            segs, _info = model.transcribe(
                buf, beam_size=1, vad_filter=True, language=args.language)
        except Exception as e:
            print(f"[live] transcription error: {e}", file=sys.stderr)
            continue

        base_s = base / SR
        added = False
        for s in segs:
            a, b = base_s + s.start, base_s + s.end
            text = s.text.strip()
            if not text or b < last_end:
                continue
            items.append((a, b, text))
            last_end = max(last_end, b)
            emit(a, text)
            added = True
        if added:
            write_srt(live_srt_path, items)


def final_transcribe(model, wav_path: str, out_stem: str, language,
                     on_info=None, on_segment=None, on_done=None):
    """Full-pass transcription of the whole recording.

    on_info(msg), on_segment(start, end, text) and on_done() are optional
    callbacks; when absent the same progress is printed to stdout."""
    def info(msg):
        if on_info is not None:
            on_info(msg)
        else:
            print(msg, flush=True)

    info("\nRunning final transcription (this can take a while)...")
    segments, info_ = model.transcribe(
        wav_path, beam_size=5, vad_filter=True, language=language,
        vad_parameters=dict(min_silence_duration_ms=500))

    info(f"language: {info_.language} (p={info_.language_probability:.2f})  "
         f"duration: {fmt_hms(info_.duration)}")

    items = []
    for s in segments:            # generator: streams as it decodes
        text = s.text.strip()
        items.append((s.start, s.end, text))
        if on_segment is not None:
            on_segment(s.start, s.end, text)
        else:
            print(f"[{fmt_hms(s.start)}] {text}", flush=True)

    full = " ".join(t for _, _, t in items)
    with open(out_stem + ".txt", "w") as f:
        f.write(full + "\n")
    write_srt(out_stem + ".srt", items)
    with open(out_stem + ".json", "w") as f:
        json.dump({
            "language": info_.language,
            "duration": round(info_.duration, 2),
            "segments": [{"start": round(a, 2), "end": round(b, 2), "text": t}
                         for a, b, t in items],
        }, f, ensure_ascii=False, indent=2)
    info(f"\nwrote: {out_stem}.wav  {out_stem}.txt  {out_stem}.srt  {out_stem}.json")
    if on_done is not None:
        on_done()


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

def main():
    ap = argparse.ArgumentParser(
        description="Record a meeting and transcribe it with Whisper "
                    "(PipeWire / ALSA via PortAudio).")
    ap.add_argument("-s", "--source", default="default",
                    help="input to record: 'default' or a substring of a "
                         "device name (see --list-sources)")
    ap.add_argument("-m", "--model", default="small",
                    help="Whisper model: tiny/base/small/medium/large-v3 "
                         "(default: small)")
    ap.add_argument("-o", "--output", default="meeting",
                    help="output file stem (default: meeting)")
    ap.add_argument("--live", type=int, default=8, metavar="SEC",
                    help="live transcription interval in seconds; 0 disables "
                         "(default: 8)")
    ap.add_argument("--language", default=None,
                    help="force language code, e.g. en, de (default: autodetect)")
    ap.add_argument("--device", default="auto",
                    help="compute device: auto / cpu / cuda (default: auto)")
    ap.add_argument("--compute-type", default="int8",
                    help="int8 / int8_float16 / float16 / float32 (default: int8)")
    ap.add_argument("--list-sources", action="store_true",
                    help="list available input devices and exit")
    args = ap.parse_args()

    if args.list_sources:
        list_sources()
        return

    device = resolve_source(args.source)

    print(f"loading Whisper model {args.model!r} "
          f"(first run downloads it to ~/.cache/huggingface)...", flush=True)
    model = WhisperModel(args.model, device=args.device,
                         compute_type=args.compute_type)

    rec = Recorder(device, args.output + ".wav")
    rec.start()
    print(f"recording to {args.output}.wav  (Ctrl+C to stop)", flush=True)

    stop = threading.Event()
    worker = None
    if args.live > 0:
        worker = threading.Thread(
            target=live_worker, args=(model, rec, args,
                                      args.output + ".live.srt", stop),
            daemon=True)
        worker.start()

    try:
        while not stop.is_set():
            stop.wait(0.5)
    except KeyboardInterrupt:
        print("\nstopping...", flush=True)
    finally:
        stop.set()
        rec.stop()
        if worker:
            worker.join(timeout=60)
        final_transcribe(model, args.output + ".wav", args.output, args.language)


if __name__ == "__main__":
    main()