"""
parse_hci.py - Parse Android btsnoop_hci.log to extract KAIYU BLE commands.

Reads a btsnoop v1 binary log and extracts:
  - BLE service / characteristic UUIDs
  - Every ATT Write Command payload (the actual LED control bytes)

Usage:
    python sniffer/parse_hci.py
    python sniffer/parse_hci.py D:\\Claude\\bugreport\\btsnoop_hci.log
    python sniffer/parse_hci.py --out results.json
"""

import struct
import sys
import json
import os
import argparse
from dataclasses import dataclass, field
from typing import Optional

BTSNOOP_MAGIC        = b"btsnoop\x00"
ATT_WRITE_CMD        = 0x52
ATT_WRITE_REQ        = 0x12
ATT_READ_BY_TYPE_RSP = 0x09
ATT_CID              = 0x0004


@dataclass
class BLEWrite:
    timestamp_us: int
    direction: str
    handle: int
    data: bytes
    data_hex: str = field(init=False)

    def __post_init__(self):
        self.data_hex = self.data.hex()


@dataclass
class ParseResult:
    writes: list = field(default_factory=list)
    handles_seen: set = field(default_factory=set)
    uuid_map: dict = field(default_factory=dict)


def parse_btsnoop(filepath: str) -> ParseResult:
    result = ParseResult()
    with open(filepath, "rb") as f:
        magic = f.read(8)
        if magic != BTSNOOP_MAGIC:
            raise ValueError(
                f"Not a btsnoop binary file (magic={magic!r}). "
                "Need the raw btsnoop_hci.log binary, not the bugreport .txt."
            )
        version, datalink = struct.unpack(">II", f.read(8))
        print(f"BTSnoop v{version}, datalink type={datalink}")

        packet_num = 0
        while True:
            hdr = f.read(24)
            if len(hdr) < 24:
                break
            orig_len, incl_len, flags, drops, timestamp_us = struct.unpack(">IIIIq", hdr)
            data = f.read(incl_len)
            packet_num += 1

            direction = "RX" if (flags & 1) else "TX"
            if len(data) < 5 or data[0] != 0x02:
                continue

            payload = data[5:]
            if len(payload) < 4:
                continue

            l2cap_len, l2cap_cid = struct.unpack("<HH", payload[:4])
            l2cap_payload = payload[4:]
            if l2cap_cid != ATT_CID or len(l2cap_payload) < 1:
                continue

            att_opcode = l2cap_payload[0]
            att_data   = l2cap_payload[1:]

            if att_opcode in (ATT_WRITE_CMD, ATT_WRITE_REQ) and direction == "TX":
                if len(att_data) < 2:
                    continue
                att_handle = struct.unpack("<H", att_data[:2])[0]
                write_data = att_data[2:]
                result.writes.append(BLEWrite(timestamp_us, direction, att_handle, write_data))
                result.handles_seen.add(att_handle)

            elif att_opcode == ATT_READ_BY_TYPE_RSP and direction == "RX":
                if not att_data:
                    continue
                item_len = att_data[0]
                items = att_data[1:]
                while len(items) >= item_len and item_len >= 7:
                    item  = items[:item_len]
                    items = items[item_len:]
                    value_handle = struct.unpack("<H", item[3:5])[0]
                    uuid_bytes   = item[5:]
                    if len(uuid_bytes) == 2:
                        uuid = f"0000{uuid_bytes[::-1].hex()}-0000-1000-8000-00805f9b34fb"
                    elif len(uuid_bytes) == 16:
                        b = uuid_bytes[::-1]
                        uuid = (f"{b[0:4].hex()}-{b[4:6].hex()}-"
                                f"{b[6:8].hex()}-{b[8:10].hex()}-{b[10:16].hex()}")
                    else:
                        uuid = uuid_bytes.hex()
                    result.uuid_map[value_handle] = uuid

    print(f"Parsed {packet_num} HCI packets.")
    return result


def try_decode(data: bytes) -> Optional[str]:
    if not data:
        return None
    if len(data) >= 7 and data[0] == 0x7e and data[-1] == 0xef:
        cmd = data[2]
        if cmd == 0x04:
            return f"7E -> Power {'ON' if data[7] != 0 else 'OFF'}"
        elif cmd == 0x05 and len(data) >= 8:
            return f"7E -> RGB({data[4]},{data[5]},{data[6]})"
        elif cmd == 0x01:
            return f"7E -> Brightness {data[3]}%"
        elif cmd == 0x03:
            return f"7E -> Effect 0x{data[3]:02x}"
        return f"7E -> cmd=0x{cmd:02x}"
    if data[0] == 0x56 and len(data) >= 7:
        if data[5] == 0xf0 and data[6] == 0xaa:
            return f"56 -> RGB({data[1]},{data[2]},{data[3]})"
        elif data[5] == 0x0f and data[6] == 0xaa:
            return f"56 -> White level={data[4]}"
    if len(data) == 3 and data[0] == 0xcc:
        return f"56 -> Power {'ON' if data[1] == 0x23 else 'OFF'}"
    return None


def analyse(result: ParseResult):
    print(f"\n{'='*60}")
    print(f"  PARSE RESULTS")
    print(f"{'='*60}")
    print(f"Write packets : {len(result.writes)}")
    print(f"Handles used  : {[f'0x{h:04x}' for h in sorted(result.handles_seen)]}")

    if result.uuid_map:
        print(f"\nCharacteristic UUIDs:")
        for handle, uuid in sorted(result.uuid_map.items()):
            marker = " <- WRITE TARGET" if handle in result.handles_seen else ""
            print(f"  0x{handle:04x} : {uuid}{marker}")

    if not result.writes:
        print("\n[!] No ATT writes found.")
        print("    Ensure HCI Snoop Log was enabled BEFORE using KAIYU app.")
        print("    Steps: enable it -> toggle BT off/on -> use app -> re-pull log.")
        return

    print(f"\n{'─'*60}")
    print(f"  ALL WRITE COMMANDS")
    print(f"{'─'*60}")
    prev_ts = result.writes[0].timestamp_us
    for i, w in enumerate(result.writes):
        delta_ms = (w.timestamp_us - prev_ts) / 1000
        prev_ts  = w.timestamp_us
        decoded  = try_decode(w.data) or ""
        print(f"[{i+1:>3}] +{delta_ms:>8.1f}ms  0x{w.handle:04x}  {w.data_hex:<30}  {decoded}")

    print(f"\n{'─'*60}")
    print(f"  UNIQUE PAYLOADS")
    print(f"{'─'*60}")
    seen = {}
    for w in result.writes:
        seen.setdefault(w.data_hex, w)
    for hex_str, w in seen.items():
        print(f"  0x{w.handle:04x}  {hex_str:<30}  {try_decode(w.data) or 'unknown format'}")

    most_used = max(result.handles_seen,
                    key=lambda h: sum(1 for w in result.writes if w.handle == h))
    uuid = result.uuid_map.get(most_used, "UNKNOWN - note the handle and check in nRF Connect")

    print(f"\n{'='*60}")
    print(f"  IDENTIFIED WRITE CHARACTERISTIC")
    print(f"{'='*60}")
    print(f"  Handle : 0x{most_used:04x}")
    print(f"  UUID   : {uuid}")
    print(f"\n  Add to controller/config.py:")
    print(f'  WRITE_CHARACTERISTIC = "{uuid}"')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("logfile", nargs="?")
    parser.add_argument("--out", "-o")
    args = parser.parse_args()

    candidates = [
        args.logfile,
        r"D:\Claude\bugreport\btsnoop_hci.log",
        r"D:\HereMyHope\btsnoop.log",
        "btsnoop_hci.log",
    ]
    logfile = next((p for p in candidates if p and os.path.exists(p)), None)
    if not logfile:
        print("No btsnoop log found. Run pull_bt_log.bat first.")
        sys.exit(1)

    print(f"Parsing : {logfile}  ({os.path.getsize(logfile):,} bytes)")
    try:
        result = parse_btsnoop(logfile)
    except ValueError as e:
        print(f"\n[ERROR] {e}")
        sys.exit(1)

    analyse(result)

    if args.out:
        with open(args.out, "w") as f:
            json.dump({
                "writes": [{"handle": f"0x{w.handle:04x}", "hex": w.data_hex,
                             "decoded": try_decode(w.data)} for w in result.writes],
                "uuid_map": {f"0x{h:04x}": u for h, u in result.uuid_map.items()},
            }, f, indent=2)
        print(f"\nSaved to {args.out}")


if __name__ == "__main__":
    main()