"""pico-bridge: Claude status -> Raspberry Pi Pico over USB serial.

Sends JSON lines that main.py on the Pico understands:
  {"t":"chats","n":3,"l":[["name","w","c"],...]}
      status: w = waiting for you, b = busy, i = idle; kind: c = Code, k = Cowork
  {"t":"lim","h5":7,"h5s":16050,"d7":39,"d7s":171450}

Sources (all undocumented, may change):
  chats  - ~/.claude/sessions/<pid>.json, written by every running Claude Code
  cowork - Windows only: the Claude app's toasts in the notification center (title = task
           name); needs the winrt-Windows.UI.Notifications.Management package
  limits - api.anthropic.com/api/oauth/usage with the CLI login in ~/.claude/.credentials.json

Usage:
  python pico_bridge.py            run (finds the Pico by USB VID)
  python pico_bridge.py --port COM3
  python pico_bridge.py --list     print what would be sent and exit
  python pico_bridge.py --pause    let the port go (for Thonny); --resume takes it back
Runs on Windows and Linux; needs pyserial. Without a console (pythonw.exe, systemd)
the log goes to state/bridge.log next to this file. Only one copy runs at a time.
"""

import ctypes
import shutil
import subprocess
import email.utils
import json
import os
import sys
import threading
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone

import serial
import serial.tools.list_ports

WINDOWS = sys.platform == "win32"
PICO_VID = 0x2E8A
CLAUDE_DIR = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.join(os.path.expanduser("~"), ".claude")
SESSIONS_DIR = os.path.join(CLAUDE_DIR, "sessions")
CREDENTIALS = os.path.join(CLAUDE_DIR, ".credentials.json")
USAGE_URL = "https://api.anthropic.com/api/oauth/usage"

MAX_CHATS = 8
NAME_LEN = 48         # the Pico scrolls names that do not fit
CHECK_S = 1          # how often the sessions folder is looked at (a stat of a few files)
PID_CHECK_S = 5      # how often running processes are re-checked
RESEND_S = 30        # full state is resent this often; the Pico shows "no link" after 90 s
RECONNECT_S = 3
# Usage API: often while you look at Claude, rarely when away.
USAGE_FOCUSED_S = 120
USAGE_ACTIVE_S = 600
USAGE_IDLE_S = 1800
IDLE_AFTER_S = 15 * 60
TOKEN_CHECK_S = 30
# The CLI login token lives ~8 h. "claude auth status" makes the CLI refresh it
# itself (no usage spent, ~1 s), shortly before it runs out or after a 401.
REFRESH_AHEAD_S = 10 * 60
REFRESH_GAP_S = 10 * 60       # not more often than this, e.g. when logged out


# Next to this file, not in AppData: the Claude app is MSIX, and what runs from it
# sees a private copy of AppData, so a pause file made there would not reach us.
DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "state")
LOG_FILE = os.path.join(DATA_DIR, "bridge.log")
LOG_MAX = 512 * 1024          # bridge.log is moved to bridge.log.1 beyond this
# While this file exists the port is left free for Thonny or pico_put.py.
PAUSE_FILE = os.path.join(DATA_DIR, "pause")
os.makedirs(DATA_DIR, exist_ok=True)


def log(*parts):
    line = " ".join([time.strftime("%m-%d %H:%M:%S")] + [str(p) for p in parts])
    try:
        if sys.stdout.isatty():           # a console; under pythonw.exe stdout is None
            print(line, flush=True)
            return
    except (AttributeError, ValueError, OSError):
        pass
    try:
        if os.path.getsize(LOG_FILE) > LOG_MAX:
            os.replace(LOG_FILE, LOG_FILE + ".1")
    except OSError:
        pass
    try:
        with open(LOG_FILE, "a", encoding="utf-8") as handle:
            handle.write(line + "\n")
    except OSError:
        pass


# --- what the Pico can draw: ASCII, Russian letters (rufont.py), a few marks ----

RUSSIAN = set("абвгдеёжзийклмнопрстуфхцчшщъыьэюя" "АБВГДЕЁЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯ" "«»—–")
# Kazakh letters the font lacks, as their nearest Russian ones.
KAZAKH = dict(zip("әғқңөұүһіӘҒҚҢӨҰҮҺІ", "агкноуухiАГКНОУУХI"))


def to_screen(text, limit=NAME_LEN):
    out = []
    for char in text:
        char = KAZAKH.get(char, char)
        if (char.isascii() and char.isprintable()) or char in RUSSIAN:
            out.append(char)
        elif char.isspace():
            if out and out[-1] != " ":
                out.append(" ")
        else:
            out.append("?")
    return "".join(out)[:limit].rstrip()


# --- system helpers -----------------------------------------------------------

if WINDOWS:
    from ctypes import wintypes

    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
    user32 = ctypes.WinDLL("user32", use_last_error=True)
    kernel32.OpenProcess.restype = wintypes.HANDLE
    PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
    STILL_ACTIVE = 259


    class FILETIME(ctypes.Structure):
        _fields_ = [("low", wintypes.DWORD), ("high", wintypes.DWORD)]


    def process_alive(pid, proc_start=None):
        """True if pid runs and, when proc_start (a FILETIME) is known, it is still the
        same process and not a newer one that got the same PID."""
        handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
        if not handle:
            return False
        try:
            code = wintypes.DWORD()
            if not kernel32.GetExitCodeProcess(handle, ctypes.byref(code)) or code.value != STILL_ACTIVE:
                return False
            if proc_start:
                created, exited, kernel, user = FILETIME(), FILETIME(), FILETIME(), FILETIME()
                if kernel32.GetProcessTimes(handle, ctypes.byref(created), ctypes.byref(exited),
                                            ctypes.byref(kernel), ctypes.byref(user)):
                    actual = created.high << 32 | created.low
                    return abs(actual - int(proc_start)) <= 20_000_000   # 2 s
            return True
        except ValueError:
            return True
        finally:
            kernel32.CloseHandle(handle)


    class LASTINPUTINFO(ctypes.Structure):
        _fields_ = [("cbSize", wintypes.UINT), ("dwTime", wintypes.DWORD)]


    def idle_seconds():
        info = LASTINPUTINFO(ctypes.sizeof(LASTINPUTINFO), 0)
        if not user32.GetLastInputInfo(ctypes.byref(info)):
            return 0
        return ((kernel32.GetTickCount() - info.dwTime) & 0xFFFFFFFF) / 1000


    def claude_focused():
        """True if the foreground window belongs to the Claude app."""
        hwnd = user32.GetForegroundWindow()
        pid = wintypes.DWORD()
        user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid))
        handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid.value)
        if not handle:
            return False
        try:
            size = wintypes.DWORD(260)
            buffer = ctypes.create_unicode_buffer(size.value)
            if kernel32.QueryFullProcessImageNameW(handle, 0, buffer, ctypes.byref(size)):
                return os.path.basename(buffer.value).lower() == "claude.exe"
            return False
        finally:
            kernel32.CloseHandle(handle)

else:
    def process_alive(pid, proc_start=None):
        try:
            os.kill(pid, 0)
        except ProcessLookupError:
            return False
        except PermissionError:
            return True
        return True

    def idle_seconds():
        return 0                  # unknown here: limits are polled as if you were active

    def claude_focused():
        return False


# --- chats --------------------------------------------------------------------

STATUS_ORDER = {"w": 0, "b": 1, "i": 2}
CODE_NAMES = set()     # every Code chat name seen since start


def read_chats(alive_cache):
    """Live chats, waiting first, then busy, then idle; newest first in each group."""
    chats = {}
    try:
        entries = list(os.scandir(SESSIONS_DIR))
    except OSError:
        return []
    for entry in entries:
        if not entry.name.endswith(".json"):
            continue
        try:
            with open(entry.path, encoding="utf-8") as handle:
                raw = json.load(handle)
            pid = int(raw["pid"])
        except (OSError, ValueError, KeyError, TypeError):
            continue          # mid-write; it is read again on the next change
        if raw.get("name"):
            CODE_NAMES.add(raw["name"].strip())   # its toasts are not Cowork
        key = (pid, raw.get("procStart"))
        if key not in alive_cache:
            alive_cache[key] = process_alive(pid, raw.get("procStart"))
        if not alive_cache[key]:
            continue
        status = {"busy": "b", "waiting": "w"}.get(raw.get("status"), "i")
        updated = raw.get("statusUpdatedAt") or raw.get("updatedAt") or 0
        name = (raw.get("name") or "").strip() or "untitled"
        # A resumed chat gets a new process; keep the newest entry per app session.
        chat_id = raw.get("hostSessionId") or raw.get("sessionId") or str(pid)
        if chat_id not in chats or chats[chat_id][2] < updated:
            chats[chat_id] = (name, status, updated)
    return sorted(chats.values(), key=lambda c: (STATUS_ORDER[c[1]], -c[2]))


# --- Cowork, from the app's Windows notifications ------------------------------

CLAUDE_AUMID_PART = "claude"   # AppUserModelId is Claude_<publisher>!Claude
COWORK_TTL_S = 30 * 60         # a Cowork task stays on the list this long after its last toast
TOAST_POLL_S = 2
WAITING_WORDS = ("ответ", "разреш", "подтверд", "answer", "approv", "permission", "input", "?")


class Toasts:
    """Polls the notification center. Code chats toast too; the main loop drops
    those by name, the rest are Cowork tasks."""

    def __init__(self, start=True):
        self.lock = threading.Lock()
        self.tasks = {}          # title -> [status, toast time, body]
        self.version = 0
        self.ok = False
        if not WINDOWS:
            return                # Cowork notifications are read on Windows only
        try:
            from winrt.windows.ui.notifications import NotificationKinds, KnownNotificationBindings
            from winrt.windows.ui.notifications.management import UserNotificationListener
        except ImportError:
            log("no winrt package: Cowork is off (pip install winrt-Windows.UI.Notifications.Management ...)")
            return
        self.kinds, self.bindings = NotificationKinds, KnownNotificationBindings
        self.listener = UserNotificationListener.current
        if self.listener.get_access_status() != 1:     # 1 = allowed
            log("no access to notifications: Cowork is off")
            return
        self.ok = True
        if start:
            threading.Thread(target=self.run, daemon=True).start()

    def snapshot(self):
        with self.lock:
            return {title: tuple(task) for title, task in self.tasks.items()}

    def read(self):
        import asyncio
        items = asyncio.run(self.listener.get_notifications_async(self.kinds.TOAST))
        newest = {}
        for item in items:
            try:
                if CLAUDE_AUMID_PART not in item.app_info.app_user_model_id.lower():
                    continue
                binding = item.notification.visual.get_binding(self.bindings.toast_generic)
                texts = [t.text for t in binding.get_text_elements()] if binding else []
                stamp = item.creation_time.timestamp()
            except (AttributeError, OSError):
                continue
            if not texts or time.time() - stamp > COWORK_TTL_S:
                continue
            title, body = texts[0].strip(), " ".join(texts[1:])
            if title not in newest or newest[title][0] < stamp:
                newest[title] = (stamp, body)
        return newest

    def run(self):
        while True:
            try:
                newest = self.read()
            except OSError as error:
                log("notifications:", error)
                newest = None
            if newest is not None:
                focused = claude_focused()
                with self.lock:
                    before = dict((k, tuple(v)) for k, v in self.tasks.items())
                    for title in list(self.tasks):
                        if title not in newest:
                            del self.tasks[title]       # expired or dismissed
                    for title, (stamp, body) in newest.items():
                        task = self.tasks.get(title)
                        if task is None or task[1] < stamp:
                            asks = any(word in body.lower() for word in WAITING_WORDS)
                            self.tasks[title] = ["w" if asks else "i", stamp, body]
                        elif task[0] == "w" and focused:
                            task[0] = "i"                # you opened Claude to answer
                    if before != dict((k, tuple(v)) for k, v in self.tasks.items()):
                        self.version += 1
            time.sleep(TOAST_POLL_S)


def merge(code_chats, toasts, code_names):
    """Code chats plus Cowork tasks, waiting first, then busy, then idle."""
    rows = [(name, status, updated, "c") for name, status, updated in code_chats]
    for title, (status, stamp, _) in toasts.items():
        # The app may title a Code chat's toast "<folder> • <chat name>".
        if not any(name and name in title for name in code_names):
            rows.append((title, status, int(stamp * 1000), "k"))
    return sorted(rows, key=lambda r: (STATUS_ORDER[r[1]], -r[2]))


def folder_signature():
    try:
        return tuple(sorted((e.name, e.stat().st_mtime_ns, e.stat().st_size)
                            for e in os.scandir(SESSIONS_DIR)))
    except OSError:
        return ()


def chats_line(chats):
    items = [[to_screen(name), status, kind] for name, status, _, kind in chats[:MAX_CHATS]]
    return json.dumps({"t": "chats", "n": len(chats), "l": items}, separators=(",", ":"),
                      ensure_ascii=False)


# --- limits -------------------------------------------------------------------

class Usage:
    """Polls the usage API in a background thread; keeps the last good answer."""

    def __init__(self, start=True):
        self.lock = threading.Lock()
        self.data = None          # (5h %, 5h reset epoch, 7d %, 7d reset epoch), server clock
        self.version = 0          # bumped on every new answer
        self.fetched_at = 0.0     # time.monotonic() of the answer
        self.server_now = 0.0
        self.wake = threading.Event()
        self.warned = False
        self.stale = False        # the last answer was 401: the token in the file expired
        self.bad_token = None     # that token; it is not sent again until the file changes
        self.expires = 0          # token expiry, PC clock (the CLI writes it with that clock)
        self.refreshed_at = -REFRESH_GAP_S
        if start:
            threading.Thread(target=self.run, daemon=True).start()

    def line(self):
        """lim message with seconds to reset, counted with the server's clock."""
        with self.lock:
            if self.data is None:
                return None
            h5, h5_reset, d7, d7_reset = self.data
            now = self.server_now + (time.monotonic() - self.fetched_at)
        message = {"t": "lim", "h5": round(h5), "h5s": max(0, int(h5_reset - now)),
                   "d7": round(d7), "d7s": max(0, int(d7_reset - now))}
        if self.stale:
            message["old"] = 1
        return json.dumps(message, separators=(",", ":"))

    def read_login(self):
        """(access token, expiry epoch by the PC clock) from the CLI login, or (None, 0)."""
        try:
            with open(CREDENTIALS, encoding="utf-8") as handle:
                login = json.load(handle)["claudeAiOauth"]
            return login["accessToken"], (login.get("expiresAt") or 0) / 1000
        except (OSError, ValueError, KeyError, TypeError):
            if not self.warned:
                log("no CLI login in", CREDENTIALS, "- limits are off (run: claude, then /login)")
                self.warned = True
            return None, 0

    def refresh_login(self, why):
        """Has the CLI refresh its own token. The daemon never writes the credentials:
        the CLI rotates the refresh token, and two refreshers could log it out."""
        if time.monotonic() - self.refreshed_at < REFRESH_GAP_S:
            return False
        self.refreshed_at = time.monotonic()
        exe = shutil.which("claude") or os.path.join(
            os.path.expanduser("~"), ".local", "bin", "claude.exe" if WINDOWS else "claude")
        try:
            result = subprocess.run([exe, "auth", "status"], stdin=subprocess.DEVNULL,
                                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                                    timeout=60,
                                    creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        except (OSError, subprocess.TimeoutExpired) as error:
            log("token refresh (%s): cannot run claude -" % why, error)
            return False
        _, expires = self.read_login()
        log("token refresh (%s): exit %d, valid %.1f h more"
            % (why, result.returncode, (expires - time.time()) / 3600))
        return True

    def fetch(self, retry=True):
        token, expires = self.read_login()
        if token is None:
            return None
        if expires and expires - time.time() < REFRESH_AHEAD_S and self.refresh_login("expiring"):
            token, expires = self.read_login()
        self.expires = expires
        if token == self.bad_token:
            return None           # wait for the CLI to put a fresh token there
        request = urllib.request.Request(USAGE_URL, headers={
            "Authorization": "Bearer " + token,
            "anthropic-beta": "oauth-2025-04-20",
            "User-Agent": "pico-bridge",
        })
        try:
            with urllib.request.urlopen(request, timeout=20) as response:
                body = json.load(response)
                date = response.headers.get("Date")
        except urllib.error.HTTPError as error:
            if error.code == 401:
                if retry and self.refresh_login("401"):
                    return self.fetch(retry=False)
                log("usage API: 401, token expired - limits wait for a fresh CLI login")
                self.bad_token = token
                with self.lock:
                    self.stale = True
                    self.version += 1
            else:
                log("usage API:", error.code)
            return None
        except (OSError, ValueError) as error:
            log("usage API:", error)
            return None
        # The PC clock is wrong, so time comes from the server's Date header.
        server_now = email.utils.parsedate_to_datetime(date).timestamp() if date else time.time()

        def window(name):
            part = body.get(name) or {}
            reset = part.get("resets_at")
            reset = datetime.fromisoformat(reset).timestamp() if reset else server_now
            return part.get("utilization") or 0, reset

        h5, h5_reset = window("five_hour")
        d7, d7_reset = window("seven_day")
        with self.lock:
            self.data = (h5, h5_reset, d7, d7_reset)
            self.server_now = server_now
            self.fetched_at = time.monotonic()
            self.stale = False
            self.version += 1
        if self.bad_token:
            log("usage API: fresh token, limits are back")
            self.bad_token = None
        return min(h5_reset, d7_reset) - server_now

    def run(self):
        while True:
            to_reset = self.fetch()
            if claude_focused():
                wait = USAGE_FOCUSED_S
            elif idle_seconds() < IDLE_AFTER_S:
                wait = USAGE_ACTIVE_S
            else:
                wait = USAGE_IDLE_S
            if to_reset is not None and 0 < to_reset + 5 < wait:
                wait = to_reset + 5          # fetch right after a window resets
            if self.expires:
                until_refresh = self.expires - REFRESH_AHEAD_S - time.time() + 1
                if until_refresh > 0:        # wake up in time to refresh it
                    wait = min(wait, until_refresh)
            if self.bad_token:
                wait = TOKEN_CHECK_S         # reading the file is cheap; no request is sent
            self.wake.wait(wait)
            self.wake.clear()


# --- serial link ----------------------------------------------------------------

def find_port():
    for port in serial.tools.list_ports.comports():
        if port.vid == PICO_VID:
            return port.device
    return None


class Link:
    def __init__(self, name):
        self.name = name
        self.port = serial.Serial(name, 115200, timeout=0.5, write_timeout=2)
        self.port.dtr = True      # MicroPython only talks to the host with DTR set
        self.alive = True
        threading.Thread(target=self.drain, daemon=True).start()

    def drain(self):
        """Show what the Pico prints, and keep its USB buffer from filling up."""
        while self.alive:
            try:
                data = self.port.readline()
            except (serial.SerialException, OSError):
                self.alive = False
                return
            text = data.decode("utf-8", "replace").strip()
            if text:
                log("pico>", text)

    def send(self, line):
        self.port.write(line.encode("utf-8") + b"\n")

    def close(self):
        self.alive = False
        try:
            self.port.close()
        except (serial.SerialException, OSError):
            pass


# --- main loop ----------------------------------------------------------------

def single_instance():
    """False if another pico-bridge already runs (a named mutex or a locked file,
    both freed when it exits)."""
    if not WINDOWS:
        import fcntl
        single_instance.handle = open(os.path.join(DATA_DIR, "lock"), "w")
        try:
            fcntl.flock(single_instance.handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except OSError:
            return False
        return True
    kernel32.CreateMutexW.restype = wintypes.HANDLE
    single_instance.handle = kernel32.CreateMutexW(None, False, r"Local\pico-bridge")
    return ctypes.get_last_error() != 183          # ERROR_ALREADY_EXISTS


def paused():
    return os.path.exists(PAUSE_FILE)


def main():
    args = sys.argv[1:]
    fixed_port = args[args.index("--port") + 1] if "--port" in args else None

    if "--pause" in args:
        open(PAUSE_FILE, "w").close()
        print("pico-bridge lets the port go (within", CHECK_S, "s). Undo: --resume")
        return
    if "--resume" in args:
        if paused():
            os.remove(PAUSE_FILE)
        print("pico-bridge takes the port again")
        return

    if "--list" in args:
        code = read_chats({})
        cowork = {}
        probe = Toasts(start=False)
        if probe.ok:
            for title, (stamp, body) in probe.read().items():
                asks = any(word in body.lower() for word in WAITING_WORDS)
                cowork[title] = ("w" if asks else "i", stamp, body)
        chats = merge(code, cowork, CODE_NAMES)
        for name, status, _, kind in chats:
            print(kind, status, name)
        print(chats_line(chats))
        usage = Usage(start=False)
        usage.fetch()
        print(usage.line())
        return

    if not single_instance():
        log("pico-bridge is already running")
        return
    log("start, watching", SESSIONS_DIR)
    usage = Usage()
    toasts = Toasts()
    toasts_seen = -1
    code_chats = []
    link = None
    last_attempt = 0.0
    sent_chats = None
    sent_usage = None         # Usage.version last sent; the Pico counts down by itself
    last_resend = 0.0
    signature = None
    alive_cache = {}
    last_pid_check = 0.0
    chats = chats_line([])
    open_error = None
    was_paused = False

    while True:
        now = time.monotonic()
        if now - last_pid_check >= PID_CHECK_S:
            alive_cache.clear()   # notice processes that are gone
            last_pid_check = now
            signature = None
        current = folder_signature()
        if current != signature or toasts.version != toasts_seen:
            if current != signature:
                signature = current
                code_chats = read_chats(alive_cache)
            toasts_seen = toasts.version
            chats = chats_line(merge(code_chats, toasts.snapshot(), CODE_NAMES))

        if link is not None and not link.alive:
            log("lost", link.name)
            link.close()
            link = None
        if paused() != was_paused:
            was_paused = not was_paused
            log("paused: port is free" if was_paused else "resumed")
            if was_paused and link is not None:
                link.close()
                link = None
        if link is None and not was_paused and now - last_attempt >= RECONNECT_S:
            last_attempt = now
            name = fixed_port or find_port()
            if name:
                try:
                    link = Link(name)
                    log("connected to", name)
                    open_error = None
                    sent_chats = sent_usage = None   # the Pico may have rebooted
                except (serial.SerialException, OSError) as error:
                    if str(error) != open_error:   # Thonny holding the port: say it once
                        open_error = str(error)
                        log("cannot open", name, "-", error)

        if link is not None:
            resend = now - last_resend >= RESEND_S
            try:
                if resend or chats != sent_chats:
                    if chats != sent_chats:
                        log("send", chats)
                    link.send(chats)
                    sent_chats = chats
                if usage.data and (resend or usage.version != sent_usage):
                    line = usage.line()
                    if usage.version != sent_usage:
                        log("send", line)
                    link.send(line)
                    sent_usage = usage.version
                if resend:
                    last_resend = now
            except (serial.SerialException, OSError) as error:
                log("lost", link.name, "-", error)
                link.close()
                link = None
        time.sleep(CHECK_S)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass
    except Exception:
        import traceback
        log("crash:", traceback.format_exc())
        sys.exit(1)               # the scheduled task starts it again