# /// script
# requires-python = ">=3.11,<3.13"
# dependencies = ["mlx-whisper", "sounddevice", "pynput", "numpy"]
# ///
"""wersper — hold Right Option, speak, release. Text appears at your cursor.

Run:  uv run wersper.py
Esc while recording cancels the take.
"""
import subprocess
import sys
import threading
import time

import numpy as np
import sounddevice as sd
import mlx_whisper
from pynput import keyboard

MODEL = "mlx-community/whisper-large-v3-turbo"
LANGUAGE = "en"  # None = autodetect (slower, occasionally wrong)
HOTKEY = keyboard.Key.alt_r
SAMPLE_RATE = 16000
# ponytail: crude silence gate; raise if background noise gets transcribed,
# lower if it eats quiet speech
MIN_RMS = 0.001

frames = []
stream = None
kb = keyboard.Controller()


def sound(name):
    subprocess.Popen(["afplay", f"/System/Library/Sounds/{name}.aiff"])


def transcribe(audio):
    result = mlx_whisper.transcribe(audio, path_or_hf_repo=MODEL, language=LANGUAGE)
    return result["text"].strip()


def paste(text):
    old = subprocess.run("pbpaste", capture_output=True).stdout
    subprocess.run("pbcopy", input=text.encode(), check=True)
    with kb.pressed(keyboard.Key.cmd):
        kb.press("v")
        kb.release("v")
    time.sleep(0.4)  # let the paste land before giving the clipboard back
    subprocess.run("pbcopy", input=old)  # ponytail: restores text only, not images


def start_recording():
    global stream
    if stream:
        return
    frames.clear()
    stream = sd.InputStream(
        samplerate=SAMPLE_RATE, channels=1, dtype="float32",
        callback=lambda data, *_: frames.append(data.copy()),
    )
    stream.start()
    sound("Pop")
    print("● recording... (Esc to cancel)", flush=True)


def stop_recording(cancel=False):
    global stream
    if not stream:
        return
    stream.stop()
    stream.close()
    stream = None
    if cancel:
        print("✕ cancelled", flush=True)
        return
    if not frames:
        return
    audio = np.concatenate(frames).flatten()
    if len(audio) < SAMPLE_RATE // 3:  # ignore accidental taps
        return
    if np.sqrt((audio ** 2).mean()) < MIN_RMS:
        print("(silence, skipped)", flush=True)
        return
    # ponytail: transcribe in a thread so the hotkey stays responsive
    threading.Thread(target=lambda: finish(audio), daemon=True).start()


def finish(audio):
    t0 = time.time()
    text = transcribe(audio)
    if text:
        print(f"→ {text}  ({time.time() - t0:.1f}s)", flush=True)
        paste(text)
        sound("Tink")


def on_press(key):
    if key == HOTKEY:
        start_recording()
    elif key == keyboard.Key.esc and stream:
        stop_recording(cancel=True)


def on_release(key):
    if key == HOTKEY:
        stop_recording()


def check():
    """Self-check: model loads, pipeline works, silence gate works."""
    assert isinstance(transcribe(np.zeros(SAMPLE_RATE, dtype=np.float32)), str)
    assert np.sqrt((np.zeros(SAMPLE_RATE) ** 2).mean()) < MIN_RMS
    print("ok: model loaded, pipeline works")


if __name__ == "__main__":
    if "--check" in sys.argv:
        check()
        sys.exit(0)
    print("loading model (first run downloads ~1.6GB)...", flush=True)
    transcribe(np.zeros(SAMPLE_RATE, dtype=np.float32))  # warm up
    print("ready — hold Right ⌥ to dictate, Esc cancels, Ctrl+C quits", flush=True)
    with keyboard.Listener(on_press=on_press, on_release=on_release) as listener:
        listener.join()