#!/usr/bin/env python3
"""
KB2040 Hidden Signer host client (WebUSB/vendor interface via pyusb).

Protocol:
  4-byte big-endian length + UTF-8 JSON body.

Default target:
  VID:PID 239a:8104

Examples:
  python3 examples/kb2040_hidden_signer_client.py status
  python3 examples/kb2040_hidden_signer_client.py set-mnemonic --mnemonic "abandon ... about"
  python3 examples/kb2040_hidden_signer_client.py get-public-key
  python3 examples/kb2040_hidden_signer_client.py sign-event --event '{"kind":1,"content":"hello"}'
"""

import argparse
import json
import struct
import sys
import time

try:
    import usb.core
    import usb.util
except ImportError:
    print("Missing pyusb. Install with: pip install pyusb", file=sys.stderr)
    sys.exit(2)


def find_vendor_interface(dev):
    cfg = dev.get_active_configuration()
    for itf in cfg:
        if itf.bInterfaceClass == 0xFF:
            ep_out = None
            ep_in = None
            for ep in itf:
                direction = usb.util.endpoint_direction(ep.bEndpointAddress)
                if direction == usb.util.ENDPOINT_OUT:
                    ep_out = ep
                elif direction == usb.util.ENDPOINT_IN:
                    ep_in = ep
            if ep_out is not None and ep_in is not None:
                return itf, ep_out, ep_in
    return None, None, None


def open_device(vid: int, pid: int):
    dev = usb.core.find(idVendor=vid, idProduct=pid)
    if dev is None:
        raise RuntimeError(f"Device {vid:04x}:{pid:04x} not found")

    try:
        dev.set_configuration()
    except usb.core.USBError:
        pass

    itf, ep_out, ep_in = find_vendor_interface(dev)
    if itf is None:
        raise RuntimeError("Vendor/WebUSB interface (class 0xFF) not found")

    itf_num = itf.bInterfaceNumber
    try:
        if dev.is_kernel_driver_active(itf_num):
            dev.detach_kernel_driver(itf_num)
    except (NotImplementedError, usb.core.USBError):
        pass

    usb.util.claim_interface(dev, itf_num)

    # WebUSB-style connect handshake (CDC SET_CONTROL_LINE_STATE / request 0x22).
    # Firmware gates IN transfers on this "connected" state.
    try:
        dev.ctrl_transfer(0x21, 0x22, 0x0001, itf_num, None)
    except usb.core.USBError:
        # Some stacks may not require/implement this; continue and let RPC I/O decide.
        pass

    return dev, itf_num, ep_out, ep_in


def close_device(dev, itf_num: int):
    try:
        usb.util.release_interface(dev, itf_num)
    except usb.core.USBError:
        pass


_rx_remainder = bytearray()


def send_frame(ep_out, payload: bytes):
    frame = struct.pack(">I", len(payload)) + payload
    ep_out.write(frame)


def recv_exact(ep_in, n: int, timeout_ms: int):
    global _rx_remainder

    out = bytearray()

    # Consume previously over-read bytes first.
    if _rx_remainder:
        take = min(n, len(_rx_remainder))
        out.extend(_rx_remainder[:take])
        del _rx_remainder[:take]

    deadline = time.time() + max(1.0, timeout_ms / 1000.0)
    packet_size = max(64, int(getattr(ep_in, "wMaxPacketSize", 64)))

    while len(out) < n and time.time() < deadline:
        chunk = ep_in.read(packet_size, timeout=timeout_ms)
        out.extend(bytes(chunk))

    if len(out) < n:
        raise TimeoutError(f"Timeout while reading {n} bytes (got {len(out)})")

    if len(out) > n:
        _rx_remainder.extend(out[n:])
        del out[n:]

    return bytes(out)


def recv_frame(ep_in, timeout_ms: int):
    hdr = recv_exact(ep_in, 4, timeout_ms)
    body_len = struct.unpack(">I", hdr)[0]
    if body_len > 64 * 1024:
        raise ValueError(f"Refusing oversized frame: {body_len}")
    body = recv_exact(ep_in, body_len, timeout_ms)
    return body


def rpc(dev_args, method: str, params: dict):
    global _rx_remainder

    dev, itf_num, ep_out, ep_in = open_device(dev_args.vid, dev_args.pid)
    try:
        _rx_remainder.clear()
        req = {"method": method}
        req.update(params)
        payload = json.dumps(req, separators=(",", ":")).encode("utf-8")
        send_frame(ep_out, payload)
        resp = recv_frame(ep_in, dev_args.timeout_ms)
        return json.loads(resp.decode("utf-8", errors="replace"))
    finally:
        close_device(dev, itf_num)


def cmd_status(args):
    print(json.dumps(rpc(args, "get_status", {}), indent=2))


def cmd_set_mnemonic(args):
    print(json.dumps(rpc(args, "set_mnemonic", {"mnemonic": args.mnemonic}), indent=2))


def cmd_set_auto_approve(args):
    value = args.value.lower() in ("1", "true", "yes", "on")
    print(json.dumps(rpc(args, "set_auto_approve", {"value": value}), indent=2))


def cmd_get_public_key(args):
    print(json.dumps(rpc(args, "get_public_key", {}), indent=2))


def cmd_sign_event(args):
    event = json.loads(args.event)
    print(json.dumps(rpc(args, "sign_event", {"event": event}), indent=2))


def build_parser():
    p = argparse.ArgumentParser()
    p.add_argument("--vid", type=lambda s: int(s, 16), default=0x239A, help="USB VID in hex (default: 239A)")
    p.add_argument("--pid", type=lambda s: int(s, 16), default=0x8104, help="USB PID in hex (default: 8104)")
    p.add_argument("--timeout-ms", type=int, default=3000, help="USB read timeout in ms")

    sub = p.add_subparsers(dest="cmd", required=True)

    s = sub.add_parser("status")
    s.set_defaults(func=cmd_status)

    s = sub.add_parser("set-mnemonic")
    s.add_argument("--mnemonic", required=True)
    s.set_defaults(func=cmd_set_mnemonic)

    s = sub.add_parser("set-auto-approve")
    s.add_argument("--value", required=True, help="true|false")
    s.set_defaults(func=cmd_set_auto_approve)

    s = sub.add_parser("get-public-key")
    s.set_defaults(func=cmd_get_public_key)

    s = sub.add_parser("sign-event")
    s.add_argument("--event", required=True, help="JSON object string")
    s.set_defaults(func=cmd_sign_event)

    return p


def main():
    parser = build_parser()
    args = parser.parse_args()
    try:
        args.func(args)
        return 0
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())