blurt/blurt.py

import base64
import io
import logging
import os
import shlex
import signal
import subprocess
import sys
import tempfile
import threading
from dataclasses import dataclass
from logging.handlers import RotatingFileHandler
from pathlib import Path
from typing import Callable

import numpy as np
import keyring
import pyperclip
import pystray
import requests
import sounddevice as sd
from PIL import Image, ImageDraw
from scipy.io import wavfile

OPENAI_TRANSCRIPTIONS_URL = "https://api.openai.com/v1/audio/transcriptions"
OPENROUTER_TRANSCRIPTIONS_URL = "https://openrouter.ai/api/v1/audio/transcriptions"
DEFAULT_TRANSCRIPTION_MODEL = "whisper-1"
DEFAULT_TRANSCRIPTION_LANGUAGE = "en"
# (model_id, menu_label, provider). The provider selects both the API endpoint
# and the request format: "openai" speaks multipart file upload; "openrouter"
# speaks JSON with base64-encoded audio. OpenRouter's whisper-large-v3-turbo is
# served by Groq (216x real-time, ~12% WER) and reuses your OpenRouter key.
TRANSCRIPTION_MODELS = [
    ("whisper-1", "OpenAI: whisper-1", "openai"),
    ("gpt-4o-mini-transcribe", "OpenAI: gpt-4o-mini-transcribe", "openai"),
    ("gpt-4o-transcribe", "OpenAI: gpt-4o-transcribe", "openai"),
    ("openai/whisper-large-v3-turbo", "OpenRouter: whisper-large-v3-turbo (Groq, fast)", "openrouter"),
]
PROVIDER_KEY_FILES = {"openai": "api_key.txt", "openrouter": "openrouter_key.txt"}
TRANSCRIPTION_LANGUAGES = [
    ("en", "English"),
    ("de", "German"),
]
SAMPLE_RATE = 16000
CHANNELS = 1
DTYPE = "int16"
APP_NAME = "Blurt"
APP_DIR = Path(__file__).resolve().parent
AGENT_PREFIX = "to my agent"
PID_FILE = Path(tempfile.gettempdir()) / "blurt.pid"
LOG_FILE = Path(tempfile.gettempdir()) / "blurt.log"
KEYRING_CREDENTIALS = [
    ("openai-api-key", "felixm"),
    ("blurt", "openai-api-key"),
]
OPENROUTER_KEYRING_CREDENTIALS = [
    ("openrouter-api-key", "felixm"),
]


def setup_logging():
    logger = logging.getLogger()
    if logger.handlers:
        return

    logger.setLevel(logging.INFO)
    handler = RotatingFileHandler(
        LOG_FILE,
        maxBytes=512 * 1024,
        backupCount=3,
        encoding="utf-8",
    )
    handler.setFormatter(
        logging.Formatter("%(asctime)s %(levelname)s [%(threadName)s] %(message)s")
    )
    logger.addHandler(handler)

    def log_thread_exception(args):
        logging.critical(
            "Uncaught thread exception in %s",
            args.thread.name if args.thread else "<unknown>",
            exc_info=(args.exc_type, args.exc_value, args.exc_traceback),
        )

    threading.excepthook = log_thread_exception


def load_config_value(env_name: str, file_name: str, default: str = "") -> str:
    env_value = os.environ.get(env_name, "").strip()
    if env_value:
        logging.info("Loaded %s from environment", env_name)
        return env_value

    config_file = APP_DIR / file_name
    try:
        value = config_file.read_text().strip()
        logging.info(
            "Loaded %s from %s: %s",
            env_name,
            config_file,
            "present" if value else "empty",
        )
        return value
    except FileNotFoundError:
        if default:
            logging.info("Config file not found, using default for %s: %s", env_name, default)
        else:
            logging.warning("Config file not found for %s: %s", env_name, config_file)
        return default


def _load_key_from_keyring(credentials: list[tuple[str, str]], label: str) -> str:
    for service_name, username in credentials:
        try:
            stored_key = keyring.get_password(service_name, username)
        except Exception:
            logging.exception(
                "Could not read %s key from keyring entry %s/%s",
                label,
                service_name,
                username,
            )
            continue
        if stored_key:
            logging.info("Loaded %s key from keyring entry %s/%s", label, service_name, username)
            return stored_key.strip()
    return ""


def load_api_key() -> str:
    env_key = os.environ.get("OPENAI_API_KEY", "").strip()
    if env_key:
        logging.info("Loaded OPENAI_API_KEY from environment")
        return env_key

    key = _load_key_from_keyring(KEYRING_CREDENTIALS, "OpenAI")
    if key:
        return key

    return load_config_value("OPENAI_API_KEY", "api_key.txt")


def load_openrouter_key() -> str:
    env_key = os.environ.get("OPENROUTER_API_KEY", "").strip()
    if env_key:
        logging.info("Loaded OPENROUTER_API_KEY from environment")
        return env_key

    key = _load_key_from_keyring(OPENROUTER_KEYRING_CREDENTIALS, "OpenRouter")
    if key:
        return key

    return load_config_value("OPENROUTER_API_KEY", "openrouter_key.txt")


def load_transcription_model() -> str:
    model = load_config_value(
        "BLURT_TRANSCRIPTION_MODEL",
        "model.txt",
        DEFAULT_TRANSCRIPTION_MODEL,
    )
    valid_models = {model_id for model_id, _, _ in TRANSCRIPTION_MODELS}
    if model not in valid_models:
        logging.warning("Unknown transcription model %s, using %s", model, DEFAULT_TRANSCRIPTION_MODEL)
        return DEFAULT_TRANSCRIPTION_MODEL
    return model


def save_transcription_model(model: str):
    model_file = APP_DIR / "model.txt"
    model_file.write_text(f"{model}\n")
    logging.info("Saved transcription model to %s: %s", model_file, model)


def load_transcription_language() -> str:
    language = load_config_value(
        "BLURT_TRANSCRIPTION_LANGUAGE",
        "language.txt",
        DEFAULT_TRANSCRIPTION_LANGUAGE,
    )
    valid_languages = {language_id for language_id, _ in TRANSCRIPTION_LANGUAGES}
    if language not in valid_languages:
        logging.warning(
            "Unknown transcription language %s, using %s",
            language,
            DEFAULT_TRANSCRIPTION_LANGUAGE,
        )
        return DEFAULT_TRANSCRIPTION_LANGUAGE
    return language


def save_transcription_language(language: str):
    language_file = APP_DIR / "language.txt"
    language_file.write_text(f"{language}\n")
    logging.info("Saved transcription language to %s: %s", language_file, language)


def get_model_label(model: str) -> str:
    for model_id, label, _ in TRANSCRIPTION_MODELS:
        if model_id == model:
            return label
    return model


def get_language_label(language: str) -> str:
    for language_id, label in TRANSCRIPTION_LANGUAGES:
        if language_id == language:
            return label
    return language


@dataclass
class Platform:
    notify: Callable[[str, str], None]
    dispatch: Callable[[str], None]
    open_logs: Callable[[], None]
    setup_toggle_trigger: Callable[["BlurtApp"], None]
    cleanup_toggle_trigger: Callable[[], None]


def _make_windows_platform() -> Platform:
    from pynput import keyboard
    from winotify import Notification

    listener = None

    def notify(title: str, message: str):
        try:
            toast = Notification(app_id=APP_NAME, title=title, msg=message)
            toast.show()
        except Exception:
            logging.exception("Windows notification failed")

    def dispatch(task: str):
        wezterm = r"C:\Program Files\WezTerm\wezterm-gui.exe"
        dispatch_sh = "/mnt/c/Users/felix.martin/AppData/Roaming/FlowLauncher/Plugins/Dispatch-1.0.0/dispatch.sh"
        safe_query = shlex.quote(task)
        cmd = [
            wezterm,
            "-e", "wsl.exe",
            "-d", "Ubuntu",
            "-e", "zsh",
            "-c", f"{dispatch_sh} {safe_query}",
        ]
        logging.info("Dispatching agent task with %d characters", len(task))
        subprocess.Popen(cmd, creationflags=subprocess.DETACHED_PROCESS)

    def open_logs():
        subprocess.Popen(["notepad.exe", str(LOG_FILE)])

    def setup_toggle_trigger(app: "BlurtApp"):
        nonlocal listener

        def on_press(key):
            if key == keyboard.Key.f20:
                app.toggle_recording()

        listener = keyboard.Listener(on_press=on_press)
        listener.start()
        logging.info("Started Windows F20 hotkey listener")

    def cleanup_toggle_trigger():
        nonlocal listener
        if listener:
            listener.stop()

    return Platform(
        notify=notify,
        dispatch=dispatch,
        open_logs=open_logs,
        setup_toggle_trigger=setup_toggle_trigger,
        cleanup_toggle_trigger=cleanup_toggle_trigger,
    )


def _make_linux_platform() -> Platform:
    prev_handler = None

    def notify(title: str, message: str):
        try:
            subprocess.Popen(
                ["notify-send", APP_NAME, f"{title}: {message}"],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        except Exception:
            logging.exception("Linux notification failed")

    def dispatch(task: str):
        dispatch_sh = Path.home() / ".local" / "bin" / "dispatch.sh"
        safe_query = shlex.quote(task)
        cmd = ["bash", str(dispatch_sh), safe_query]
        logging.info("Dispatching agent task with %d characters", len(task))
        subprocess.Popen(
            cmd,
            start_new_session=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )

    def open_logs():
        subprocess.Popen(
            ["xdg-open", str(LOG_FILE)],
            start_new_session=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )

    def setup_toggle_trigger(app: "BlurtApp"):
        nonlocal prev_handler
        PID_FILE.write_text(str(os.getpid()))
        logging.info("Wrote toggle PID file: %s", PID_FILE)

        def on_sigusr1(signum, frame):
            threading.Thread(target=app.toggle_recording, daemon=True).start()

        prev_handler = signal.signal(signal.SIGUSR1, on_sigusr1)

    def cleanup_toggle_trigger():
        nonlocal prev_handler
        if prev_handler is not None:
            signal.signal(signal.SIGUSR1, prev_handler)
        PID_FILE.unlink(missing_ok=True)

    return Platform(
        notify=notify,
        dispatch=dispatch,
        open_logs=open_logs,
        setup_toggle_trigger=setup_toggle_trigger,
        cleanup_toggle_trigger=cleanup_toggle_trigger,
    )


def get_platform() -> Platform:
    if sys.platform == "win32":
        return _make_windows_platform()
    return _make_linux_platform()


def send_toggle():
    """Send toggle signal to running blurt instance (Linux only)."""
    if sys.platform == "win32":
        print("'blurt toggle' is not supported on Windows.", file=sys.stderr)
        sys.exit(1)
    try:
        pid = int(PID_FILE.read_text().strip())
        logging.info("Sending toggle signal to PID %s", pid)
        os.kill(pid, signal.SIGUSR1)
    except (FileNotFoundError, ValueError, ProcessLookupError):
        print("No running blurt instance found.", file=sys.stderr)
        sys.exit(1)


class BlurtApp:
    def __init__(self, platform: Platform):
        self.platform = platform
        self.recording = False
        self.audio_frames: list[np.ndarray] = []
        self.stream: sd.InputStream | None = None
        self.lock = threading.Lock()
        self.icon: pystray.Icon | None = None
        self.device_index: int | None = None
        self.input_devices = self._get_input_devices()
        self.transcription_model = load_transcription_model()
        self.transcription_language = load_transcription_language()
        self.api_keys = {
            "openai": load_api_key(),
            "openrouter": load_openrouter_key(),
        }
        self.session = requests.Session()
        self.pending_transcriptions = 0
        logging.info(
            "BlurtApp initialized with %d input devices, model=%s, language=%s",
            len(self.input_devices),
            self.transcription_model,
            self.transcription_language,
        )

    def _get_input_devices(self) -> list[dict]:
        devices = sd.query_devices()
        input_devices = [
            {"index": i, "name": d["name"]}
            for i, d in enumerate(devices)
            if d["max_input_channels"] > 0
        ]
        for dev in input_devices:
            logging.info("Input device %s: %s", dev["index"], dev["name"])
        return input_devices

    def _create_icon_image(self, color: tuple[int, int, int]) -> Image.Image:
        img = Image.new("RGBA", (64, 64), (0, 0, 0, 0))
        draw = ImageDraw.Draw(img)
        draw.ellipse([8, 8, 56, 56], fill=color)
        return img

    def _update_icon(self):
        if self.icon:
            if self.recording:
                color = (220, 40, 40)  # red
                state = "Recording..."
            elif self.pending_transcriptions > 0:
                color = (40, 100, 220)  # blue
                state = "Transcribing..."
            else:
                color = (120, 120, 120)  # gray
                state = "Ready"
            self.icon.icon = self._create_icon_image(color)
            self.icon.title = f"{APP_NAME} - {state}"
            self.icon.update_menu()

    def _provider_for_model(self, model: str) -> str:
        for model_id, _, provider in TRANSCRIPTION_MODELS:
            if model_id == model:
                return provider
        return "openai"

    def _current_api_key(self) -> str:
        return self.api_keys.get(self._provider_for_model(self.transcription_model), "")

    def _audio_callback(self, indata, frames, time_info, status):
        if status:
            print(f"sounddevice: {status}", file=sys.stderr)
        self.audio_frames.append(indata.copy())

    def _start_recording(self):
        if not self._current_api_key():
            provider = self._provider_for_model(self.transcription_model)
            key_file = PROVIDER_KEY_FILES.get(provider, "api_key.txt")
            self.platform.notify("Error", f"No {provider} API key. Put your key in {key_file}")
            return
        self.audio_frames = []
        try:
            self.stream = sd.InputStream(
                device=self.device_index,
                samplerate=SAMPLE_RATE,
                channels=CHANNELS,
                dtype=DTYPE,
                callback=self._audio_callback,
            )
            self.stream.start()
            logging.info("Started recording with device_index=%s", self.device_index)
        except Exception as e:
            logging.exception("Could not start recording")
            self.platform.notify("Error", f"Could not start recording: {e}")
            return
        self.recording = True
        self._update_icon()
        self.platform.notify("Recording", "Press hotkey or click tray to stop.")

    def _stop_recording(self) -> np.ndarray | None:
        if self.stream:
            self.stream.stop()
            self.stream.close()
            self.stream = None
        self.recording = False
        self._update_icon()
        if not self.audio_frames:
            logging.warning("Stopped recording with no audio frames")
            return None
        audio_data = np.concatenate(self.audio_frames, axis=0)
        self.audio_frames = []
        logging.info(
            "Stopped recording: samples=%d shape=%s dtype=%s",
            len(audio_data),
            audio_data.shape,
            audio_data.dtype,
        )
        return audio_data

    def _transcribe(self, audio_data: np.ndarray):
        self.platform.notify(
            "Transcribing",
            f"Using {get_model_label(self.transcription_model)} ({self.transcription_language})",
        )
        logging.info("Starting transcription: samples=%d", len(audio_data))

        buffer = io.BytesIO()
        wavfile.write(buffer, SAMPLE_RATE, audio_data)
        wav_bytes = buffer.getvalue()
        logging.info("Encoded WAV in memory: %d bytes", len(wav_bytes))

        try:
            if self._provider_for_model(self.transcription_model) == "openrouter":
                result = self._transcribe_openrouter(wav_bytes)
            else:
                result = self._transcribe_openai(wav_bytes)
            text = result.get("text", "")
            logging.info("Transcription text length: %d", len(text))
            if "usage" in result:
                logging.info("Transcription usage: %s", result["usage"])

            if not text:
                error_msg = result.get("error", {}).get("message", "Unknown error")
                logging.error("Transcription response did not include text: %s", result)
                self.platform.notify("Error", f"Transcription failed: {error_msg}")
                return

            if text.lower().startswith(AGENT_PREFIX):
                task = text[len(AGENT_PREFIX):].lstrip(" ,.:;-")
                if task:
                    pyperclip.copy(task)
                    try:
                        self.platform.dispatch(task)
                    except Exception as e:
                        self.platform.notify("Dispatch Error", str(e))
                    preview = task[:50] + ("..." if len(task) > 50 else "")
                    self.platform.notify("Dispatched", preview)
                    logging.info("Dispatched transcription task with %d characters", len(task))
                else:
                    pyperclip.copy(text)
                    self.platform.notify("Error", f"No task after '{AGENT_PREFIX}'")
                    logging.warning("Agent prefix detected with no task")
            else:
                pyperclip.copy(text)
                preview = text[:50] + ("..." if len(text) > 50 else "")
                self.platform.notify("Copied", preview)
                logging.info("Copied transcription to clipboard")

        except requests.RequestException as e:
            logging.exception("API request failed")
            self.platform.notify("Error", f"API request failed: {e}")
        except Exception as e:
            logging.exception("Unexpected transcription failure")
            self.platform.notify("Error", f"Unexpected transcription failure: {e}")
        finally:
            self.pending_transcriptions -= 1
            self._update_icon()

    def _transcribe_openai(self, wav_bytes: bytes) -> dict:
        response = self.session.post(
            OPENAI_TRANSCRIPTIONS_URL,
            headers={"Authorization": f"Bearer {self.api_keys['openai']}"},
            files={"file": ("blurt.wav", wav_bytes, "audio/wav")},
            data={
                "model": self.transcription_model,
                "language": self.transcription_language,
            },
            timeout=30,
        )
        self._log_transcription_response(response)
        response.raise_for_status()
        return response.json()

    def _transcribe_openrouter(self, wav_bytes: bytes) -> dict:
        # OpenRouter's transcription endpoint takes a JSON body with base64-encoded
        # audio (not OpenAI's multipart upload). The response still has a "text" field.
        audio_b64 = base64.b64encode(wav_bytes).decode("ascii")
        response = self.session.post(
            OPENROUTER_TRANSCRIPTIONS_URL,
            headers={"Authorization": f"Bearer {self.api_keys['openrouter']}"},
            json={
                "model": self.transcription_model,
                "input_audio": {"data": audio_b64, "format": "wav"},
                "language": self.transcription_language,
            },
            timeout=30,
        )
        self._log_transcription_response(response)
        response.raise_for_status()
        return response.json()

    def _log_transcription_response(self, response: requests.Response):
        logging.info(
            "Transcription response: status=%s content_type=%s bytes=%d",
            response.status_code,
            response.headers.get("content-type", ""),
            len(response.content),
        )
        if not response.ok:
            logging.error("Transcription error response: %s", response.text[:1000])

    def toggle_recording(self):
        with self.lock:
            if not self.recording:
                self._start_recording()
            else:
                audio_data = self._stop_recording()
                if audio_data is not None:
                    self.pending_transcriptions += 1
                    self._update_icon()
                    t = threading.Thread(
                        target=self._transcribe, args=(audio_data,), daemon=True
                    )
                    t.start()
                else:
                    self.platform.notify("Error", "No audio recorded.")

    def _on_click(self, icon, item):
        self.toggle_recording()

    def _show_logs(self, icon, item):
        try:
            logging.info("Opening log file: %s", LOG_FILE)
            self.platform.open_logs()
        except Exception as e:
            logging.exception("Could not open log file")
            self.platform.notify("Error", f"Could not open log file: {e}")

    def _select_model(self, model: str):
        def handler(icon, item):
            self.transcription_model = model
            self.api_keys = {
                "openai": load_api_key(),
                "openrouter": load_openrouter_key(),
            }
            save_transcription_model(model)
            self.icon.update_menu()
            self.platform.notify("Model Selected", get_model_label(model))
        return handler

    def _select_language(self, language: str):
        def handler(icon, item):
            self.transcription_language = language
            save_transcription_language(language)
            self.icon.update_menu()
            self.platform.notify("Language Selected", get_language_label(language))
        return handler

    def _is_model_selected(self, model: str):
        def check(item):
            return self.transcription_model == model
        return check

    def _is_language_selected(self, language: str):
        def check(item):
            return self.transcription_language == language
        return check

    def _select_device(self, device_index):
        def handler(icon, item):
            self.device_index = device_index
            self.icon.update_menu()
        return handler

    def _is_device_selected(self, device_index):
        def check(item):
            return self.device_index == device_index
        return check

    def _build_menu(self):
        device_items = []
        for dev in self.input_devices:
            device_items.append(
                pystray.MenuItem(
                    dev["name"],
                    self._select_device(dev["index"]),
                    checked=self._is_device_selected(dev["index"]),
                    radio=True,
                )
            )

        model_items = []
        for model, label, _ in TRANSCRIPTION_MODELS:
            model_items.append(
                pystray.MenuItem(
                    label,
                    self._select_model(model),
                    checked=self._is_model_selected(model),
                    radio=True,
                )
            )

        language_items = []
        for language, label in TRANSCRIPTION_LANGUAGES:
            language_items.append(
                pystray.MenuItem(
                    label,
                    self._select_language(language),
                    checked=self._is_language_selected(language),
                    radio=True,
                )
            )

        return pystray.Menu(
            pystray.MenuItem(
                lambda item: "Stop Recording" if self.recording else "Start Recording",
                self._on_click,
                default=True,
            ),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("Input Device", pystray.Menu(*device_items)),
            pystray.MenuItem("Model", pystray.Menu(*model_items)),
            pystray.MenuItem("Language", pystray.Menu(*language_items)),
            pystray.MenuItem("Show Logs", self._show_logs),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("Quit", self._on_quit),
        )

    def _on_quit(self, icon, item):
        if self.recording:
            self._stop_recording()
        self.platform.cleanup_toggle_trigger()
        icon.stop()

    def _setup(self, icon):
        icon.visible = True

    def run(self):
        self.platform.setup_toggle_trigger(self)
        logging.info("Starting tray icon")
        self.icon = pystray.Icon(
            APP_NAME,
            icon=self._create_icon_image((120, 120, 120)),
            title=f"{APP_NAME} - Ready",
            menu=self._build_menu(),
        )
        self.icon.run(setup=self._setup)


def _silence_benign_gtk_warnings():
    """Drop a harmless, repeating Gdk-CRITICAL emitted by Gtk.StatusIcon.

    pystray's GTK backend (the one qtile's XEmbed Systray widget needs) uses the
    deprecated Gtk.StatusIcon, which logs this assertion on icon/menu updates:
        gdk_window_thaw_toplevel_updates: assertion
        'window->update_and_descendants_freeze_count > 0' failed
    It is cosmetic. Suppress only that message and re-emit every other Gdk log so
    genuine problems still surface.
    """
    try:
        from gi.repository import GLib
    except Exception:
        return

    benign = "gdk_window_thaw_toplevel_updates"

    def handler(domain, level, message, _user_data):
        if benign not in message:
            print(f"{domain}: {message}", file=sys.stderr)

    GLib.log_set_handler(
        "Gdk",
        GLib.LogLevelFlags.LEVEL_CRITICAL
        | GLib.LogLevelFlags.LEVEL_WARNING
        | GLib.LogLevelFlags.FLAG_FATAL
        | GLib.LogLevelFlags.FLAG_RECURSION,
        handler,
        None,
    )


def main():
    setup_logging()
    logging.info("Starting %s argv=%s platform=%s log_file=%s", APP_NAME, sys.argv, sys.platform, LOG_FILE)
    if len(sys.argv) > 1 and sys.argv[1] == "toggle":
        send_toggle()
        return

    _silence_benign_gtk_warnings()
    platform = get_platform()
    app = BlurtApp(platform)
    app.run()


if __name__ == "__main__":
    main()