aman/src/leld.py

#!/usr/bin/env python3
import argparse
import json
import logging
import os
import signal
import sys
import threading
import time
from pathlib import Path

from config import Config, load, redacted_dict
from recorder import start_recording, stop_recording
from stt import FasterWhisperSTT, STTConfig
from aiprocess import AIConfig, build_processor
from inject import inject
from x11_hotkey import listen
from tray import run_tray


class State:
    IDLE = "idle"
    RECORDING = "recording"
    TRANSCRIBING = "transcribing"
    PROCESSING = "processing"
    OUTPUTTING = "outputting"


class Daemon:
    def __init__(self, cfg: Config):
        self.cfg = cfg
        self.lock = threading.Lock()
        self.state = State.IDLE
        self.proc = None
        self.record = None
        self.timer = None
        self.stt = FasterWhisperSTT(
            STTConfig(
                model=cfg.whisper_model,
                language=cfg.whisper_lang,
                device=cfg.whisper_device,
                vad_filter=True,
            )
        )
        self.ai = None
        if cfg.ai_enabled:
            self.ai = build_processor(
                AIConfig(
                    provider=cfg.ai_provider,
                    model=cfg.ai_model,
                    temperature=cfg.ai_temperature,
                    system_prompt_file=cfg.ai_system_prompt_file,
                    base_url=cfg.ai_base_url,
                    api_key=cfg.ai_api_key,
                    timeout_sec=cfg.ai_timeout_sec,
                )
            )

    def set_state(self, state: str):
        with self.lock:
            self.state = state

    def get_state(self):
        with self.lock:
            return self.state

    def toggle(self):
        with self.lock:
            if self.state == State.IDLE:
                self._start_recording_locked()
                return
            if self.state == State.RECORDING:
                self.state = State.TRANSCRIBING
                threading.Thread(target=self._stop_and_process, daemon=True).start()
                return
            logging.info("busy (%s), trigger ignored", self.state)

    def _start_recording_locked(self):
        try:
            proc, record = start_recording(self.cfg.ffmpeg_input, self.cfg.ffmpeg_path)
        except Exception as exc:
            logging.error("record start failed: %s", exc)
            return
        self.proc = proc
        self.record = record
        self.state = State.RECORDING
        logging.info("recording started (%s)", record.wav_path)
        if self.timer:
            self.timer.cancel()
        self.timer = threading.Timer(self.cfg.record_timeout_sec, self._timeout_stop)
        self.timer.daemon = True
        self.timer.start()

    def _timeout_stop(self):
        with self.lock:
            if self.state != State.RECORDING:
                return
            self.state = State.TRANSCRIBING
        threading.Thread(target=self._stop_and_process, daemon=True).start()

    def _stop_and_process(self):
        proc = self.proc
        record = self.record
        self.proc = None
        self.record = None
        if self.timer:
            self.timer.cancel()
            self.timer = None

        if not proc or not record:
            self.set_state(State.IDLE)
            return

        logging.info("stopping recording (user)")
        try:
            stop_recording(proc)
        except Exception as exc:
            logging.error("record stop failed: %s", exc)
            self.set_state(State.IDLE)
            return

        if not Path(record.wav_path).exists():
            logging.error("no audio captured")
            self.set_state(State.IDLE)
            return

        try:
            self.set_state(State.TRANSCRIBING)
            text = self.stt.transcribe(record.wav_path)
        except Exception as exc:
            logging.error("whisper failed: %s", exc)
            self.set_state(State.IDLE)
            return

        logging.info("transcript: %s", text)

        if self.ai:
            self.set_state(State.PROCESSING)
            try:
                text = self.ai.process(text) or text
            except Exception as exc:
                logging.error("ai process failed: %s", exc)

        logging.info("output: %s", text)

        try:
            self.set_state(State.OUTPUTTING)
            inject(text, self.cfg.injection_backend)
        except Exception as exc:
            logging.error("output failed: %s", exc)
        finally:
            self.set_state(State.IDLE)

    def stop_recording(self):
        with self.lock:
            if self.state != State.RECORDING:
                return
            self.state = State.TRANSCRIBING
        threading.Thread(target=self._stop_and_process, daemon=True).start()


def _lock_single_instance():
    runtime_dir = Path(os.getenv("XDG_RUNTIME_DIR", "/tmp")) / "lel"
    runtime_dir.mkdir(parents=True, exist_ok=True)
    lock_path = runtime_dir / "lel.lock"
    f = open(lock_path, "w")
    try:
        import fcntl

        fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except Exception:
        raise SystemExit("another instance is running")
    return f


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="", help="path to config.json")
    parser.add_argument("--no-tray", action="store_true", help="disable tray icon")
    parser.add_argument("--dry-run", action="store_true", help="log hotkey only")
    args = parser.parse_args()

    logging.basicConfig(stream=sys.stderr, level=logging.INFO, format="leld: %(asctime)s %(message)s")
    cfg = load(args.config)

    _lock_single_instance()

    logging.info("ready (hotkey: %s)", cfg.hotkey)
    logging.info("config (%s):\n%s", args.config or str(Path.home() / ".config" / "lel" / "config.json"), json.dumps(redacted_dict(cfg), indent=2))

    daemon = Daemon(cfg)

    def on_quit():
        os._exit(0)

    def handle_signal(_sig, _frame):
        logging.info("signal received, shutting down")
        daemon.stop_recording()
        end = time.time() + 5
        while time.time() < end and daemon.get_state() != State.IDLE:
            time.sleep(0.1)
        os._exit(0)

    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    if args.no_tray:
        listen(cfg.hotkey, lambda: logging.info("hotkey pressed (dry-run)") if args.dry_run else daemon.toggle())
        return

    threading.Thread(target=lambda: listen(cfg.hotkey, lambda: logging.info("hotkey pressed (dry-run)") if args.dry_run else daemon.toggle()), daemon=True).start()
    run_tray(daemon.get_state, on_quit)


if __name__ == "__main__":
    main()