#!/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: python3 assistant.py        (--test runs the self-check)
"""

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

CONFIG_DIR = os.path.join(os.environ.get("XDG_CONFIG_HOME") or os.path.expanduser("~/.config"),
                          "desktop-assistant")
RAW = os.path.join(os.environ.get("XDG_RUNTIME_DIR", "/tmp"), "desktop-assistant.raw")
WAV = RAW[:-3] + "wav"
DEFAULTS = {
    "card": "Microphone",                 # ALSA card id of the USB mic
    "llm": "http://localhost:8181/v1",
    "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."),
    "tool_rounds": 3,
    "whisper": {"model": "/data/voice-models/whisper/ggml-large-v3-turbo-q5_0.bin",
                "device": 0, "port": 8182},
    "kokoro": {"model": "/data/voice-models/kokoro/kokoro-v1.0.onnx",
               "voices": "/data/voice-models/kokoro/voices-v1.0.bin"},
    "voice": {"it": {"lang": "it", "blend": {"if_sara": 0.8, "af_bella": 0.2}},
              "en": {"lang": "en-us", "blend": {"af_heart": 1.0}}},
    "env": {},
}
OPEN = {"env", "blend"}                   # free-form tables: replaced whole, not merged


def merge(base, over, where=""):
    # A typo must fail loudly, not fall back to the default in silence.
    out = dict(base)
    for k, v in over.items():
        if k not in base:
            raise SystemExit(f"config: unknown key {where}{k}")
        if isinstance(base[k], dict) and k not in OPEN:
            if not isinstance(v, dict):
                raise SystemExit(f"config: {where}{k} must be a table")
            v = merge(base[k], v, f"{where}{k}.")
        out[k] = v
    return out


def load_config(path):
    try:
        with open(path, "rb") as f:
            return merge(DEFAULTS, tomllib.load(f))
    except FileNotFoundError:
        return DEFAULTS


def load_tools(d):
    tools = {}
    for path in sorted(glob.glob(os.path.join(d, "*.toml"))):
        with open(path, "rb") as f:
            t = tomllib.load(f)
        if not isinstance(t.get("description"), str) or not t.get("command"):
            raise SystemExit(f"{path}: needs description and command")
        tools[os.path.basename(path)[:-5]] = t
    return tools


def schema(tools):
    return [{"type": "function", "function": {
        "name": name, "description": t["description"],
        "parameters": {"type": "object", "properties": {
            p: {"type": s.get("type", "string"), "description": s.get("description", "")}
            for p, s in t.get("params", {}).items()}}}}
        for name, t in tools.items()]


CFG = load_config(os.path.join(CONFIG_DIR, "config.toml"))
WHISPER = f"http://127.0.0.1:{CFG['whisper']['port']}"
KLANG = {l: v["lang"] for l, v in CFG["voice"].items()}
TOOLS_DIR = os.path.join(CONFIG_DIR, "tools")
TOOLS = load_tools(TOOLS_DIR)
TOOL_SCHEMA = schema(TOOLS) or None


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


def mic_on():
    out = subprocess.run(["amixer", "-c", CFG["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 run_tool(tools, name, arguments, d=TOOLS_DIR):
    # Whatever goes wrong becomes the tool's answer, so the model can say it
    # could not find out; a tool never takes the assistant down.
    t = tools.get(name)
    if t is None:
        return f"error: no tool named {name}"
    try:
        args = json.loads(arguments or "{}")
    except json.JSONDecodeError:
        args = None
    if not isinstance(args, dict):
        return "error: arguments are not valid JSON"
    vals = {p: "" if args.get(p) is None else str(args[p]) for p in t.get("params", {})}
    # No shell: each {name} fills in inside one argv element.
    argv = [re.sub(r"\{(\w+)\}", lambda m: vals.get(m.group(1), ""), a) for a in t["command"]]
    if "/" not in argv[0] and os.path.exists(os.path.join(d, argv[0])):
        argv[0] = os.path.join(d, argv[0])
    env = {**os.environ, **{k: str(v) for k, v in CFG["env"].items()}}
    try:
        # Own process group: a timeout kills the script's children too (curl),
        # which would otherwise keep the pipes open and block us.
        p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
                             env=env, start_new_session=True)
    except OSError as e:
        return f"error: cannot run {argv[0]}: {e.strerror}"
    try:
        out, err = p.communicate(timeout=t.get("timeout", 10))
    except subprocess.TimeoutExpired:
        os.killpg(p.pid, signal.SIGKILL)
        p.communicate()
        return "error: timed out"
    if p.returncode:
        last = err.strip().splitlines()[-1:]
        return ("error: " + (last[0] if last else f"exit status {p.returncode}"))[:2000]
    return out.strip()[:2000]


def collect_calls(calls, fragments):
    # A streamed tool call arrives in pieces keyed by index: id and name in
    # the first, the JSON arguments spread over the rest.
    for f in fragments:
        c = calls.setdefault(f["index"], {"id": "", "name": "", "arguments": ""})
        fn = f.get("function") or {}
        c["id"] = f.get("id") or c["id"]
        c["name"] = fn.get("name") or c["name"]
        c["arguments"] += fn.get("arguments") or ""


def announcement(tools, calls, lang):
    # Something to hear while a slow tool runs: there is no screen.
    for c in calls.values():
        phrases = tools.get(c["name"], {}).get("announce", {}).get(lang)
        if phrases:
            return random.choice(phrases)
    return None


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(CFG["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, calls, tools=None):
    # Yields reply text as it streams; tool-call fragments go into calls.
    body = {"model": llm_model(), "messages": messages, "stream": True,
            "chat_template_kwargs": {"enable_thinking": False}}
    if tools:
        body["tools"] = tools
    req = urllib.request.Request(CFG["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 not choices:
                continue
            delta = choices[0]["delta"]
            collect_calls(calls, delta.get("tool_calls") or [])
            if delta.get("content"):
                yield delta["content"]


def speaker(q):
    from kokoro_onnx import Kokoro
    k = Kokoro(CFG["kokoro"]["model"], CFG["kokoro"]["voices"])
    # Each voice is a weighted blend of Kokoro style vectors.
    voices = {l: sum(w * k.get_voice_style(n) for n, w in v["blend"].items())
              for l, v in CFG["voice"].items()}
    play = None
    while True:
        item = q.get()
        if item is None:                  # end of reply
            if play:
                play.stdin.close()
                play.wait()
                play = None
            continue
        if isinstance(item, float):       # seconds of silence
            if play:
                play.stdin.write(bytes(int(sr * item) * 2))
            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", CFG["whisper"]["model"],
                          "-dev", str(CFG["whisper"]["device"]),
                          "--host", "127.0.0.1", "--port", str(CFG["whisper"]["port"])],
                         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})
    # ponytail: last 20 messages, summarise if long sessions need memory
    # The model has no clock: without the date it makes one up for tools.
    now = time.strftime("It is %A %d %B %Y, %H:%M.")
    msgs = [{"role": "system", "content": f"{CFG['system']} {now}"}] + history[-20:]
    t, said, voice, announced = time.perf_counter(), [], lang, False

    def say(s):
        nonlocal voice
        if not said:
            log(f"first sentence after {time.perf_counter() - t:.2f}s")
        said.append(s)
        voice = lang_of(s, voice)
        q.put((re.sub(r"[*_#`]", "", s), voice))

    for rnd in range(CFG["tool_rounds"] + 1):
        calls, buf = {}, ""
        # The last round offers no tools, which forces a spoken answer.
        for delta in ask(msgs, calls, TOOL_SCHEMA if rnd < CFG["tool_rounds"] else None):
            buf += delta
            done, buf = split_sentences(buf)
            for s in done:
                say(s)
        if buf.strip():
            say(buf.strip())
        if not calls:
            break
        msgs.append({"role": "assistant", "content": "", "tool_calls": [
            {"id": c["id"], "type": "function",
             "function": {"name": c["name"], "arguments": c["arguments"]}}
            for c in calls.values()]})
        phrase = None if announced else announcement(TOOLS, calls, lang)
        if phrase:
            # The question's language: the reply has no text yet to follow.
            q.put((phrase, lang))
            # A beat after it: straight into the answer sounds rushed.
            q.put(0.6)
            announced = True
        for c in calls.values():
            result = run_tool(TOOLS, c["name"], c["arguments"])
            log(f"tool {c['name']}({c['arguments']}): {result[:120]!r}")
            msgs.append({"role": "tool", "tool_call_id": c["id"], "content": result})
    q.put(None)
    # Tool calls stay inside the turn: history keeps only what was said.
    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", CFG["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"
    c = merge(DEFAULTS, {"card": "X", "whisper": {"port": 9},
                         "voice": {"it": {"blend": {"a": 1.0}}}, "env": {"K": "v"}})
    assert c["card"] == "X" and c["whisper"] == {**DEFAULTS["whisper"], "port": 9}
    assert c["voice"]["it"] == {"lang": "it", "blend": {"a": 1.0}}      # blend replaced whole
    assert c["voice"]["en"] == DEFAULTS["voice"]["en"] and c["env"] == {"K": "v"}
    assert DEFAULTS["card"] == "Microphone"                             # defaults untouched
    for bad in ({"whisper": {"prot": 1}}, {"cards": "X"}, {"whisper": 3}):
        try:
            merge(DEFAULTS, bad)
            assert False, bad
        except SystemExit as e:
            assert "config:" in str(e)
    import tempfile
    with tempfile.TemporaryDirectory() as d:
        with open(os.path.join(d, "echo.sh"), "w") as f:
            f.write('#!/bin/sh\nprintf "%s|" "$@"\n')
        os.chmod(os.path.join(d, "echo.sh"), 0o755)
        with open(os.path.join(d, "now.toml"), "w") as f:
            f.write('description = "d"\ncommand = ["date"]\n'
                    '[params.x]\ntype = "string"\ndescription = "px"\n')
        assert load_tools(d) == {"now": {"description": "d", "command": ["date"],
                                         "params": {"x": {"type": "string", "description": "px"}}}}
        assert schema(load_tools(d)) == [{"type": "function", "function": {
            "name": "now", "description": "d", "parameters": {"type": "object", "properties": {
                "x": {"type": "string", "description": "px"}}}}}]
        tools = {"echo": {"command": ["echo.sh", "{a}", "x{b}y"], "params": {"a": {}, "b": {}}},
                 "slow": {"command": ["sh", "-c", "sleep 5; true"], "timeout": 0.3},
                 "fail": {"command": ["sh", "-c", "echo oops >&2; exit 3"]},
                 "gone": {"command": ["no-such-binary-here"]}}
        # A value is one argv element, never a shell command or a second argument.
        assert run_tool(tools, "echo", '{"a": "x; rm -rf ~", "b": 2}', d) == "x; rm -rf ~|x2y|"
        assert run_tool(tools, "echo", "{}", d) == "|xy|"
        t0 = time.time()
        assert run_tool(tools, "slow", "{}", d) == "error: timed out"
        assert time.time() - t0 < 2                 # the grandchild sleep was killed too
        assert run_tool(tools, "fail", "{}", d) == "error: oops"
        assert run_tool(tools, "gone", "{}", d).startswith("error: cannot run")
        assert run_tool(tools, "nope", "{}", d) == "error: no tool named nope"
        assert run_tool(tools, "echo", "{bad", d) == "error: arguments are not valid JSON"
        assert run_tool(tools, "echo", "[1]", d) == "error: arguments are not valid JSON"
    calls = {}
    collect_calls(calls, [{"index": 0, "id": "c1", "type": "function",
                           "function": {"name": "weather", "arguments": "{"}}])
    for piece in ['"place', '":', '"', "Rome", '"', "}"]:
        collect_calls(calls, [{"index": 0, "function": {"arguments": piece}}])
    collect_calls(calls, [{"index": 1, "id": "c2", "function": {"name": "now", "arguments": "{}"}}])
    assert calls == {0: {"id": "c1", "name": "weather", "arguments": '{"place":"Rome"}'},
                     1: {"id": "c2", "name": "now", "arguments": "{}"}}
    tools = {"now": {}, "search": {"announce": {"it": ["Cerco."], "en": ["Searching."]}}}
    assert announcement(tools, {0: {"name": "now"}}, "it") is None
    assert announcement(tools, {0: {"name": "now"}, 1: {"name": "search"}}, "en") == "Searching."
    assert announcement(tools, {0: {"name": "unknown"}}, "it") is None
    print("ok")


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