# claude buddy v1 — claude desktop hardware companion
# target: Adafruit ESP32-S3 Reverse TFT Feather (board #5691), CircuitPython 10.1.4
#
# what it does:
#   - advertises Nordic UART Service as "Claude-Feather"
#   - parses heartbeat/turn/time/owner/status JSON from Claude Desktop
#   - shows state on the 240x135 TFT: disconnected / idle / busy / approval / turn
#   - D0 = approve once, D1 = deny, D2 = next info screen
#   - NeoPixel = ambient status light (blue adv / green idle / yellow busy / red-pulse approval)
#
# protocol ref: github.com/anthropics/claude-desktop-buddy REFERENCE.md

import time
import json
import board
import digitalio
import terminalio
import displayio
import neopixel
from adafruit_display_text import label
from adafruit_display_shapes.rect import Rect

from adafruit_ble import BLERadio
from adafruit_ble.advertising.standard import ProvideServicesAdvertisement
from adafruit_ble.services.nordic import UARTService

# ---------------- hardware ----------------
display = board.DISPLAY
display.brightness = 0.9

# neopixel
pix_pwr = digitalio.DigitalInOut(board.NEOPIXEL_POWER)
pix_pwr.switch_to_output(value=True)
pix = neopixel.NeoPixel(board.NEOPIXEL, 1, brightness=0.25, auto_write=True)

# buttons: D0=BOOT (pull-up, active-low), D1/D2 (pull-down, active-high on Reverse TFT)
btn_d0 = digitalio.DigitalInOut(board.D0); btn_d0.switch_to_input(pull=digitalio.Pull.UP)
btn_d1 = digitalio.DigitalInOut(board.D1); btn_d1.switch_to_input(pull=digitalio.Pull.DOWN)
btn_d2 = digitalio.DigitalInOut(board.D2); btn_d2.switch_to_input(pull=digitalio.Pull.DOWN)

def _d0_pressed(): return not btn_d0.value     # BOOT is active-low
def _d1_pressed(): return btn_d1.value         # active-high
def _d2_pressed(): return btn_d2.value         # active-high

# ---------------- theme ----------------
BG        = 0x000000
FG        = 0xFFFFFF
MUTED     = 0x888888
OK        = 0x00FF66
BUSY      = 0xFFCC00
ALERT     = 0xFF3366
HINT      = 0x66CCFF

# ---------------- display: build root group ----------------
W, H = 240, 135
root = displayio.Group()
display.root_group = root

bg = Rect(0, 0, W, H, fill=BG)
root.append(bg)

# status bar (top)
status_bar = Rect(0, 0, W, 16, fill=0x222222)
root.append(status_bar)
lbl_status = label.Label(terminalio.FONT, text="booting…", color=MUTED, x=4, y=8)
root.append(lbl_status)
lbl_name   = label.Label(terminalio.FONT, text="Claude-Feather", color=HINT, x=W-96, y=8)
root.append(lbl_name)

# main area — 3 stacked lines (big)
lbl_l1 = label.Label(terminalio.FONT, text="", color=FG, x=8, y=38, scale=2)
lbl_l2 = label.Label(terminalio.FONT, text="", color=FG, x=8, y=64, scale=2)
lbl_l3 = label.Label(terminalio.FONT, text="", color=MUTED, x=8, y=88, scale=1)
root.append(lbl_l1); root.append(lbl_l2); root.append(lbl_l3)

# button legend (bottom)
lbl_btns  = label.Label(terminalio.FONT, text="", color=MUTED, x=4, y=H-8)
root.append(lbl_btns)

def set_screen(status, s_color, l1="", l2="", l3="", btns=""):
    lbl_status.text = status[:28]
    lbl_status.color = s_color
    lbl_l1.text = l1[:18]
    lbl_l2.text = l2[:18]
    lbl_l3.text = l3[:38]
    lbl_btns.text = btns[:38]

# ---------------- BLE ----------------
ble = BLERadio()
mac = ble.address_bytes
short = "".join("%02X" % b for b in mac[-2:])
ble.name = "Claude-Feather-" + short  # distinguish multiple boards
uart = UARTService()
adv = ProvideServicesAdvertisement(uart)

# ---------------- runtime state ----------------
state = {
    "phase":       "DISCONNECTED",  # DISCONNECTED / IDLE / BUSY / APPROVAL / TURN
    "owner":       "",
    "total":       0,
    "running":     0,
    "waiting":     0,
    "msg":         "",
    "tokens_today":0,
    "prompt":      None,
    "last_rx":     0,
    "time_epoch":  0,
    "stats": {"appr": 0, "deny": 0},
    "view":        0,  # D2 cycles between info pages
}

def now_s():
    return time.monotonic()

# ---------------- rendering ----------------
def render():
    p = state["phase"]
    if p == "DISCONNECTED":
        pix[0] = (0, 0, 40)
        set_screen("advertising", HINT, "Claude", "buddy", "waiting for desktop connection",
                   "Developer > Open Hardware Buddy > Connect")
        return

    # connected branches
    owner = state["owner"] or "you"
    if p == "APPROVAL":
        tool = ""
        pid = ""
        if state["prompt"]:
            tool = state["prompt"].get("tool", "?")
            pid = state["prompt"].get("id", "")[-6:]
        pix[0] = (60, 0, 0) if int(now_s()*2) % 2 == 0 else (0, 0, 0)  # pulse red
        set_screen("APPROVAL NEEDED", ALERT,
                   "approve?", tool,
                   "req " + pid,
                   "D0=once  D1=deny  D2=info")
        return

    if p == "BUSY":
        pix[0] = (40, 30, 0)  # yellow
    elif p == "IDLE":
        pix[0] = (0, 40, 0)   # green
    else:
        pix[0] = (20, 20, 20)

    # info views cycle via D2
    v = state["view"]
    if v == 0:
        set_screen(p.lower(), OK if p=="IDLE" else BUSY,
                   "sessions", "%d run %d wait" % (state["running"], state["waiting"]),
                   "total %d   hi %s" % (state["total"], owner),
                   "D2=next")
    elif v == 1:
        toks = state["tokens_today"]
        ts = "%dk" % (toks//1000) if toks >= 1000 else str(toks)
        set_screen("tokens", HINT,
                   ts, "today",
                   "cumulative: %d" % toks,
                   "D2=next")
    else:
        ap = state["stats"].get("appr", 0)
        dn = state["stats"].get("deny", 0)
        set_screen("approvals", HINT,
                   "y: %d" % ap, "n: %d" % dn,
                   state["msg"][:38] or "all quiet",
                   "D2=next")

# ---------------- protocol handling ----------------
def send_line(obj):
    try:
        s = json.dumps(obj) + "\n"
        uart.write(s.encode("utf-8"))
    except Exception as e:
        print("send err:", e)

def handle_message(obj):
    """Apply one parsed desktop message to state."""
    # acks for commands
    if "cmd" in obj:
        cmd = obj["cmd"]
        if cmd == "status":
            send_line({"ack":"status","ok":True,"data":{
                "name": ble.name, "sec": False,
                "sys": {"up": int(now_s())},
                "stats": state["stats"],
            }})
            return
        if cmd == "name":
            state["owner"] = obj.get("name", state["owner"])
            send_line({"ack":"name","ok":True,"n":0})
            return
        if cmd == "owner":
            state["owner"] = obj.get("name", "")
            send_line({"ack":"owner","ok":True,"n":0})
            return
        if cmd == "unpair":
            send_line({"ack":"unpair","ok":True,"n":0})
            return
        # unknown command → ack false
        send_line({"ack":cmd,"ok":False,"error":"unknown"})
        return

    # time sync (one-shot)
    if "time" in obj and isinstance(obj["time"], list):
        state["time_epoch"] = obj["time"][0]
        return

    # turn event
    if obj.get("evt") == "turn":
        # brief acknowledgment — render will pick up next loop
        state["msg"] = "turn: " + obj.get("role", "?")
        return

    # heartbeat snapshot
    if "total" in obj or "prompt" in obj:
        state["total"]        = obj.get("total", state["total"])
        state["running"]      = obj.get("running", state["running"])
        state["waiting"]      = obj.get("waiting", state["waiting"])
        state["msg"]          = obj.get("msg", state["msg"])
        state["tokens_today"] = obj.get("tokens_today", state["tokens_today"])
        state["prompt"]       = obj.get("prompt", None)
        # phase decision
        if state["prompt"]:
            state["phase"] = "APPROVAL"
        elif state["running"] > 0:
            state["phase"] = "BUSY"
        else:
            state["phase"] = "IDLE"
        return

# ---------------- button debounce ----------------
class Btn:
    def __init__(self, fn):
        self.fn = fn
        self.prev = fn()
        self.last_change = 0
    def edge_pressed(self):
        v = self.fn()
        t = now_s()
        if v != self.prev and (t - self.last_change) > 0.04:
            self.last_change = t
            self.prev = v
            return v  # returns True on press edge
        self.prev = v
        return False

b_d0 = Btn(_d0_pressed)
b_d1 = Btn(_d1_pressed)
b_d2 = Btn(_d2_pressed)

def send_permission(decision):
    if not state["prompt"]: return
    pid = state["prompt"].get("id")
    if not pid: return
    send_line({"cmd":"permission","id":pid,"decision":decision})
    if decision == "once":
        state["stats"]["appr"] += 1
    elif decision == "deny":
        state["stats"]["deny"] += 1
    # optimistically clear prompt locally; next heartbeat may re-set if desktop still waiting
    state["prompt"] = None
    state["phase"] = "IDLE" if state["running"] == 0 else "BUSY"

# ---------------- main loop ----------------
print("claude-buddy v1 starting, name=", ble.name)
set_screen("booting…", MUTED, "Claude", "buddy", "initializing BLE…", "")

rx_buf = b""
last_render = 0

while True:
    # advertise until connected
    state["phase"] = "DISCONNECTED"
    state["prompt"] = None
    render()
    if not ble.advertising:
        ble.start_advertising(adv)
        print("advertising as", ble.name)

    while not ble.connected:
        render()
        time.sleep(0.2)

    ble.stop_advertising()
    print("connected")
    state["phase"] = "IDLE"
    state["last_rx"] = now_s()
    rx_buf = b""

    while ble.connected:
        # --- BLE RX ---
        if uart.in_waiting:
            chunk = uart.read(uart.in_waiting)
            if chunk:
                rx_buf += chunk
                state["last_rx"] = now_s()
                while b"\n" in rx_buf:
                    line, rx_buf = rx_buf.split(b"\n", 1)
                    line = line.strip()
                    if not line: continue
                    try:
                        obj = json.loads(line.decode("utf-8"))
                        handle_message(obj)
                    except Exception as e:
                        print("parse err:", e, line[:80])

        # --- buttons ---
        if b_d0.edge_pressed():
            if state["phase"] == "APPROVAL":
                send_permission("once")
            else:
                print("D0 (no-op outside approval)")
        if b_d1.edge_pressed():
            if state["phase"] == "APPROVAL":
                send_permission("deny")
            else:
                print("D1 (no-op outside approval)")
        if b_d2.edge_pressed():
            state["view"] = (state["view"] + 1) % 3
            print("D2 -> view", state["view"])

        # --- heartbeat watchdog (30s per spec) ---
        if now_s() - state["last_rx"] > 30 and state["phase"] != "APPROVAL":
            # don't force disconnect; just note stale
            state["msg"] = "stale heartbeat"

        # --- render at ~5Hz ---
        t = now_s()
        if t - last_render > 0.2:
            render()
            last_render = t

        time.sleep(0.01)

    print("disconnected")
    pix[0] = (40, 0, 0)
    time.sleep(0.5)