#!/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, mono WAV at the input's native rate
    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

Audio is recorded at the input device's native sample rate; both the live
window and the engines get it resampled to Whisper's 16 kHz internally.

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
    python meetrec.py --engine whisper-cpp -o meeting         # GPU via Vulkan (AMD/NVIDIA)

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

import argparse
import json
import os
import queue
import shutil
import subprocess
import sys
import tempfile
import threading
import time
import urllib.request
import wave
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import sounddevice as sd

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 resample_16k(audio: np.ndarray, rate: int) -> np.ndarray:
    """Linearly resample mono float32 audio to 16 kHz (Whisper's rate).

    Only needed for the live rolling buffer, which is fed to the engines as
    a raw array (they assume 16 kHz there). WAV files are resampled by the
    engines themselves.
    """
    audio = np.asarray(audio, dtype=np.float32)
    if rate == SR or audio.size == 0:
        return audio
    n_out = max(1, int(round(audio.size * SR / rate)))
    t_out = np.arange(n_out, dtype=np.float64) * (rate / SR)
    return np.interp(t_out, np.arange(audio.size, dtype=np.float64),
                     audio).astype(np.float32)


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 float32 at the input's native rate to a WAV file and
    keeps a rolling buffer for live transcription."""

    def __init__(self, device, wav_path: str):
        self.device = device
        self.wav_path = wav_path
        dev = sd.query_devices(kind="input") if device is None \
            else sd.query_devices(device)
        self.rate = int(dev["default_samplerate"])   # input's native rate
        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 * self.rate
            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 / self.rate

    def _write(self):
        with wave.open(self.wav_path, "wb") as w:
            w.setnchannels(1)
            w.setsampwidth(2)
            w.setframerate(self.rate)
            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 window resampled to 16 kHz, its start time in s)."""
        with self.lock:
            return resample_16k(self.buf, self.rate), self.base / self.rate

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

    def stop(self):
        if self.stop_event.is_set():
            return          # idempotent: the GUI close path may call this twice
        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_s = rec.snapshot()   # 16 kHz audio, absolute start time
        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

        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()


# --------------------------------------------------------------------------
# Engines
# --------------------------------------------------------------------------
#
# Two interchangeable transcription backends:
#   faster-whisper  CTranslate2 — NVIDIA CUDA or CPU (default)
#   whisper-cpp     whisper.cpp CLI — GPU via Vulkan (AMD, Intel, NVIDIA)
#                   or CPU; preferred on AMD hardware (e.g. Ryzen AI laptops)
#
# Both expose:  transcribe(audio, beam_size, vad_filter, language,
#                        vad_parameters) -> (segments, info)
# where `audio` is a 1-D float32 @16 kHz array or a WAV path, segments have
# .start/.end/.text and info has .language/.language_probability/.duration.

WHISPER_CPP_MODEL_FILES = {
    "tiny": "ggml-tiny.bin",
    "base": "ggml-base.bin",
    "small": "ggml-small.bin",
    "medium": "ggml-medium.bin",
    "large-v3": "ggml-large-v3.bin",
}
WHISPER_CPP_MODEL_URL = ("https://huggingface.co/ggerganov/whisper.cpp/"
                         "resolve/main/{name}")
WHISPER_CPP_VAD_FILE = "ggml-silero-v6.2.0.bin"
WHISPER_CPP_VAD_URL = ("https://huggingface.co/ggml-org/whisper-vad/"
                       "resolve/main/" + WHISPER_CPP_VAD_FILE)


def _download_file(url: str, dest: Path) -> None:
    tmp = dest.with_suffix(dest.suffix + ".part")
    req = urllib.request.Request(url, headers={"User-Agent": "meetrec"})
    with urllib.request.urlopen(req, timeout=60) as r:
        total = int(r.headers.get("Content-Length") or 0)
        done = 0
        with open(tmp, "wb") as f:
            while True:
                chunk = r.read(1024 * 1024)
                if not chunk:
                    break
                f.write(chunk)
                done += len(chunk)
                if total:
                    print(f"\r  {done/1e6:.0f}/{total/1e6:.0f} MB", end="",
                          flush=True)
    if total:
        print(flush=True)
    os.replace(tmp, dest)


def find_whisper_cpp_bin() -> str:
    """Locate the whisper.cpp 'whisper-cli' binary."""
    candidates = [
        os.environ.get("WHISPER_CPP_BIN"),
        shutil.which("whisper-cli"),
        str(Path.home() / ".local/share/meetrec/whisper-cpp/bin/whisper-cli"),
    ]
    for c in candidates:
        if c and os.path.isfile(c) and os.access(c, os.X_OK):
            return c
    raise RuntimeError(
        "whisper-cli not found. Build whisper.cpp (e.g. './install.sh "
        "--whisper-cpp', or cmake -DWHISPER_VULKAN=ON) or set WHISPER_CPP_BIN "
        "to the binary.")


def _whisper_cpp_cache() -> Path:
    cache = Path(os.environ.get(
        "MEETREC_CACHE", str(Path.home() / ".cache/meetrec")))
    d = cache / "whisper-cpp"
    d.mkdir(parents=True, exist_ok=True)
    return d


def _cached_ggml(fname: str, url: str) -> str:
    """Return the path of a cached GGML file, downloading it if missing."""
    p = _whisper_cpp_cache() / fname
    if p.exists() and p.stat().st_size > 0:
        return str(p)
    print(f"downloading {fname} to {p} ...", flush=True)
    _download_file(url, p)
    return str(p)


def whisper_cpp_model_path(model_name: str) -> str:
    """Return the local path of a whisper.cpp GGML model, downloading it."""
    try:
        fname = WHISPER_CPP_MODEL_FILES[model_name]
    except KeyError:
        raise RuntimeError(
            f"unknown model {model_name!r}; use one of: "
            f"{', '.join(WHISPER_CPP_MODEL_FILES)}")
    return _cached_ggml(fname, WHISPER_CPP_MODEL_URL.format(name=fname))


def write_wav_float(path: str, audio: np.ndarray) -> None:
    """Write a 1-D float32 mono @16 kHz array as 16-bit PCM WAV."""
    audio = np.asarray(audio, dtype=np.float32).ravel()
    pcm = (np.clip(audio, -1.0, 1.0) * 32767).astype(np.int16)
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(SR)
        w.writeframes(pcm.tobytes())


def _parse_ts(s: str) -> float:
    """whisper.cpp timestamp 'HH:MM:SS,mmm' -> seconds."""
    h, m, rest = s.split(":")
    sec, ms = rest.split(",")
    return int(h) * 3600 + int(m) * 60 + int(sec) + int(ms) / 1000.0


class WhisperCppEngine:
    """Transcription via the whisper.cpp 'whisper-cli' binary.

    Build whisper.cpp with -DWHISPER_VULKAN=ON to run on a GPU (Radeon,
    Intel, or NVIDIA); it falls back to CPU otherwise. GPU vs. CPU is decided
    by the build, so --device/--compute-type do not apply. vad_filter=True
    enables whisper.cpp's Silero VAD (auto-downloaded); vad_parameters are
    accepted but not mapped.
    """

    def __init__(self, model_name: str, cli: str = None):
        self.model_name = model_name
        self.cli = cli or find_whisper_cpp_bin()
        self.model_path = whisper_cpp_model_path(model_name)

    def transcribe(self, audio, beam_size=5, vad_filter=False,
                   language=None, vad_parameters=None):
        # vad_parameters accepted for interface compatibility; Silero VAD uses
        # whisper.cpp's own defaults
        with tempfile.TemporaryDirectory(prefix="meetrec-") as td:
            if isinstance(audio, (str, os.PathLike)):
                wav = str(audio)
                with wave.open(wav, "rb") as w:
                    duration = w.getnframes() / w.getframerate()
            else:
                wav = os.path.join(td, "in.wav")
                write_wav_float(wav, audio)
                duration = len(audio) / SR

            cmd = [self.cli, "-m", self.model_path, "-f", wav,
                   "-oj", "-of", os.path.join(td, "out"),
                   "-bs", str(max(1, int(beam_size))), "-np",
                   "-l", language or "auto"]
            if vad_filter:
                cmd += ["--vad", "-vm", _cached_ggml(
                    WHISPER_CPP_VAD_FILE, WHISPER_CPP_VAD_URL)]
            proc = subprocess.run(cmd, capture_output=True, text=True)
            if proc.returncode != 0:
                raise RuntimeError(
                    (proc.stderr or proc.stdout).strip()
                    or f"whisper-cli exited {proc.returncode}")
            with open(os.path.join(td, "out.json")) as f:
                data = json.load(f)

        result = data.get("result", {})
        segs = []
        for item in data.get("transcription", []):
            off = item.get("offsets", {})
            if "from" in off and "to" in off:      # milliseconds (current)
                a, b = off["from"] / 1000.0, off["to"] / 1000.0
            else:                                    # 'HH:MM:SS,mmm' (older)
                ts = item.get("timestamps", {})
                a = _parse_ts(ts.get("from", "00:00:00,000"))
                b = _parse_ts(ts.get("to", "00:00:00,000"))
            text = item.get("text", "").strip()
            if text:
                segs.append(SimpleNamespace(start=a, end=b, text=text))
        info = SimpleNamespace(
            language=result.get("language") or language or "unknown",
            language_probability=1.0,   # not reported by whisper.cpp
            duration=duration,
        )
        return segs, info


def load_engine(model_name, engine="faster-whisper",
                device="auto", compute_type="int8"):
    """Create a transcription engine. See the Engines section above."""
    if engine == "whisper-cpp":
        return WhisperCppEngine(model_name)
    if engine == "faster-whisper":
        from faster_whisper import WhisperModel
        return WhisperModel(model_name, device=device,
                            compute_type=compute_type)
    raise RuntimeError(f"unknown engine {engine!r}; use 'faster-whisper' "
                       "or 'whisper-cpp'")


# --------------------------------------------------------------------------
# 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("--engine", default="faster-whisper",
                    choices=["faster-whisper", "whisper-cpp"],
                    help="transcription backend (default: faster-whisper). "
                         "whisper-cpp can use a GPU via Vulkan — preferred "
                         "on AMD hardware")
    ap.add_argument("--device", default="auto",
                    help="compute device: auto / cpu / cuda (default: auto); "
                         "faster-whisper only")
    ap.add_argument("--compute-type", default="int8",
                    help="int8 / int8_float16 / float16 / float32 "
                         "(default: int8); faster-whisper only")
    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} via {args.engine!r} "
          f"(first run downloads it to ~/.cache)...")
    try:
        model = load_engine(args.model, engine=args.engine,
                            device=args.device,
                            compute_type=args.compute_type)
    except Exception as e:
        print(f"error: {e}", file=sys.stderr)
        sys.exit(1)

    rec = Recorder(device, args.output + ".wav")
    rec.start()
    print(f"recording to {args.output}.wav at {rec.rate} Hz "
          f"(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()