#!/usr/bin/python3
# sg-dictate -- voice typing, the native half (Windows' Win+H).
#
# The engine is Parakeet TDT 0.6B v3 on the CPU (sgspeech.py). The Windows
# side is sg-shell's sg-dictate.exe, the dictation toolbar: Wine gives it no
# AF_UNIX and no pipes to a native program, so it re-launches itself as
# `sg-dictate --bridge wine <itself> --bridged ...` and this process serves it
# over its standard handles, as sg-netctl's bridge does.
#
#   sg-dictate --transcribe-file WAV...   print what was said (the gate's mode)
#   sg-dictate --listen [--once] [--seconds N] [--partials]
#                                         dictate from the microphone to stdout
#                                         (--partials: the partial results on stderr)
#   --language en-US|de-DE|fr-FR|es-ES|auto  the text rules' language (default en-US)
#   sg-dictate --bridge PROGRAM [ARGS]    serve the dictation toolbar
#   sg-dictate --mics [--out FILE]        MIC <name>\t<description>, DEFAULT <name>
#   sg-dictate --meter FILE [--device D] [--seconds N]
#                                         the microphone's level, for "Test
#                                         microphone": LEVEL <0-100> in FILE,
#                                         until N seconds pass or FILE.stop exists
#   sg-dictate --status                   MODEL installed|missing|downloading|failed
#   sg-dictate --download                 fetch the speech model (asks sg-speechd)
#   sg-dictate --remove                   remove it again (administrators)
#   sg-dictate --serve                    sg-speechd: root, one per connection
#
# Privacy: audio is held in memory only while it is recognised, nothing heard
# or typed is logged, and nothing but the explicit model download touches the
# network.
#
# Copyright (C) 2026 Stained Glass OS contributors
# SPDX-License-Identifier: AGPL-3.0-or-later

import json
import os
import queue
import select
import signal
import socket
import subprocess
import sys
import threading
import time

sys.path.insert(0, os.environ.get("SG_SPEECH_LIB", "/usr/lib/stained-glass/speech"))
import sgspeech  # noqa: E402

SOCKET = os.environ.get("SG_SPEECHD_SOCKET", "/run/stained-glass-speech/speechd.sock")
STATUS = os.path.join(sgspeech.MODEL_DIR, "status")
ADMIN_GROUP = "sg-admins"
USER_GROUP = os.environ.get("SG_WINE_GROUP", "sgwine")
CHUNK = 512 * 2  # bytes: 32 ms of 16 kHz s16le mono, one VAD frame


def log(msg):
    """Diagnostics only -- never what was heard."""
    print("sg-dictate: " + msg, file=sys.stderr, flush=True)


# ---- the microphone -----------------------------------------------------------------

def list_mics():
    """[(name, description)], default name. Monitors of outputs are not
    microphones and are left out."""
    try:
        out = subprocess.run(["pactl", "--format=json", "list", "sources"],
                             capture_output=True, text=True, timeout=10).stdout
        sources = json.loads(out or "[]")
    except (OSError, ValueError, subprocess.TimeoutExpired):
        return [], ""
    mics = []
    for s in sources:
        props = s.get("properties", {})
        if props.get("device.class") == "monitor" or s.get("monitor_source") or \
           s.get("name", "").endswith(".monitor"):
            continue
        mics.append((s.get("name", ""), s.get("description") or s.get("name", "")))
    try:
        default = subprocess.run(["pactl", "get-default-source"], capture_output=True,
                                 text=True, timeout=10).stdout.strip()
    except (OSError, subprocess.TimeoutExpired):
        default = ""
    return mics, default


class Mic:
    """16 kHz mono s16le from PipeWire (through its PulseAudio server, which
    every desktop has), on a thread, to `sink(samples)`.

    SG_DICTATE_AUDIO_FILE stands a WAV in for the microphone, played at real
    time and then silence, for the gates."""

    def __init__(self, sink, device=""):
        self.sink = sink
        self.device = device
        self.proc = None
        self.stopped = threading.Event()
        self.error = None
        self.thread = threading.Thread(target=self._run, daemon=True)

    def start(self):
        fake = os.environ.get("SG_DICTATE_AUDIO_FILE")
        if not fake:
            dev = [f"--device={self.device}"] if self.device else []
            cmds = [["parec", "--raw", "--format=s16le", "--rate=16000", "--channels=1",
                     "--latency-msec=30", "--client-name=Voice typing"] + dev,
                    ["pw-record", "--raw", "--rate", "16000", "--channels", "1", "--format", "s16",
                     "--latency", "30ms"] + (["--target", self.device] if self.device else []) + ["-"]]
            for cmd in cmds:
                try:
                    self.proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                                 stderr=subprocess.DEVNULL, stdin=subprocess.DEVNULL)
                    break
                except OSError:
                    continue
            if not self.proc:
                raise OSError("no audio capture program (parec or pw-record)")
        self.thread.start()

    def _run(self):
        import numpy as np
        fake = os.environ.get("SG_DICTATE_AUDIO_FILE")
        if fake:
            audio = sgspeech.read_wav(fake)
            n = 512
            t0 = time.monotonic()
            i = 0
            while not self.stopped.is_set():
                chunk = audio[i:i + n] if i < len(audio) else np.zeros(n, np.float32)
                self.sink(chunk)
                i += n
                delay = t0 + i / sgspeech.SAMPLE_RATE - time.monotonic()
                if delay > 0:
                    time.sleep(delay)
            return
        pending = b""
        fd = self.proc.stdout.fileno()
        while not self.stopped.is_set():
            try:
                data = os.read(fd, 4096)
            except OSError:
                data = b""
            if not data:
                if not self.stopped.is_set():
                    self.error = "the microphone stopped"
                    self.sink(None)
                return
            pending += data
            usable = len(pending) - len(pending) % 2
            if usable:
                self.sink(np.frombuffer(pending[:usable], "<i2").astype(np.float32) / 32768)
                pending = pending[usable:]

    def stop(self):
        self.stopped.set()
        if self.proc:
            try:
                self.proc.terminate()
                self.proc.wait(timeout=2)
            except (OSError, subprocess.TimeoutExpired):
                self.proc.kill()
            self.proc = None


# ---- dictation: one engine behind the bridge or --listen ---------------------------------

class Dictation:
    """Listens, cuts speech into utterances, recognises them in order, and
    hands each one's text to `emit`. Loads the model on a thread at once, so
    audio captured while it loads is not lost.

    While someone speaks, the utterance so far is recognised again every
    PARTIAL_EVERY seconds of new audio (or less often, when that takes long)
    and sent as PARTIAL: the toolbar shows it, and never types it. The
    utterance's final TEXT replaces it."""

    NO_SPEECH_TIMEOUT = {False: 8.0, True: 120.0}  # seconds; continuous or not
    PARTIAL_EVERY = 0.5     # seconds of new speech between partial results
    PARTIAL_FIRST = 0.4     # speech heard before the first one

    def __init__(self, emit):
        self.emit = emit
        self.lock = threading.Lock()
        self.model_ready = threading.Event()
        self.rec = self.vad = self.seg = None
        self.load_error = None
        self.mic = None
        self.opts = {}
        self.backlog = []       # audio heard before the model was ready
        self.utterances = queue.Queue()
        self.tail = ""          # the last character typed, to join the next text
        self.listening = False
        self.heard_at = 0.0
        self.level_at = 0.0
        self.levels = []
        self.loading = False
        self.serial = 0         # utterances finished so far: a partial of an older one is stale
        self.partial_job = None  # (serial, audio, opts): the newest utterance-so-far to show
        self.partial_at = 0     # samples of the utterance when the last partial was asked for
        self.partial_cost = 0.0  # seconds the last partial took, to keep partials cheap
        self.partial_shown = False
        self.last_tail = ""     # the tail before the last text typed, for "delete that"
        threading.Thread(target=self._recognise, daemon=True).start()

    def preload(self):
        """Load the model now, on a thread, if it is not loaded or loading."""
        with self.lock:
            if self.loading or self.model_ready.is_set():
                return
            self.loading = True
        threading.Thread(target=self._load, daemon=True).start()

    def unload(self):
        """Give the memory back (the hold-to-talk listener, long unused)."""
        with self.lock:
            if self.listening or self.loading or not self.model_ready.is_set():
                return
            self.rec = self.vad = self.seg = None
            self.load_error = None
            self.model_ready.clear()
        import gc
        gc.collect()
        log("model unloaded")

    def _load(self):
        if not sgspeech.model_installed() and not os.environ.get("SG_SPEECH_MODEL"):
            with self.lock:
                self.load_error = "nomodel"
                self.loading = False
                self.model_ready.set()
                if self.listening:
                    self._stop_locked()
                    self.emit("STATE nomodel")
            return
        try:
            t = time.monotonic()
            d = os.environ.get("SG_SPEECH_MODEL") or sgspeech.model_path()
            self.rec = sgspeech.Recognizer(d)
            self.vad = sgspeech.SileroVAD(os.path.join(d, "silero_vad.onnx"))
            self.seg = sgspeech.Segmenter(self.vad)
            log("model loaded in %.1fs" % (time.monotonic() - t))
        except Exception as e:  # noqa: BLE001 -- any failure is "cannot dictate"
            self.load_error = "error " + str(e).splitlines()[0][:200]
        with self.lock:
            self.loading = False
            self.model_ready.set()
            if self.listening:
                if self.load_error:
                    self._stop_locked()
                    self.emit("STATE " + self.load_error)
                else:
                    backlog, self.backlog = self.backlog, []
                    for a in backlog:
                        self._segment(a)
                    self.emit("STATE listening")

    def start(self, opts):
        with self.lock:
            if self.load_error and not self.loading:
                # Try again: the model may have been installed since.
                self.load_error = None
                self.model_ready.clear()
        self.preload()
        with self.lock:
            if self.listening:
                return
            self.opts = opts
            if opts.get("fresh"):
                self.tail = ""
            if isinstance(opts.get("tail"), str):
                # What is before the caret, when the toolbar could tell.
                self.tail = opts["tail"][-1:]
            if self.model_ready.is_set() and self.load_error:
                self.emit("STATE " + self.load_error)
                return
            self.mic = Mic(self._audio, opts.get("mic", ""))
            try:
                self.mic.start()
            except OSError as e:
                self.mic = None
                self.emit("STATE nomic " + str(e))
                return
            self.listening = True
            self.heard_at = time.monotonic()
            if self.seg:
                self.seg.reset()
            self.emit("STATE listening" if self.model_ready.is_set() else "STATE loading")

    def _audio(self, samples):
        if samples is None:  # the capture program died
            with self.lock:
                if self.listening:
                    err = self.mic.error if self.mic else "the microphone stopped"
                    self._stop_locked()
                    self.emit("STATE nomic " + err)
            return
        now = time.monotonic()
        with self.lock:
            if not self.listening:
                return
            self.levels.append(sgspeech.level(samples))
            if now - self.level_at >= 0.1:
                self.emit("LEVEL %d" % max(self.levels))
                self.levels, self.level_at = [], now
            if not self.model_ready.is_set():
                self.backlog.append(samples)
                del self.backlog[:-2000]  # a minute, whatever happens
                return
            self._segment(samples)
            continuous = bool(self.opts.get("continuous", True))
            if self.seg.in_speech:
                self.heard_at = now
            elif now - self.heard_at > self.NO_SPEECH_TIMEOUT[continuous]:
                self._finish_locked("timeout")

    def _segment(self, samples):
        for utt in self.seg.push(samples):
            self._utterance_done(utt)
            self.heard_at = time.monotonic()
            if not self.opts.get("continuous", True):
                # One utterance, as Windows' voice typing without continuous
                # dictation: stop listening once the speaker pauses.
                self._finish_locked(None, flush=False)
                return
        self._maybe_partial()

    def _utterance_done(self, audio):
        self.serial += 1
        self.partial_job = None
        self.partial_at = 0
        self.utterances.put((audio, dict(self.opts)))

    def _maybe_partial(self):
        """Ask for a partial result of the utterance so far, when enough new
        speech has been heard since the last one. Bounded: the interval grows
        to twice what recognising took, so a slow machine is never behind."""
        if not self.opts.get("partials", True) or not self.seg or not self.seg.in_speech:
            return
        n = sum(len(f) for f in self.seg.speech)
        every = max(self.PARTIAL_EVERY, 2 * self.partial_cost) * sgspeech.SAMPLE_RATE
        first = self.PARTIAL_FIRST * sgspeech.SAMPLE_RATE
        if n < first or n - self.partial_at < (every if self.partial_at else first):
            return
        self.partial_at = n
        import numpy as np
        pending = self.partial_job is not None
        self.partial_job = (self.serial, np.concatenate(self.seg.speech), dict(self.opts))
        if not pending:
            self.utterances.put(("partial", None))

    def _finish_locked(self, why, flush=True):
        tail = self.seg.flush() if (flush and self.seg) else None
        if self.seg and not flush:
            self.seg.reset()
        if tail is not None:
            self._utterance_done(tail)
        self._stop_locked()
        self.utterances.put(("idle", why))

    def _stop_locked(self):
        self.listening = False
        self.backlog = []
        self.serial += 1        # whatever partial is in flight is stale now
        self.partial_job = None
        self.partial_at = 0
        if self.mic:
            self.mic.stop()
            self.mic = None

    def stop(self, cancel=False):
        """Stop listening. What was said up to now is still typed, unless
        this is a cancel."""
        with self.lock:
            if not self.listening:
                return
            if cancel:
                if self.seg:
                    self.seg.reset()
                self._stop_locked()
                self._clear_partial_locked()
                self.emit("STATE idle")
                return
            if not self.model_ready.is_set():
                # Stopped before the model was ready: recognise the backlog.
                backlog, self.backlog = self.backlog, []
                self._stop_locked()
                self.utterances.put(("backlog", (backlog, dict(self.opts))))
                self.utterances.put(("idle", None))
                return
            self._finish_locked(None)

    def _recognise(self):
        import numpy as np
        while True:
            item, extra = self.utterances.get()
            if isinstance(item, str) and item == "idle":
                with self.lock:
                    self._clear_partial_locked()
                self.emit("STATE idle" + (" " + extra if extra else ""))
                continue
            if isinstance(item, str) and item == "partial":
                self._partial()
                continue
            if isinstance(item, str) and item == "backlog":
                backlog, opts = extra
                self.model_ready.wait()
                if self.load_error or not backlog:
                    continue
                audio = np.concatenate(backlog)
                seg = sgspeech.Segmenter(sgspeech.SileroVAD(os.path.join(
                    os.environ.get("SG_SPEECH_MODEL") or sgspeech.model_path(), "silero_vad.onnx")))
                pieces = seg.push(audio)
                last = seg.flush()
                for p in pieces + ([last] if last is not None else []):
                    self._type(p, opts)
                continue
            self._type(item, extra)

    def _clear_partial_locked(self):
        if self.partial_shown:
            self.partial_shown = False
            self.emit("PARTIAL \"\"")

    def _postprocess(self, raw, opts):
        return sgspeech.postprocess(
            raw, spoken_punctuation=opts.get("spoken", True),
            auto_punctuation=opts.get("auto", True),
            remove_fillers=opts.get("fillers", True), numbers=opts.get("numbers", True),
            language=opts.get("language", "en-US"))

    def _partial(self):
        """Recognise the utterance so far and show it, unless it has been
        finished (or listening stopped) meanwhile: then the final text, or
        nothing, is what counts."""
        with self.lock:
            job, self.partial_job = self.partial_job, None
        if not job or not self.rec:
            return
        serial, audio, opts = job
        t = time.monotonic()
        try:
            raw = self.rec.transcribe(audio)
        except Exception as e:  # noqa: BLE001 -- a partial is only a preview
            log("partial recognition failed: %s" % e)
            return
        took = time.monotonic() - t
        text = self._postprocess(raw, opts) if raw else ""
        if os.environ.get("SG_DICTATE_TIMING"):
            log("partial: %.2fs of audio in %.2fs" % (len(audio) / sgspeech.SAMPLE_RATE, took))
        with self.lock:
            self.partial_cost = took
            if serial != self.serial or not text:
                return
            self.partial_shown = True
            self.emit("PARTIAL " + json.dumps(text, ensure_ascii=False))

    def _type(self, audio, opts):
        t = time.monotonic()
        try:
            raw = self.rec.transcribe(audio)
        except Exception as e:  # noqa: BLE001 -- say so, and keep serving
            log("recognition failed: %s" % e)
            self.emit("STATE error recognition failed")
            return
        took = time.monotonic() - t
        if os.environ.get("SG_DICTATE_TIMING"):
            log("recognised %.2fs of audio in %.2fs" % (len(audio) / sgspeech.SAMPLE_RATE, took))
        # "commands": the toolbar's own switch (always on since sg-shell
        # 0.1.0-92, where spoken punctuation is off by default); an older
        # toolbar tied them to spoken punctuation
        commands = opts.get("commands", opts.get("spoken", True))
        what = sgspeech.command(raw, opts.get("language", "en-US")) if commands else None
        if what:
            # "Delete that", "Das löschen", "Efface ça", "Borra eso"...: a
            # command, never text. The toolbar does delete/undo; stop is ours.
            with self.lock:
                self._clear_partial_locked()
                if what == "delete":
                    self.tail = self.last_tail
                if what == "stop":
                    if self.listening:
                        if self.seg:
                            self.seg.reset()
                        self._stop_locked()
                        self.utterances.put(("idle", "stopped"))
                    return
            self.emit("CMD " + what)
            return
        text = self._postprocess(raw, opts)
        with self.lock:
            if not text:
                self._clear_partial_locked()
                return
            self.last_tail = self.tail
            text = sgspeech.join(self.tail, text)
            self.tail = text[-1:]
            # the final replaces the partial: the toolbar clears it on TEXT
            self.partial_shown = False
        self.emit("TEXT " + json.dumps(text, ensure_ascii=False))

    def quit(self):
        with self.lock:
            self._stop_locked()


# ---- --bridge: the toolbar's engine ----------------------------------------------------

def bridge(program):
    """Run the dictation toolbar with pipes as its standard handles. It sends
    one JSON object a line ({"cmd": "start", ...}, "stop", "cancel", "quit");
    this answers with STATE/LEVEL/PARTIAL/TEXT/CMD lines whenever there is news. Ends
    when the program does."""
    down_r, down_w = os.pipe()
    up_r, up_w = os.pipe()
    child = subprocess.Popen(program, stdin=down_r, stdout=up_w, close_fds=True)
    os.close(down_r)
    os.close(up_w)
    to_child = os.fdopen(down_w, "w", encoding="utf-8", newline="\n")
    out_lock = threading.Lock()

    def emit(line):
        with out_lock:
            try:
                to_child.write(line + "\n")
                to_child.flush()
            except (BrokenPipeError, ValueError):
                pass

    engine = Dictation(emit)
    pending = b""
    while True:
        ready, _, _ = select.select([up_r], [], [], 1.0)
        if not ready:
            if child.poll() is not None:
                break
            continue
        chunk = os.read(up_r, 65536)
        if not chunk:
            break
        pending += chunk
        if len(pending) > 65536 and b"\n" not in pending:
            pending = b""
        while b"\n" in pending:
            raw, pending = pending.split(b"\n", 1)
            line = raw.decode("utf-8", "replace").strip("\r")
            if not line.startswith("{"):
                continue  # Wine may write other things to stdout
            try:
                req = json.loads(line)
                cmd = req.get("cmd")
            except (ValueError, AttributeError):
                continue
            if cmd == "start":
                engine.start(req)
            elif cmd == "stop":
                engine.stop()
            elif cmd == "cancel":
                engine.stop(cancel=True)
            elif cmd == "preload":
                engine.preload()
            elif cmd == "unload":
                engine.unload()
            elif cmd == "ping":
                emit("PONG")
            elif cmd == "quit":
                engine.quit()
                break
    engine.quit()
    try:
        return child.wait(timeout=10)
    except subprocess.TimeoutExpired:
        return 0


# ---- the other modes -------------------------------------------------------------------

def transcribe_files(paths, raw, post):
    d = os.environ.get("SG_SPEECH_MODEL") or sgspeech.model_path()
    if not os.path.exists(os.path.join(d, "encoder-model.int8.onnx")):
        log("the speech model is not installed (sg-dictate --download)")
        return 3
    t = time.monotonic()
    rec = sgspeech.Recognizer(d)
    log("model loaded in %.1fs" % (time.monotonic() - t))
    for p in paths:
        audio = sgspeech.read_wav(p)
        t = time.monotonic()
        text = rec.transcribe(audio)
        took = time.monotonic() - t
        if not raw:
            what = sgspeech.command(text, post.get("language", "en")) if post["spoken_punctuation"] else None
            text = "<command %s>" % what if what else sgspeech.postprocess(text, **post)
        if os.environ.get("SG_DICTATE_TIMING"):
            log("%s: %.2fs of audio in %.2fs" % (p, len(audio) / sgspeech.SAMPLE_RATE, took))
        print(json.dumps(text, ensure_ascii=False) if len(paths) == 1 else
              p + "\t" + json.dumps(text, ensure_ascii=False), flush=True)
    return 0


def listen(post, continuous=True, seconds=None, show_partials=False):
    """Dictate from the microphone to stdout until Ctrl+C, `seconds` pass,
    or (not continuous) the speaker pauses. `show_partials`: the partial
    results too, on stderr ("partial: <JSON>")."""
    done = threading.Event()

    def emit(line):
        if line.startswith("TEXT "):
            sys.stdout.write(json.loads(line[5:]))
            sys.stdout.flush()
        elif line.startswith("PARTIAL ") and show_partials:
            # What the toolbar would show while the speaker is still talking.
            print("partial: " + line[8:], file=sys.stderr, flush=True)
        elif line.startswith("CMD "):
            log("command " + line[4:])
        elif line.startswith("STATE "):
            log(line[6:])
            if line.startswith("STATE idle") or line.startswith("STATE no") or \
               line.startswith("STATE error"):
                done.set()

    engine = Dictation(emit)
    opts = dict(post)
    opts["continuous"] = continuous
    engine.start(opts)
    try:
        if not done.wait(seconds):
            engine.stop()
            done.wait(60)
    except KeyboardInterrupt:
        engine.stop()
        done.wait(30)
    print()
    return 0


def meter(path, device, seconds):
    """For "Test microphone": the level, rewritten every 100 ms, until
    `seconds` pass or PATH.stop appears. Only a number is written."""
    state = {"levels": [], "at": 0.0}

    def write(text):
        tmp = path + ".tmp"
        with open(tmp, "w") as f:
            f.write(text)
        os.replace(tmp, path)

    def sink(samples):
        if samples is None:
            return
        state["levels"].append(sgspeech.level(samples))
        now = time.monotonic()
        if now - state["at"] >= 0.1:
            write("LEVEL %d\n" % max(state["levels"]))
            state["levels"], state["at"] = [], now

    mic = Mic(sink, device)
    try:
        mic.start()
    except OSError as e:
        write("ERROR %s\n" % e)
        return 1
    end = time.monotonic() + seconds
    while time.monotonic() < end and not os.path.exists(path + ".stop") and mic.error is None:
        time.sleep(0.1)
    mic.stop()
    write("DONE\n" if mic.error is None else "ERROR %s\n" % mic.error)
    try:
        os.unlink(path + ".stop")
    except OSError:
        pass
    return 0


def read_status():
    """MODEL installed | missing | downloading DONE TOTAL | failed MESSAGE."""
    if sgspeech.model_installed():
        return "MODEL installed"
    try:
        with open(STATUS) as f:
            fields = dict(line.rstrip("\n").split(" ", 1) for line in f if " " in line)
    except OSError:
        fields = {}
    state = fields.get("STATE", "missing")
    if state == "downloading":
        return "MODEL downloading %s %s" % (fields.get("DONE", "0"), fields.get("TOTAL", sgspeech.TOTAL_BYTES))
    if state == "failed":
        return "MODEL failed " + fields.get("MESSAGE", "")
    return "MODEL missing"


# ---- the model download (sg-speechd, root) -----------------------------------------------

def write_status(**fields):
    tmp = STATUS + ".tmp"
    with open(tmp, "w") as f:
        for k, v in fields.items():
            f.write("%s %s\n" % (k.upper(), v))
    os.chmod(tmp, 0o644)
    os.replace(tmp, STATUS)


def fetch_model(progress, skip_if_packaged=False):
    """Download what sgspeech.FILES lists into the model directory, resuming
    anything partly there, checking each file's SHA-256, and mark the model
    verified. `progress(done, total)` hears about it."""
    import fcntl
    import hashlib
    import urllib.request

    if skip_if_packaged and sgspeech.model_packaged():
        # Installed as a package already: nothing to download. (sg-image's
        # build-deb.sh calls this without it to fill its cache: that must
        # really fetch, whatever the build machine has installed.)
        progress(sgspeech.TOTAL_BYTES, sgspeech.TOTAL_BYTES)
        return
    target = sgspeech.download_path()
    os.makedirs(target, mode=0o755, exist_ok=True)
    lock = open(os.path.join(sgspeech.MODEL_DIR, ".lock"), "w")
    fcntl.flock(lock, fcntl.LOCK_EX)  # a second request waits for the first
    if sgspeech.model_downloaded():
        progress(sgspeech.TOTAL_BYTES, sgspeech.TOTAL_BYTES)
        return
    mirror = os.environ.get("SG_SPEECH_MIRROR") if os.getuid() != 0 else None
    total = sgspeech.TOTAL_BYTES
    done_before = 0
    last = [0.0]

    def report(done, force=False):
        now = time.monotonic()
        if force or now - last[0] >= 0.25:
            last[0] = now
            write_status(state="downloading", done=done, total=total)
            progress(done, total)

    try:
        for name, url, size, sha in sgspeech.FILES:
            final = os.path.join(target, name)
            part = final + ".part"
            if os.path.exists(final) and os.path.getsize(final) == size:
                done_before += size
                report(done_before, True)
                continue
            for attempt in range(6):
                have = os.path.getsize(part) if os.path.exists(part) else 0
                if have > size:
                    os.unlink(part)
                    have = 0
                if have < size:
                    src = ("file://" + os.path.join(mirror, name)) if mirror else url
                    req = urllib.request.Request(src, headers={"User-Agent": "sg-speechd"})
                    if have and not mirror:
                        req.add_header("Range", "bytes=%d-" % have)
                    try:
                        with urllib.request.urlopen(req, timeout=30) as r, open(part, "ab") as f:
                            if have and getattr(r, "status", 200) != 206:
                                f.truncate(0)  # the server ignored the range
                                f.seek(0)
                                have = 0
                            if mirror and have:
                                r.read(have)
                            while True:
                                block = r.read(1 << 16)
                                if not block:
                                    break
                                f.write(block)
                                have += len(block)
                                report(done_before + have)
                    except OSError as e:
                        log("download of %s interrupted (%s); retrying" % (name, e))
                        time.sleep(min(30, 2 ** attempt))
                        continue
                if os.path.getsize(part) == size:
                    break
            else:
                raise OSError("could not download " + name)
            h = hashlib.sha256()
            with open(part, "rb") as f:
                for block in iter(lambda: f.read(1 << 20), b""):
                    h.update(block)
            if h.hexdigest() != sha:
                os.unlink(part)
                raise OSError("%s did not match its checksum" % name)
            os.chmod(part, 0o644)
            os.replace(part, final)
            done_before += size
            report(done_before, True)
        with open(os.path.join(target, "NOTICE"), "w") as f:
            f.write("Speech recognition: NVIDIA Parakeet TDT 0.6B v3 (int8 ONNX export by\n"
                    "istupakov), licensed CC BY 4.0: https://creativecommons.org/licenses/by/4.0/\n"
                    "https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3\n"
                    "Voice activity detection: Silero VAD v5, MIT licence.\n"
                    "https://github.com/snakers4/silero-vad\n")
        with open(os.path.join(target, ".verified"), "w") as f:
            for name, _, _, sha in sgspeech.FILES:
                f.write("%s  %s\n" % (sha, name))
        write_status(state="installed", done=total, total=total)
    except OSError as e:
        write_status(state="failed", message=str(e))
        raise


PACKAGED_MESSAGE = ("the speech model is installed as the sg-speech-model-parakeet package; "
                    "remove that package to remove it")


def remove_model():
    """Remove sg-speechd's download. The packaged model is dpkg's: callers
    refuse that before getting here."""
    import shutil
    target = sgspeech.download_path()
    try:
        os.unlink(os.path.join(target, ".verified"))
    except OSError:
        pass
    shutil.rmtree(target, ignore_errors=True)
    write_status(state="missing")


def serve():
    """sg-speechd: one request on stdin (a socket, SO_PEERCRED), as root.
    Anyone with a Windows session may install the speech model, as anyone may
    add a speech pack on Windows; removing it is an administrator's."""
    import grp
    import pwd
    import struct
    signal.signal(signal.SIGPIPE, signal.SIG_IGN)
    sock = socket.socket(fileno=0)
    creds = sock.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i"))
    _, uid, gid = struct.unpack("3i", creds)
    out = os.fdopen(1, "w", encoding="utf-8", newline="\n")

    def say(line):
        try:
            out.write(line + "\n")
            out.flush()
        except OSError:
            pass  # the asker went away; the download carries on

    def groups():
        try:
            name = pwd.getpwuid(uid).pw_name
            return {grp.getgrgid(g).gr_name for g in os.getgrouplist(name, gid)}
        except (KeyError, OSError):
            return set()

    line = sys.stdin.buffer.readline(4096)
    try:
        argv = json.loads(line)["argv"]
    except (ValueError, KeyError, TypeError):
        say("ERROR invalid malformed request")
        return 2
    g = groups()
    admin = uid == 0 or ADMIN_GROUP in g
    if argv == ["status"]:
        say(read_status())
        say("OK")
        return 0
    if argv == ["download"]:
        if not (admin or USER_GROUP in g):
            say("ERROR denied only users of this computer may install the speech model")
            return 3
        try:
            fetch_model(lambda d, t: say("PROGRESS %d %d" % (d, t)), skip_if_packaged=True)
        except OSError as e:
            say("ERROR failed " + str(e))
            return 1
        say("OK")
        return 0
    if argv == ["remove"]:
        if not admin:
            say("ERROR denied you need to be an administrator to remove the speech model")
            return 3
        if sgspeech.model_packaged():
            say("ERROR unsupported " + PACKAGED_MESSAGE)
            return 4
        remove_model()
        say("OK")
        return 0
    say("ERROR invalid unknown request")
    return 2


def ask(argv, show):
    """A request to sg-speechd; `show` gets each line before the last."""
    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        s.connect(SOCKET)
    except OSError as e:
        log("the speech service is not available (%s): %s" % (SOCKET, e))
        return 1
    s.sendall((json.dumps({"argv": argv}) + "\n").encode())
    f = s.makefile("r", encoding="utf-8")
    for line in f:
        line = line.rstrip("\n")
        if line == "OK":
            return 0
        if line.startswith("ERROR "):
            log(line.split(" ", 2)[-1])
            return 3 if line.startswith("ERROR denied") else 4 if line.startswith("ERROR unsupported") else 1
        show(line)
    log("the speech service went away")
    return 1


def local_or_ask(argv):
    """Where the model directory is ours (a developer's, a gate's), do it
    here; otherwise ask the service."""
    def show(line):
        print(line, flush=True)
    if argv == ["remove"] and sgspeech.model_packaged():
        log(PACKAGED_MESSAGE)
        return 4
    if os.access(sgspeech.MODEL_DIR, os.W_OK) or (
            not os.path.exists(sgspeech.MODEL_DIR) and os.environ.get("SG_SPEECH_DIR")):
        os.makedirs(sgspeech.MODEL_DIR, exist_ok=True)
        try:
            if argv == ["download"]:
                fetch_model(lambda d, t: show("PROGRESS %d %d" % (d, t)), skip_if_packaged=True)
            else:
                remove_model()
        except OSError as e:
            log(str(e))
            return 1
        return 0
    return ask(argv, show)


def main():
    args = sys.argv[1:]
    language = "en-US"
    if "--language" in args:
        i = args.index("--language")
        language = args[i + 1] if i + 1 < len(args) else "en-US"
        args = args[:i] + args[i + 2:]
    post = {"spoken_punctuation": "--no-spoken-punctuation" not in args,
            "auto_punctuation": "--no-auto-punctuation" not in args,
            "remove_fillers": "--keep-fillers" not in args,
            "numbers": "--no-numbers" not in args,
            "language": language}
    flags = [a for a in args if a.startswith("--no-") or a in ("--keep-fillers", "--partials")]
    rest = [a for a in args if a not in flags]
    mode = rest[0] if rest else "--help"
    if mode == "--bridge" and len(rest) > 1:
        return bridge(rest[1:])
    if mode == "--transcribe-file" and len(rest) > 1:
        raw = "--raw" in rest
        return transcribe_files([p for p in rest[1:] if p != "--raw"], raw, post)
    if mode == "--listen":
        seconds = float(rest[rest.index("--seconds") + 1]) if "--seconds" in rest else None
        return listen({"spoken": post["spoken_punctuation"], "auto": post["auto_punctuation"],
                       "fillers": post["remove_fillers"], "numbers": post["numbers"],
                       "language": language, "partials": "--no-partials" not in args},
                      continuous="--once" not in rest, seconds=seconds,
                      show_partials="--partials" in args)
    if mode == "--mics":
        # --out FILE: for a Windows program, which cannot read our stdout.
        mics, default = list_mics()
        text = "".join("MIC %s\t%s\n" % m for m in mics) + ("DEFAULT %s\n" % default if default else "")
        if "--out" in rest:
            path = rest[rest.index("--out") + 1]
            with open(path + ".tmp", "w", encoding="utf-8") as f:
                f.write(text + "END\n")
            os.replace(path + ".tmp", path)
        else:
            sys.stdout.write(text)
        return 0
    if mode == "--meter" and len(rest) > 1:
        device, seconds = "", 15.0
        if "--device" in rest:
            device = rest[rest.index("--device") + 1]
        if "--seconds" in rest:
            seconds = float(rest[rest.index("--seconds") + 1])
        return meter(rest[1], device, seconds)
    if mode == "--status":
        print(read_status())
        print("SIZE %d" % sgspeech.TOTAL_BYTES)
        if sgspeech.model_installed():
            # packaged: dpkg's (sg-speech-model-parakeet); downloaded: sg-speechd's
            print("SOURCE %s %s" % (("packaged", sgspeech.packaged_path()) if sgspeech.model_packaged()
                                    else ("downloaded", sgspeech.download_path())))
        return 0
    if mode == "--download":
        return local_or_ask(["download"])
    if mode == "--remove":
        return local_or_ask(["remove"])
    if mode == "--serve":
        return serve()
    print(__doc__ or "usage: see the header of " + sys.argv[0], file=sys.stderr)
    return 2


if __name__ == "__main__":
    sys.exit(main())
