# SPDX-License-Identifier: MIT
"""
DIY "Codex Micro" for the Adafruit MacroPad RP2040.

A multi-mode macro keypad for driving Codex, Cursor, Claude Code, and Salesforce
DX. Press the rotary encoder to cycle modes; each mode remaps all 12 keys, recolors
the NeoPixels, and updates the OLED legend.

Config is loaded from /config.json if present (managed by the companion web app over
USB serial), otherwise the built-in DEFAULT_MODES below are used.

Serial protocol (over usb_cdc.data, enabled in boot.py), newline-delimited:
    PING            -> PONG MACROPAD-CTL v<FIRMWARE_VERSION>  (app compares version)
    GET             -> CONFIG <json>
    SET <json>      -> OK saved | OK ram | ERR <reason>

Required libraries in CIRCUITPY/lib:
  adafruit_macropad, adafruit_debouncer, adafruit_simple_text_display,
  neopixel, adafruit_display_text/, adafruit_hid/, adafruit_ticks
"""

import json
import time

import displayio
import terminalio
import usb_cdc
import usb_hid
from adafruit_display_text import label
from adafruit_hid.mouse import Mouse
from adafruit_macropad import MacroPad

# ---------------------------------------------------------------------------
# Tunables
# ---------------------------------------------------------------------------
BRIGHTNESS = 0.30          # NeoPixel base brightness (0.0 - 1.0)
ACCENT = (255, 255, 255)   # color a key flashes while held
DIAL_CYCLES_MODES = True   # turn the dial to switch modes (press also switches)
ENCODER_ARROWS = False     # only used if DIAL_CYCLES_MODES is False: turn sends Up/Down arrows
MODE_TONE = True           # short beep when switching modes
KEY_CLICK = False          # tiny beep on each keypress (adds a little latency)
STEP_DELAY = 0.06          # pause between macro steps (lets palettes/terminals catch up)
CONFIG_FILE = "/config.json"
FIRMWARE_VERSION = 2       # bump when code.py gains features; web app compares this

macropad = MacroPad()      # rotation=0: USB up, encoder top-right
KC = macropad.Keycode
serial = usb_cdc.data      # None if boot.py did not enable it

# Mouse HID (for mapping mouse-button actions, e.g. middle-click). Guarded in case
# the mouse device is not exposed; mouse steps become no-ops if so.
try:
    mouse = Mouse(usb_hid.devices)
except Exception:  # noqa: BLE001 - any HID setup failure just disables mouse steps
    mouse = None
MOUSE_BUTTONS = {
    "left": Mouse.LEFT_BUTTON,
    "middle": Mouse.MIDDLE_BUTTON,
    "right": Mouse.RIGHT_BUTTON,
}

# ---------------------------------------------------------------------------
# Built-in default keymap (used when /config.json is absent)
# ---------------------------------------------------------------------------
# Each key's "action" is a list of STEPS run in order. A step is either:
#   * a tuple/list of Keycodes pressed together then released -> (KC.GUI, KC.N)
#   * a string typed out verbatim                             -> "/clear"
#   * a dict describing a mouse click                         -> {"mouse": "middle"}
# Sequences are multiple steps; text macros that submit = ["/clear", (KC.RETURN,)]

CURSOR = {
    "name": "CURSOR",
    "color": (40, 90, 255),
    "keys": [
        {"label": "Agent",   "action": [(KC.GUI, KC.I)]},
        {"label": "Ask",     "action": [(KC.GUI, KC.L)]},
        {"label": "NewChat", "action": [(KC.GUI, KC.N)]},
        {"label": "Send",    "action": [(KC.RETURN,)]},
        {"label": "Accept",  "action": [(KC.GUI, KC.RETURN)]},
        {"label": "Stop",    "action": [(KC.ESCAPE,)]},
        {"label": "Sidebar", "action": [(KC.GUI, KC.B)]},
        {"label": "Term",    "action": [(KC.CONTROL, KC.GRAVE_ACCENT)]},
        {"label": "CmdPal",  "action": [(KC.GUI, KC.SHIFT, KC.P)]},
        {"label": "GoTo",    "action": [(KC.GUI, KC.P)]},
        {"label": "Undo",    "action": [(KC.GUI, KC.Z)]},
        {"label": "Save",    "action": [(KC.GUI, KC.S)]},
    ],
}

CODEX = {
    "name": "CODEX",
    "color": (0, 200, 160),
    "keys": [
        {"label": "NewChat",  "action": [(KC.GUI, KC.N)]},
        {"label": "Send",     "action": [(KC.RETURN,)]},
        {"label": "Search",   "action": [(KC.GUI, KC.K)]},
        {"label": "Approve",  "action": [(KC.RETURN,)]},
        {"label": "Decline",  "action": [(KC.ESCAPE,)]},
        {"label": "Stop",     "action": [(KC.ESCAPE,)]},
        {"label": "Sidebar",  "action": [(KC.GUI, KC.SHIFT, KC.S)]},
        {"label": "NewWin",   "action": [(KC.GUI, KC.SHIFT, KC.N)]},
        {"label": "Back",     "action": [(KC.GUI, KC.LEFT_BRACKET)]},
        {"label": "Fwd",      "action": [(KC.GUI, KC.RIGHT_BRACKET)]},
        {"label": "Copy",     "action": [(KC.GUI, KC.C)]},
        {"label": "Settings", "action": [(KC.GUI, KC.COMMA)]},
    ],
}

CLAUDE = {
    "name": "CLAUDE CODE",
    "color": (255, 130, 0),
    "keys": [
        {"label": "Mode",     "action": [(KC.SHIFT, KC.TAB)]},
        {"label": "Stop",     "action": [(KC.ESCAPE,)]},
        {"label": "EditPrev", "action": [(KC.ESCAPE,), (KC.ESCAPE,)]},
        {"label": "Send",     "action": [(KC.RETURN,)]},
        {"label": "Approve",  "action": [(KC.RETURN,)]},
        {"label": "Decline",  "action": [(KC.ESCAPE,)]},
        {"label": "History",  "action": [(KC.UP_ARROW,)]},
        {"label": "/clear",   "action": ["/clear", (KC.RETURN,)]},
        {"label": "/compact", "action": ["/compact", (KC.RETURN,)]},
        {"label": "/model",   "action": ["/model", (KC.RETURN,)]},
        {"label": "Verbose",  "action": [(KC.CONTROL, KC.R)]},
        {"label": "Quit",     "action": [(KC.CONTROL, KC.C)]},
    ],
}

SALESFORCE = {
    "name": "SALESFORCE",
    "color": (0, 180, 220),
    "keys": [
        {"label": "Term",    "action": [(KC.CONTROL, KC.GRAVE_ACCENT)]},
        {"label": "Manifst", "action": ["sf project generate manifest --from-org "]},
        {"label": "Deploy",  "action": ["sf project deploy start", (KC.RETURN,)]},
        {"label": "Retrve",  "action": ["sf project retrieve start", (KC.RETURN,)]},
        {"label": "DplyMan", "action": ["sf project deploy start --manifest manifest/package.xml", (KC.RETURN,)]},
        {"label": "Tests",   "action": ["sf apex run test --code-coverage --result-format human", (KC.RETURN,)]},
        {"label": "OpenOrg", "action": ["sf org open", (KC.RETURN,)]},
        {"label": "Login",   "action": ["sf org login web", (KC.RETURN,)]},
        {"label": "OrgList", "action": ["sf org list", (KC.RETURN,)]},
        {"label": "NewProj", "action": ["sf project generate --name "]},
        {"label": "OrgBrsr", "action": [(KC.GUI, KC.SHIFT, KC.P), "SFDX: Open Org Browser", (KC.RETURN,)]},
        {"label": "Cancel",  "action": [(KC.CONTROL, KC.C)]},
    ],
}

# Wispr Flow is voice dictation. Both dictation keys use Wispr's default Middle Click
# "Push to talk" binding, so nothing needs re-binding in Wispr:
#   * "Dictate" -> "latch": True  -> tap once to start (button held), tap again to stop
#                                    (emulates hands-free using the push-to-talk binding)
#   * "Talk"    -> "hold": True   -> momentary push-to-talk, held only while pressed
# Note: Wispr's fn-based shortcuts (Hands-free fn+Space, Command fn+Ctrl) can't be sent
# over USB HID, so rebind them to a normal hotkey in Wispr if you want dedicated keys.
WISPR = {
    "name": "WISPR FLOW",
    "color": (170, 70, 230),
    "keys": [
        {"label": "Dictate",  "action": [{"mouse": "middle"}], "latch": True},
        {"label": "Talk",     "action": [{"mouse": "middle"}], "hold": True},
        {"label": "Undo",     "action": [(KC.GUI, KC.Z)]},
        {"label": "Redo",     "action": [(KC.GUI, KC.SHIFT, KC.Z)]},
        {"label": "Copy",     "action": [(KC.GUI, KC.C)]},
        {"label": "Paste",    "action": [(KC.GUI, KC.V)]},
        {"label": "Cut",      "action": [(KC.GUI, KC.X)]},
        {"label": "SelAll",   "action": [(KC.GUI, KC.A)]},
        {"label": "Enter",    "action": [(KC.RETURN,)]},
        {"label": "Del",      "action": [(KC.BACKSPACE,)]},
        {"label": "Space",    "action": [(KC.SPACEBAR,)]},
        {"label": "Esc",      "action": [(KC.ESCAPE,)]},
    ],
}

# Figma is a GUI app: single-letter tool shortcuts + common object/view actions.
FIGMA = {
    "name": "FIGMA",
    "color": (235, 70, 150),
    "keys": [
        {"label": "Move",    "action": [(KC.V,)]},
        {"label": "Frame",   "action": [(KC.F,)]},
        {"label": "Rect",    "action": [(KC.R,)]},
        {"label": "Ellipse", "action": [(KC.O,)]},
        {"label": "Text",    "action": [(KC.T,)]},
        {"label": "Pen",     "action": [(KC.P,)]},
        {"label": "Comment", "action": [(KC.C,)]},
        {"label": "Compnt",  "action": [(KC.GUI, KC.ALT, KC.K)]},
        {"label": "Group",   "action": [(KC.GUI, KC.G)]},
        {"label": "Ungroup", "action": [(KC.GUI, KC.SHIFT, KC.G)]},
        {"label": "ZoomFit", "action": [(KC.SHIFT, KC.ONE)]},
        {"label": "DevMode", "action": [(KC.SHIFT, KC.D)]},
    ],
}

# Google Meet (browser-based, macOS). Meet's own shortcuts cover mic/camera/hand/
# push-to-talk; the rest are browser + tab controls for managing the call tab.
MEET = {
    "name": "GOOGLE MEET",
    "color": (0, 170, 90),
    "keys": [
        {"label": "Mic",     "action": [(KC.GUI, KC.D)]},
        {"label": "Camera",  "action": [(KC.GUI, KC.E)]},
        {"label": "Talk",    "action": [(KC.SPACEBAR,)], "hold": True},
        {"label": "Hand",    "action": [(KC.CONTROL, KC.GUI, KC.H)]},
        {"label": "FullScr", "action": [(KC.CONTROL, KC.GUI, KC.F)]},
        {"label": "PrevTab", "action": [(KC.GUI, KC.SHIFT, KC.LEFT_BRACKET)]},
        {"label": "NextTab", "action": [(KC.GUI, KC.SHIFT, KC.RIGHT_BRACKET)]},
        {"label": "Reload",  "action": [(KC.GUI, KC.R)]},
        {"label": "Address", "action": [(KC.GUI, KC.L)]},
        {"label": "Copy",    "action": [(KC.GUI, KC.C)]},
        {"label": "Paste",   "action": [(KC.GUI, KC.V)]},
        {"label": "Close",   "action": [(KC.GUI, KC.W)]},
    ],
}

# Zoom (macOS desktop app) meeting controls. "Talk" holds Space to unmute while held.
ZOOM = {
    "name": "ZOOM",
    "color": (45, 140, 255),
    "keys": [
        {"label": "Mute",   "action": [(KC.GUI, KC.SHIFT, KC.A)]},
        {"label": "Video",  "action": [(KC.GUI, KC.SHIFT, KC.V)]},
        {"label": "Talk",   "action": [(KC.SPACEBAR,)], "hold": True},
        {"label": "Share",  "action": [(KC.GUI, KC.SHIFT, KC.S)]},
        {"label": "Pause",  "action": [(KC.GUI, KC.SHIFT, KC.T)]},
        {"label": "Record", "action": [(KC.GUI, KC.SHIFT, KC.R)]},
        {"label": "Hand",   "action": [(KC.ALT, KC.Y)]},
        {"label": "People", "action": [(KC.GUI, KC.U)]},
        {"label": "Chat",   "action": [(KC.GUI, KC.SHIFT, KC.H)]},
        {"label": "View",   "action": [(KC.GUI, KC.SHIFT, KC.W)]},
        {"label": "Invite", "action": [(KC.GUI, KC.I)]},
        {"label": "Leave",  "action": [(KC.GUI, KC.W)]},
    ],
}

DEFAULT_MODES = [CODEX, CURSOR, CLAUDE, SALESFORCE, WISPR, FIGMA, MEET, ZOOM]


# ---------------------------------------------------------------------------
# Config (de)serialization
# ---------------------------------------------------------------------------
def _blank_key():
    return {"label": "", "action": [], "hold": False, "latch": False}


def parse_modes(data):
    """Turn parsed JSON into a usable modes list, or return None if invalid."""
    try:
        raw_modes = data["modes"]
    except (KeyError, TypeError):
        return None
    modes = []
    for raw in raw_modes:
        keys = []
        for raw_key in raw.get("keys", []):
            action = []
            for step in raw_key.get("action", []):
                if isinstance(step, str):
                    action.append(step)
                elif isinstance(step, dict):
                    action.append(step)  # mouse step, e.g. {"mouse": "middle"}
                else:
                    action.append(tuple(step))
            keys.append({
                "label": str(raw_key.get("label", "")),
                "action": action,
                "hold": bool(raw_key.get("hold", False)),
                "latch": bool(raw_key.get("latch", False)),
            })
        keys = (keys + [_blank_key() for _ in range(12)])[:12]  # pad/truncate to 12
        color = raw.get("color", [80, 80, 80])
        modes.append({"name": str(raw.get("name", "MODE")), "color": tuple(color), "keys": keys})
    return modes or None


def serialize_modes():
    out = {"modes": []}
    for mode in MODES:
        keys = []
        for key in mode["keys"]:
            steps = [s if isinstance(s, (str, dict)) else list(s) for s in key["action"]]
            keys.append({
                "label": key["label"],
                "action": steps,
                "hold": bool(key.get("hold", False)),
                "latch": bool(key.get("latch", False)),
            })
        out["modes"].append({"name": mode["name"], "color": list(mode["color"]), "keys": keys})
    return json.dumps(out)


def load_modes():
    try:
        with open(CONFIG_FILE) as handle:
            modes = parse_modes(json.load(handle))
            if modes:
                return modes
    except (OSError, ValueError):
        pass
    return DEFAULT_MODES


MODES = load_modes()
mode_index = 0

# ---------------------------------------------------------------------------
# OLED: title bar + 3x4 legend grid
# ---------------------------------------------------------------------------
macropad.pixels.brightness = BRIGHTNESS

group = displayio.Group()

bar_bmp = displayio.Bitmap(128, 13, 1)
bar_pal = displayio.Palette(1)
bar_pal[0] = 0xFFFFFF
group.append(displayio.TileGrid(bar_bmp, pixel_shader=bar_pal, x=0, y=0))

title = label.Label(
    terminalio.FONT, text="", color=0x000000,
    anchor_point=(0.5, 0.5), anchored_position=(64, 6),
)
group.append(title)

key_labels = []
for i in range(12):
    row, col = i // 3, i % 3
    lbl = label.Label(
        terminalio.FONT, text="", color=0xFFFFFF,
        anchor_point=(0, 0.5), anchored_position=(col * 43 + 1, 17 + row * 12),
    )
    group.append(lbl)
    key_labels.append(lbl)

try:
    macropad.display.root_group = group  # CircuitPython 8+
except AttributeError:
    macropad.display.show(group)  # CircuitPython 7.x


def show_mode():
    mode = MODES[mode_index]
    title.text = mode["name"]
    for i, lbl in enumerate(key_labels):
        lbl.text = mode["keys"][i]["label"][:7]
    macropad.pixels.fill(mode["color"])


def switch_mode(delta):
    global mode_index
    release_all_latches()  # never leave a key/button stuck when leaving a mode
    mode_index = (mode_index + delta) % len(MODES)
    show_mode()
    if MODE_TONE:
        macropad.play_tone(440 + 110 * mode_index, 0.08)


def mouse_button(step):
    return MOUSE_BUTTONS.get(step.get("mouse")) if isinstance(step, dict) else None


def run_action(steps):
    for step in steps:
        if isinstance(step, str):
            macropad.keyboard_layout.write(step)
        elif isinstance(step, dict):
            btn = mouse_button(step)
            if mouse is not None and btn is not None:
                mouse.click(btn)
        else:
            macropad.keyboard.press(*step)
            macropad.keyboard.release_all()
        time.sleep(STEP_DELAY)


# --- Latching keys -----------------------------------------------------------
# A "latch" key holds its first key/mouse step DOWN on the first tap and releases
# it on the next tap (no need to keep the pad key held). latched maps a key number
# to ("kbd", keycodes) or ("mouse", button) so we can release exactly what we held.
latched = {}


def _first_holdable_step(key_def):
    for step in key_def["action"]:
        if not isinstance(step, str):
            return step
    return None


def _release_latch(key):
    kind, payload = latched.pop(key)
    if kind == "mouse":
        if mouse is not None:
            mouse.release(payload)
    else:
        macropad.keyboard.release(*payload)


def toggle_latch(key, key_def):
    """Tap on/off a held key/button. Returns True if now latched ON."""
    if key in latched:
        _release_latch(key)
        return False
    step = _first_holdable_step(key_def)
    if step is None:
        return False
    if isinstance(step, dict):
        btn = mouse_button(step)
        if mouse is None or btn is None:
            return False
        mouse.press(btn)
        latched[key] = ("mouse", btn)
    else:
        macropad.keyboard.press(*step)
        latched[key] = ("kbd", tuple(step))
    return True


def release_all_latches():
    for key in list(latched):
        _release_latch(key)


# ---------------------------------------------------------------------------
# Serial config protocol
# ---------------------------------------------------------------------------
_rx = b""
test_mode = False   # when True, keys report events instead of firing macros


def send_line(text):
    if serial is None:
        return
    data = (text + "\n").encode("utf-8")
    while data:
        written = serial.write(data)
        data = data[written:] if written else data


def apply_config(payload):
    global MODES, mode_index
    try:
        data = json.loads(payload)
    except ValueError:
        send_line("ERR badjson")
        return
    new_modes = parse_modes(data)
    if not new_modes:
        send_line("ERR nomodes")
        return
    MODES = new_modes
    if mode_index >= len(MODES):
        mode_index = 0
    show_mode()
    try:
        with open(CONFIG_FILE, "w") as handle:
            handle.write(payload)
        send_line("OK saved")
    except OSError:
        send_line("OK ram")  # filesystem read-only (KEY1 held at boot); applied in RAM only


def handle_line(raw):
    global test_mode
    try:
        text = raw.decode("utf-8").strip()
    except UnicodeError:
        return
    if text == "PING":
        send_line("PONG MACROPAD-CTL v%d" % FIRMWARE_VERSION)
    elif text == "GET":
        send_line("CONFIG " + serialize_modes())
    elif text.startswith("SET "):
        apply_config(text[4:])
    elif text == "TEST 1":
        test_mode = True
        send_line("TESTON " + str(mode_index))
    elif text == "TEST 0":
        test_mode = False
        show_mode()  # restore LED colors for the active mode
        send_line("TESTOFF")


def poll_serial():
    global _rx
    if serial is None or not serial.in_waiting:
        return
    _rx += serial.read(serial.in_waiting)
    while b"\n" in _rx:
        line, _rx = _rx.split(b"\n", 1)
        handle_line(line)


if serial is not None:
    serial.timeout = 0

# ---------------------------------------------------------------------------
# Main loop
# ---------------------------------------------------------------------------
show_mode()
last_position = macropad.encoder

while True:
    poll_serial()

    macropad.encoder_switch_debounced.update()
    if test_mode:
        if macropad.encoder_switch_debounced.pressed:
            send_line("EV S 1")
        if macropad.encoder_switch_debounced.released:
            send_line("EV S 0")
    elif macropad.encoder_switch_debounced.pressed:
        switch_mode(1)  # pressing the dial also cycles forward

    position = macropad.encoder
    if position != last_position:
        delta = position - last_position
        if test_mode:
            # Cycle modes just like normal use so the OLED/LEDs update, and tell
            # the web app the new mode so its on-screen mockup follows along.
            switch_mode(delta)
            send_line("EV M " + str(mode_index))
            send_line("EV E " + str(position))
        elif DIAL_CYCLES_MODES:
            switch_mode(delta)
        elif ENCODER_ARROWS:
            keycode = KC.DOWN_ARROW if delta > 0 else KC.UP_ARROW
            for _ in range(min(abs(delta), 8)):
                macropad.keyboard.press(keycode)
                macropad.keyboard.release_all()
        last_position = position

    event = macropad.keys.events.get()
    if event:
        mode = MODES[mode_index]
        key = event.key_number
        if test_mode:
            if event.pressed:
                macropad.pixels[key] = ACCENT
                send_line("EV P " + str(key))
            else:
                macropad.pixels[key] = mode["color"]
                send_line("EV R " + str(key))
        elif event.pressed:
            if KEY_CLICK:
                macropad.play_tone(660, 0.01)
            key_def = mode["keys"][key]
            if key_def.get("latch"):
                # Tap to toggle: hold the key/button down until the next tap. Keep the
                # LED lit while latched ON so it's obvious the button is being held.
                macropad.pixels[key] = ACCENT if toggle_latch(key, key_def) else mode["color"]
            elif key_def.get("hold"):
                # Push-to-talk: hold the first key combo (or mouse button) down until release.
                macropad.pixels[key] = ACCENT
                step = _first_holdable_step(key_def)
                if isinstance(step, dict):
                    btn = mouse_button(step)
                    if mouse is not None and btn is not None:
                        mouse.press(btn)
                elif step is not None:
                    macropad.keyboard.press(*step)
            else:
                macropad.pixels[key] = ACCENT
                run_action(key_def["action"])
        else:
            key_def = mode["keys"][key]
            if key_def.get("latch"):
                pass  # LED reflects latch state; leave it as set on press
            elif key_def.get("hold"):
                macropad.pixels[key] = mode["color"]
                macropad.keyboard.release_all()
                if mouse is not None:
                    mouse.release_all()
            else:
                macropad.pixels[key] = mode["color"]
