#!/usr/bin/env python3
# desktop-assistant: push-to-talk local voice assistant.
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License version 2 as published
# by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
"""Mic button on: record. Mic button off: transcribe (whisper-server),
ask the loaded llama-server model, speak the reply (Kokoro).

Run: .venv/bin/python assistant.py        (--test runs the self-check)
"""

import ctypes
import json
import os
import queue
import re
import signal
import subprocess
import sys
import threading
import time
import unicodedata
import urllib.request
import wave

CARD = "Microphone"                       # ALSA card id of the USB mic
MODELS = "/data/voice-models"
WHISPER = "http://127.0.0.1:8182"
LLM = "http://localhost:8181/v1"
RAW = os.path.join(os.environ.get("XDG_RUNTIME_DIR", "/tmp"), "desktop-assistant.raw")
WAV = RAW[:-3] + "wav"
SYSTEM = ("You are a voice assistant. Your replies are spoken aloud, so answer "
          "briefly in plain sentences: no markdown, no lists, no emoji. "
          "Reply in the language the user speaks, Italian or English.")
KLANG = {"it": "it", "en": "en-us"}


def log(*a):
    print(time.strftime("%H:%M:%S"), *a, flush=True)


def mic_on():
    out = subprocess.run(["amixer", "-c", CARD, "cget", "numid=2"],
                         capture_output=True, text=True).stdout
    return out.rstrip().endswith("values=on")


def presenting():
    # During calls the user turns on presentation mode; the button then
    # belongs to the call, not to us.
    r = subprocess.run([os.path.expanduser("~/bin/statusctl"), "presentation", "get"],
                       capture_output=True, text=True)
    return r.stdout.strip() == "1"


def latin(text):
    return all(unicodedata.name(c, "").startswith("LATIN") for c in text if c.isalpha())


def transcribe(lang="auto"):
    # Always send language: whisper-server's own default is en, not auto.
    out = subprocess.run(["curl", "-s", "-F", f"file=@{WAV}", "-F", f"language={lang}",
                          "-F", "response_format=verbose_json", WHISPER + "/inference"],
                         capture_output=True, text=True).stdout
    d = json.loads(out)
    return d.get("language_probabilities", {}), " ".join(d["text"].split())


def stt():
    # large-v3-turbo's detector is reliable for this speaker (small's was
    # not). Only it and en are spoken: if detection lands elsewhere, or the
    # text is non-Latin (a hallucination, both are Latin script), redo the
    # pass forced to the likelier of the two.
    p, text = transcribe()
    lang = max(KLANG, key=lambda l: p.get(l, 0))
    if max(p, key=p.get, default=lang) != lang or not latin(text):
        text = transcribe(lang)[1]
    return lang, text


IT_WORDS = set("il lo la gli le di che è non per una sono ho hai con mi ti ci del della "
               "come questo questa anche ma più se nel alla puoi posso sei".split())
EN_WORDS = set("the is are you i to of and it that for with not this what be have "
               "can your my do there".split())


def lang_of(sentence, fallback):
    # The voice must follow the reply's language, not the user's: the model
    # may answer in the other one. ponytail: stopword count, a real language
    # identifier if this misfires.
    words = re.findall(r"[a-zàèéìòù]+", sentence.lower())
    it = sum(w in IT_WORDS for w in words)
    en = sum(w in EN_WORDS for w in words)
    return "it" if it > en else "en" if en > it else fallback


def split_sentences(buf):
    parts = re.split(r"(?<=[.!?…])\s+|\n+", buf)
    return [p.strip() for p in parts[:-1] if p.strip()], parts[-1]


def llm_model():
    # The router keeps one model resident; asking for another would swap it
    # out under other clients, so use whatever is not unloaded.
    with urllib.request.urlopen(LLM + "/models") as r:
        models = json.load(r)["data"]
    return next((m["id"] for m in models if m["status"]["value"] != "unloaded"),
                models[0]["id"])


def ask(messages):
    body = {"model": llm_model(), "messages": messages, "stream": True,
            "chat_template_kwargs": {"enable_thinking": False}}
    req = urllib.request.Request(LLM + "/chat/completions", json.dumps(body).encode(),
                                 {"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as r:
        for line in r:
            if not line.startswith(b"data: ") or line.strip() == b"data: [DONE]":
                continue
            choices = json.loads(line[6:])["choices"]
            if choices and choices[0]["delta"].get("content"):
                yield choices[0]["delta"]["content"]


def speaker(q):
    from kokoro_onnx import Kokoro, EspeakConfig
    # The espeakng-loader wheel's bundled lib has a broken data path; use
    # the system espeak-ng.
    k = Kokoro(f"{MODELS}/kokoro/kokoro-v1.0.onnx", f"{MODELS}/kokoro/voices-v1.0.bin",
               espeak_config=EspeakConfig(lib_path="/usr/lib64/libespeak-ng.so.1",
                                          data_path="/usr/share/espeak-ng-data"))
    voices = {"it": 0.8 * k.get_voice_style("if_sara") + 0.2 * k.get_voice_style("af_bella"),
              "en": k.get_voice_style("af_heart")}
    play = None
    while True:
        item = q.get()
        if item is None:                  # end of reply
            if play:
                play.stdin.close()
                play.wait()
                play = None
            continue
        text, lang = item
        audio, sr = k.create(text, voice=voices[lang], lang=KLANG[lang])
        if play is None:
            play = subprocess.Popen(["aplay", "-q", "-r", str(sr), "-f", "S16_LE", "-c", "1"],
                                    stdin=subprocess.PIPE)
            # A new stream wakes the suspended sink and the first ~0.2 s is
            # lost: lead with silence so it eats that, not the first word.
            play.stdin.write(bytes(int(sr * 0.3) * 2))
        # The pipe buffer lets synthesis of the next sentence overlap playback.
        play.stdin.write((audio.clip(-1, 1) * 32767).astype("<i2").tobytes())


def die_with_parent():
    # Runs in the child before exec: the kernel sends SIGTERM when we die,
    # however we die, so no whisper-server or amixer is ever orphaned.
    ctypes.CDLL("libc.so.6").prctl(1, signal.SIGTERM)   # PR_SET_PDEATHSIG


def start_whisper():
    def up():
        try:
            urllib.request.urlopen(WHISPER, timeout=1)
            return True
        except OSError:
            return False
    if up():
        return
    p = subprocess.Popen(["whisper-server", "-m", f"{MODELS}/whisper/ggml-large-v3-turbo-q5_0.bin",
                          "-dev", "0", "--host", "127.0.0.1", "--port", WHISPER.rsplit(":", 1)[1]],
                         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                         preexec_fn=die_with_parent)
    for _ in range(60):
        if up():
            return
        time.sleep(0.5)
    p.kill()
    sys.exit("whisper-server did not come up")


def reply(history, q):
    pcm = open(RAW, "rb").read()
    if len(pcm) < 16000 * 2 * 0.3:               # under 0.3 s: a stray press
        return
    # Trailing silence: whisper drops the last word when audio ends mid-word,
    # which it does when the button is pressed while still talking.
    with wave.open(WAV, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(16000)
        w.writeframes(pcm + bytes(16000 * 2 // 2))
    t = time.perf_counter()
    lang, text = stt()
    if not text or text.startswith("["):         # [BLANK_AUDIO] and friends
        return
    log(f"you ({lang}, {time.perf_counter() - t:.2f}s): {text}")
    history.append({"role": "user", "content": text})
    t, said, buf = time.perf_counter(), [], ""
    for delta in ask([{"role": "system", "content": SYSTEM}] + history[-20:]):
        # ponytail: last 20 messages, summarise if long sessions need memory
        buf += delta
        done, buf = split_sentences(buf)
        if not said and done:
            log(f"first sentence after {time.perf_counter() - t:.2f}s")
        for s in done:
            said.append(s)
            lang = lang_of(s, lang)
            q.put((re.sub(r"[*_#`]", "", s), lang))
    if buf.strip():
        if not said:
            log(f"first sentence after {time.perf_counter() - t:.2f}s")
        said.append(buf.strip())
        q.put((re.sub(r"[*_#`]", "", buf.strip()), lang_of(buf, lang)))
    q.put(None)
    history.append({"role": "assistant", "content": " ".join(said)})
    log("assistant:", " ".join(said))


def main():
    # SIGTERM must run the cleanup below too, or the recording is left behind.
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
    start_whisper()
    q = queue.Queue()
    threading.Thread(target=speaker, args=(q,), daemon=True).start()
    history, rec = [], None
    events = subprocess.Popen(["stdbuf", "-oL", "amixer", "-c", CARD, "events"],
                              stdout=subprocess.PIPE, text=True, preexec_fn=die_with_parent)
    prev = mic_on()
    log(f"ready, mic is {'on' if prev else 'off'}")
    try:
        for line in events.stdout:
            if "Mic Capture Switch" not in line:
                continue
            on = mic_on()
            if on == prev:
                continue
            prev = on
            if on and presenting():
                log("presentation mode: ignored")
            elif on:
                # ponytail: no barge-in, a new press does not cut the reply
                rec = subprocess.Popen(["arecord", "-q", "-t", "raw", "-f", "S16_LE", "-r", "16000",
                                        "-c", "1", RAW], stderr=subprocess.DEVNULL)
            elif rec:
                rec.send_signal(signal.SIGINT)
                rec.wait()
                rec = None
                reply(history, q)
    except KeyboardInterrupt:
        pass
    finally:
        for f in (RAW, WAV):                    # the user's voice: never keep it
            if os.path.exists(f):
                os.remove(f)


def test():
    assert split_sentences("Ciao! Come stai? Io") == (["Ciao!", "Come stai?"], "Io")
    assert split_sentences("Pi is 3.14 today") == ([], "Pi is 3.14 today")
    assert split_sentences("one\n\ntwo") == (["one"], "two")
    assert latin("Perché no? Hi!") and not latin("早, 早!")
    assert lang_of("Non posso fornire il meteo in tempo reale.", "en") == "it"
    assert lang_of("That is a very exciting place to be.", "it") == "en"
    assert lang_of("Tokyo!", "it") == "it"
    print("ok")


if __name__ == "__main__":
    test() if sys.argv[1:] == ["--test"] else main()