#!/usr/bin/env python3

"""Dictation: hold a key, speak, and the words are typed where you were typing.

    panama-dictate start       begin recording (bound to key press)
    panama-dictate stop        stop, transcribe, and type the result
    panama-dictate cancel      stop and discard
    panama-dictate status      JSON: model, server, and image
    panama-dictate setup       fetch the image and the model, with progress

There is no daemon here, and that is the point. Transcription is done by
whisper.cpp running in a container -- see config/containers/whisper -- which
holds the model on the GPU between utterances. This script only records, asks
that server for the text, and types it, so `start` and `stop` are two short
runs of a script rather than two messages to something long-lived. What they
share is a pidfile.

Neither the image nor the model ships with Panama, and neither is fetched during
installation. Together they are more than two gigabytes, and an installer that
quietly spends ten minutes on them is the surprise the interview exists to
prevent. Settings offers both as one action; until then, dictation says what is
missing rather than hanging.

The image is upstream's own Vulkan build of whisper.cpp -- Vulkan so that one
image serves the AMD, Intel and NVIDIA machines this config runs on, where a
ROCm build would serve only the AMD ones at twenty times the size.
"""

from __future__ import annotations

import json
import os
import re
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
from pathlib import Path

MODEL = "small"
MODEL_DIR = Path(os.environ.get("XDG_DATA_HOME", Path.home() / ".local/share")) / "panama/whisper"
MODEL_FILE = MODEL_DIR / f"ggml-{MODEL}.bin"
MODEL_URL = f"https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-{MODEL}.bin"

IMAGE = "ghcr.io/ggml-org/whisper.cpp:main-vulkan"
SERVICE = "panama-whisper.service"
ENDPOINT = "http://127.0.0.1:8791/inference"

RUNTIME_DIR = Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}"))
STATE_FILE = RUNTIME_DIR / "panama-dictate.json"

# 16 kHz mono is what Whisper resamples to anyway, so recording it directly
# avoids a conversion and keeps the file small.
SAMPLE_RATE = "16000"

# Anything shorter is a key tapped by accident rather than something said.
# Transcribing a fifth of a second of room tone reliably invents a word, and
# that word then gets typed.
MIN_SECONDS = 0.35

REASONS = {
    "not-downloaded": "Speech model not downloaded — see Settings",
    "no-image": "Speech server not installed — see Settings",
    "server-unavailable": "The speech server did not answer",
    "transcribe-failed": "Transcription failed",
    "too-short": "Too short",
    "no-speech": "Nothing was said",
    "not-recording": "Not recording",
    "already-recording": "Already listening",
    "deliver-failed": "Could not deliver the text",
}


def osd(message: str | None) -> None:
    """Show a line on the on-screen display, or clear it.

    Never fatal: the shell can be restarting, and dictation that refused to work
    because it could not draw a label would be worse than a silent one.
    """
    argv = (["qs", "ipc", "call", "osd", "hide"] if message is None
            else ["qs", "ipc", "call", "osd", "message", "microphone", message])
    try:
        subprocess.run(argv, capture_output=True, timeout=3, check=False)
    except (OSError, subprocess.SubprocessError):
        pass


def emit(result: dict) -> int:
    print(json.dumps(result))
    return 0 if result.get("ok") else 2


def have_image() -> bool:
    try:
        done = subprocess.run(["podman", "image", "exists", IMAGE],
                              capture_output=True, timeout=15, check=False)
        return done.returncode == 0
    except (OSError, subprocess.SubprocessError):
        return False


def server_ready(timeout: float = 1.5) -> bool:
    try:
        request = urllib.request.Request(ENDPOINT, method="OPTIONS")
        urllib.request.urlopen(request, timeout=timeout)
        return True
    except urllib.error.HTTPError:
        # Answering at all is what is being asked. A 404 or 405 to OPTIONS still
        # means something is listening and loaded.
        return True
    except (urllib.error.URLError, OSError):
        return False


def start_server() -> bool:
    """Bring the container up and wait for the model to load.

    Started here rather than at login: whisper-server holds the model from the
    moment it starts, and a session where nobody dictates should not be paying
    half a gigabyte of VRAM for the possibility.
    """
    if server_ready():
        return True
    try:
        subprocess.run(["systemctl", "--user", "start", SERVICE],
                       capture_output=True, timeout=30, check=False)
    except (OSError, subprocess.SubprocessError):
        return False

    # The first start loads the model, which takes seconds rather than
    # milliseconds. Recording is already running by then, so this wait costs
    # nothing that was going to be spoken.
    deadline = time.monotonic() + 120
    while time.monotonic() < deadline:
        if server_ready():
            return True
        time.sleep(0.25)
    return False


def deliver(text: str) -> str:
    """Get text into the focused window. Returns how it got there.

    Typing is preferred: it lands wherever the cursor is, works in terminals and
    editors that ignore a paste, and leaves the clipboard alone -- which matters
    here, because Panama keeps clipboard history and dictation would otherwise
    fill it with everything ever said.

    The clipboard is the fallback rather than the default, for the applications
    that refuse synthetic keystrokes. Which one happened is reported back, so
    the caller can say so: a fallback nobody is told about is how "it typed
    nothing" becomes a mystery.
    """
    if shutil.which("wtype"):
        try:
            done = subprocess.run(["wtype", "--", text],
                                  capture_output=True, timeout=60, check=False)
            if done.returncode == 0:
                return "typed"
        except (OSError, subprocess.SubprocessError):
            pass

    if shutil.which("wl-copy"):
        try:
            subprocess.run(["wl-copy", "--", text],
                           capture_output=True, timeout=10, check=False)
            return "copied"
        except (OSError, subprocess.SubprocessError):
            pass

    return "failed"


def transcribe(path: Path) -> tuple[str | None, str | None]:
    """Ask the container for the text. Returns (text, error)."""
    boundary = "----panama-dictate"
    body = b"".join([
        f'--{boundary}\r\nContent-Disposition: form-data; name="file"; '
        f'filename="speech.wav"\r\nContent-Type: audio/wav\r\n\r\n'.encode(),
        path.read_bytes(),
        f'\r\n--{boundary}\r\nContent-Disposition: form-data; name="response_format"'
        f'\r\n\r\ntext\r\n--{boundary}--\r\n'.encode(),
    ])
    request = urllib.request.Request(
        ENDPOINT, data=body, method="POST",
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"})
    try:
        with urllib.request.urlopen(request, timeout=300) as response:
            answer = response.read().decode(errors="replace").strip()
    except (urllib.error.URLError, OSError, urllib.error.HTTPError):
        return None, "transcribe-failed"

    # response_format=text is asked for, but a server built from a different
    # upstream revision may answer JSON regardless. Reading both costs three
    # lines and saves a version bump turning dictation into gibberish.
    if answer.startswith("{"):
        try:
            answer = str(json.loads(answer).get("text", "")).strip()
        except (json.JSONDecodeError, AttributeError):
            return None, "transcribe-failed"

    # One line, always. The server returns a line per segment, and a newline
    # typed into a window is an Enter press -- which sends the half-finished
    # message, submits the form, or runs the command. Whitespace is collapsed
    # rather than only stripped, because the break is in the middle of the
    # sentence, not at its ends.
    return " ".join(answer.split()), None


# Whisper describes what it heard when it did not hear words: [BLANK_AUDIO] for
# silence, and things like (wind blowing) or [MUSIC] for sound that is not
# speech. They arrive as ordinary text, so without this they get typed -- the
# first test of this dictated the literal string "[BLANK_AUDIO]" into the
# clipboard, which is precisely the kind of thing nobody notices until it lands
# in the middle of a commit message.
NON_SPEECH = re.compile(r"[\[(][^\])]*[\])]")


def is_speech(text: str) -> bool:
    """False when the transcription is only Whisper describing the absence of speech.

    Deliberately conservative: bracketed groups are removed and what remains is
    what decides. A sentence that merely contains an aside keeps its brackets and
    is still typed in full -- only a transcription that is nothing BUT markers is
    discarded, because throwing away real words would be far worse than typing a
    stray marker.
    """
    return NON_SPEECH.sub("", text).strip() != ""


def read_state() -> dict | None:
    try:
        return json.loads(STATE_FILE.read_text())
    except (OSError, json.JSONDecodeError):
        return None


def handle_start() -> dict:
    if not MODEL_FILE.exists():
        osd(REASONS["not-downloaded"])
        return {"ok": False, "error": "not-downloaded"}
    if not have_image():
        osd(REASONS["no-image"])
        return {"ok": False, "error": "no-image"}

    existing = read_state()
    if existing is not None:
        # A stale pidfile from a crash should not block dictation forever, so
        # it is only respected while the process it names is actually alive.
        try:
            os.kill(existing["pid"], 0)
            return {"ok": False, "error": "already-recording"}
        except (OSError, KeyError, TypeError):
            STATE_FILE.unlink(missing_ok=True)

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    handle, name = tempfile.mkstemp(prefix="panama-dictate-", suffix=".wav")
    os.close(handle)

    # pw-record rather than a library: pipewire-utils is already installed --
    # it backs the privacy indicators and the sound test -- and it records from
    # whatever WirePlumber calls the default source, so changing input device
    # in Settings changes what dictation hears, with no code here.
    recorder = subprocess.Popen(
        ["pw-record", "--rate", SAMPLE_RATE, "--channels", "1", "--format", "s16", name],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)

    STATE_FILE.write_text(json.dumps({
        "pid": recorder.pid, "path": name, "startedAt": time.time()}))
    STATE_FILE.chmod(0o600)
    osd("Listening…")
    return {"ok": True, "recording": True}


def finish_recording() -> tuple[Path | None, float]:
    """Stop the recorder and hand back its file and how long it ran."""
    state = read_state()
    STATE_FILE.unlink(missing_ok=True)
    if state is None:
        return None, 0.0

    path = Path(state.get("path", ""))
    elapsed = max(0.0, time.time() - float(state.get("startedAt", time.time())))
    try:
        # SIGINT rather than SIGKILL: pw-record finalises the WAV header when
        # interrupted, and a file whose header still claims zero length is one
        # the server reads as silence.
        os.kill(int(state["pid"]), signal.SIGINT)
    except (OSError, KeyError, TypeError, ValueError):
        return (path if path.exists() else None), elapsed

    deadline = time.monotonic() + 5
    while time.monotonic() < deadline:
        try:
            os.kill(int(state["pid"]), 0)
        except OSError:
            break
        time.sleep(0.05)
    return (path if path.exists() else None), elapsed


def handle_stop() -> dict:
    path, elapsed = finish_recording()
    if path is None:
        osd(None)
        return {"ok": False, "error": "not-recording"}

    try:
        if elapsed < MIN_SECONDS:
            osd(REASONS["too-short"])
            return {"ok": False, "error": "too-short"}

        osd("Transcribing…")
        if not start_server():
            osd(REASONS["server-unavailable"])
            return {"ok": False, "error": "server-unavailable"}

        text, error = transcribe(path)
        if error is not None:
            osd(REASONS[error])
            return {"ok": False, "error": error}
        if not text or not is_speech(text):
            osd(REASONS["no-speech"])
            return {"ok": False, "error": "no-speech"}

        how = deliver(text)
        if how == "failed":
            osd(REASONS["deliver-failed"])
            return {"ok": False, "error": "deliver-failed", "text": text}

        # Announced only when it did NOT go where it was meant to. A successful
        # dictation needs no notification: the words are already on screen.
        osd("Pasted — press Ctrl+V" if how == "copied" else None)
        return {"ok": True, "delivered": how, "characters": len(text)}
    finally:
        path.unlink(missing_ok=True)


def handle_cancel() -> dict:
    path, _ = finish_recording()
    if path is not None:
        path.unlink(missing_ok=True)
    osd(None)
    return {"ok": True, "cancelled": True}


def handle_status() -> dict:
    state = read_state()
    return {
        "ok": True,
        "model": MODEL,
        "modelInstalled": MODEL_FILE.exists(),
        "modelPath": str(MODEL_FILE),
        "modelBytes": MODEL_FILE.stat().st_size if MODEL_FILE.exists() else 0,
        "imageBuilt": have_image(),
        "serverReady": server_ready(),
        "recording": state is not None,
        "typingAvailable": shutil.which("wtype") is not None,
    }


def pull_image() -> bool:
    """Fetch the speech server image. Progress is podman's own, on stderr."""
    if have_image():
        return True
    print(json.dumps({"ok": True, "state": "pulling", "image": IMAGE}), flush=True)
    try:
        done = subprocess.run(["podman", "pull", IMAGE],
                              capture_output=True, timeout=3600, check=False)
        return done.returncode == 0
    except (OSError, subprocess.SubprocessError):
        return False


def download_model() -> bool:
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    if MODEL_FILE.exists():
        return True

    # Written beside the target and moved into place, so an interrupted download
    # never leaves a file that looks like a model. Half a model is worse than
    # none: everything downstream believes it is there and fails at load.
    partial = MODEL_FILE.with_suffix(".partial")
    print(json.dumps({"ok": True, "state": "downloading", "model": MODEL}), flush=True)
    try:
        with urllib.request.urlopen(MODEL_URL, timeout=60) as response:
            total = int(response.headers.get("content-length", 0))
            done = 0
            last = 0.0
            with partial.open("wb") as handle:
                while True:
                    chunk = response.read(1 << 20)
                    if not chunk:
                        break
                    handle.write(chunk)
                    done += len(chunk)
                    now = time.monotonic()
                    # Once a second, not once a megabyte: five hundred progress
                    # lines is not progress, it is a flood.
                    if total and now - last >= 1.0:
                        last = now
                        print(json.dumps({"ok": True, "state": "downloading",
                                          "bytes": done, "total": total}), flush=True)
    except (urllib.error.URLError, OSError):
        partial.unlink(missing_ok=True)
        return False

    partial.replace(MODEL_FILE)
    return True


def handle_setup() -> int:
    """Everything dictation needs that Panama does not install.

    The image first: a model is useless without something to run it, and being
    told to fetch half a gigabyte before hearing that an image is also needed
    would be two trips.
    """
    if not shutil.which("podman"):
        print(json.dumps({"ok": False, "error": "no-podman"}), flush=True)
        return 2
    if not pull_image():
        print(json.dumps({"ok": False, "error": "pull-failed"}), flush=True)
        return 2
    if not download_model():
        print(json.dumps({"ok": False, "error": "download-failed"}), flush=True)
        return 2

    print(json.dumps({"ok": True, "state": "ready",
                      "bytes": MODEL_FILE.stat().st_size}), flush=True)
    return 0


HANDLERS = {
    "start": handle_start,
    "stop": handle_stop,
    "cancel": handle_cancel,
    "status": handle_status,
}


def main(argv: list[str]) -> int:
    command = argv[0] if argv else "status"
    if command == "setup":
        return handle_setup()
    handler = HANDLERS.get(command)
    if handler is None:
        return emit({"ok": False, "error": "unknown-command"})
    return emit(handler())


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))