#!/usr/bin/env python3
"""Capture Richmond remote presses with a CC1101 on a CJMCU-232H (FT232H) board, no ESP32.

The FT232H talks SPI to the CC1101 to set it up (433.92 MHz, OOK, asynchronous receive, Flipper's
AM650 settings), then samples the CC1101's demodulated data at 50 kHz: with CSN high, the SO pin
carries GDO1 (IOCFG1 = 0x0D), and an SPI read on a second, unconnected chip select (AD4) clocks it
in, one bit per sample. Each press is shown as it arrives; the timings go to an ESPHome-style log
and keeloq-check.py tests them against the public key list at the end.

Wiring (CC1101 header pin -> CJMCU-232H pin), 6 jumpers.
The CC1101's header prints only 1, 2, 7 and 8; pin 1 is the square pad, odd pins on the board edge:
  1 GND  -> GND on the 4-pin header    2 VCC -> +3.3V on the 4-pin header (never +5V)
  4 CSN  -> AD3    5 SCK -> AD0    6 MOSI -> AD1    7 MISO/GDO1 -> AD2
  3 GDO0, 8 GDO2, AD4: not connected

Usage, in this folder, with pyftdi in .venv:
       .venv/bin/python -I keeloq-capture-ft232h.py [--seconds 60]
       .venv/bin/python -I keeloq-capture-ft232h.py --self-test
"""
import argparse
import importlib.util
import queue
import random
import re
import sys
import threading
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
_spec = importlib.util.spec_from_file_location("keeloq_check", HERE / "keeloq-check.py")
kc = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(kc)

RATE = 50_000  # samples per second: 20 µs steps, about 20 per TE
CHUNK = RATE * 2 // 8  # bytes per USB read: 2 s of samples
GLITCH_US = 100  # shorter runs are noise (ESPHome rig: filter 120us)
GAP_US = 8000  # a low run this long ends a burst (ESPHome rig: idle 8ms)
LOG = Path("/tmp/keeloq-capture.log")

# CC1101 OOK asynchronous receive, Flipper's AM650 preset (650 kHz bandwidth)
REGS = {
    0x00: 0x2E,  # IOCFG2: GDO2 off (high impedance)
    0x01: 0x0D,  # IOCFG1: GDO1 = async serial data, on SO while CSN is high
    0x02: 0x2E,  # IOCFG0: GDO0 off
    0x03: 0x07,  # FIFOTHR
    0x08: 0x32,  # PKTCTRL0: asynchronous serial, infinite length
    0x0B: 0x06,  # FSCTRL1: IF 152 kHz
    0x10: 0x17,  # MDMCFG4: RX bandwidth 650 kHz
    0x11: 0x32,  # MDMCFG3
    0x12: 0x30,  # MDMCFG2: ASK/OOK, no sync word
    0x13: 0x00,  # MDMCFG1
    0x14: 0x00,  # MDMCFG0
    0x18: 0x18,  # MCSM0: calibrate on IDLE -> RX
    0x19: 0x18,  # FOCCFG
    0x1B: 0x07,  # AGCCTRL2
    0x1C: 0x00,  # AGCCTRL1
    0x1D: 0x91,  # AGCCTRL0
    0x20: 0xFB,  # WORCTRL
    0x21: 0xB6,  # FREND1
    0x22: 0x11,  # FREND0
}
FREQ = round(433.92e6 * 65536 / 26e6)  # 26 MHz crystal
REGS |= {0x0D: FREQ >> 16 & 0xFF, 0x0E: FREQ >> 8 & 0xFF, 0x0F: FREQ & 0xFF}
SRES, SRX, SIDLE = 0x30, 0x34, 0x36
PARTNUM, VERSION, MARCSTATE = 0x30, 0x31, 0x35


def to_durations(data, rate=RATE):
    """Samples (bytes, MSB first) -> ESPHome-style raw timings: +µs high (carrier), -µs low."""
    bits = format(int.from_bytes(data, "big"), f"0{len(data) * 8}b")
    minlen = GLITCH_US * rate // 1_000_000
    runs = []  # [level, samples]
    for m in re.finditer(r"0+|1+", bits):
        level, n = m[0][0] == "1", len(m[0])
        if runs and (n < minlen or runs[-1][0] == level):
            runs[-1][1] += n  # a glitch joins the run before it
        else:
            runs.append([level, n])
    return [round(n * 1e6 / rate) * (1 if level else -1) for level, n in runs]


def bursts(durations):
    """Split at long silences; keep bursts long enough to hold a KeeLoq frame (about 158 timings)."""
    out, cur = [], []
    for d in durations:
        if d <= -GAP_US:
            if len(cur) >= 100:
                out.append(cur)
            cur = []
        elif cur or d > 0:
            cur.append(d)
    if len(cur) >= 100:
        out.append(cur)
    return out


def setup(url):
    from pyftdi.spi import SpiController

    spi = SpiController(cs_count=2)
    spi.configure(url)
    radio = spi.get_port(cs=0, freq=1e6, mode=0)  # AD3 -> CSN
    sampler = spi.get_port(cs=1, freq=RATE, mode=0)  # AD4, not connected: CSN stays high

    def status(addr):
        return radio.exchange([0xC0 | addr, 0], duplex=True)[1]

    radio.write([SRES])
    time.sleep(0.01)
    part, version = status(PARTNUM), status(VERSION)
    print(f"🪵 CC1101 PARTNUM 0x{part:02x} VERSION 0x{version:02x}")
    if version in (0x00, 0xFF):
        raise SystemExit("🔴 no CC1101 answers: check VCC, GND, SCK, MISO, CSN\n\n🔴 FAILED")
    for addr, value in REGS.items():
        radio.write([addr, value])
    bad = {a: v for a, v in REGS.items() if radio.exchange([0x80 | a, 0], duplex=True)[1] != v}
    if bad:
        raise SystemExit(f"🔴 registers read back wrong {sorted(bad)}: check MOSI\n\n🔴 FAILED")
    radio.write([SIDLE])
    radio.write([SRX])
    time.sleep(0.01)
    state = status(MARCSTATE) & 0x1F
    if state != 0x0D:
        raise SystemExit(f"🔴 CC1101 not in receive (MARCSTATE 0x{state:02x})\n\n🔴 FAILED")
    print("🟢 CC1101 set up: 433.92 MHz OOK, receiving")
    return spi, radio, sampler


def connect(url):
    """setup(), with a plain message when the board can't be opened (rf-scan-ft232h.py uses it too)."""
    try:
        return setup(url)
    except SystemExit:
        raise
    except Exception as e:  # pyftdi / pyusb: missing board or no USB permission
        raise SystemExit(f"🔴 can't open the FT232H ({url}): {e}\n"
                         "   plugged in? udev rule for 0403:6014 installed?\n\n🔴 FAILED")


def capture(url, seconds, log):
    spi, _, sampler = connect(url)
    chunks, stop = queue.Queue(), threading.Event()

    def reader():  # keeps the USB reads back to back while the main thread decodes
        while not stop.is_set():
            chunks.put(bytes(sampler.read(CHUNK)))
        chunks.put(None)

    threading.Thread(target=reader, daemon=True).start()
    print(f"🌱 listening for {seconds} s: press the remote's spare button 4 times, about 1 s each, "
          "a few seconds apart (Ctrl+C stops early)")
    seen, end = set(), time.monotonic() + seconds
    with log.open("w") as f:
        try:
            while (data := chunks.get()) is not None:
                for b in bursts(to_durations(data)):
                    f.write("[I][remote.raw]: Received Raw: " + ", ".join(map(str, b)) + "\n")
                    for fr in kc.decode_raw(b):
                        key = (fr["serial"], fr["button"], fr["hop"])
                        if key not in seen:
                            seen.add(key)
                            print(f"🟢 press {len(seen)}: serial 0x{fr['serial']:07x} button 0x{fr['button']:x} "
                                  f"hop 0x{fr['hop']:08x} TE {fr['te']} µs")
                f.flush()
                if time.monotonic() > end:
                    stop.set()
        except KeyboardInterrupt:
            stop.set()
    spi.terminate()
    print(f"🪵 timings saved to {log}")


def self_test():
    """Render two synthetic presses, with noise and glitches, into 50 kHz samples and decode them back."""
    rng = random.Random(1)
    keys = kc.load_keys(kc.KEYFILE) if kc.KEYFILE.exists() else kc.TEST_KEYS
    mk, _, name = keys[len(keys) // 2]
    serial, button = 0x0ABCDEF, 0x2
    dev = kc.normal(button << 28 | serial, mk)
    raw = []
    for cnt in (0x0101, 0x0102):
        raw += [rng.choice([1, -1]) * rng.randint(60, 300) for _ in range(400)]  # noise between presses
        hop = kc.encrypt(button << 28 | (serial & 0x3FF) << 16 | cnt, dev)
        v = (button << 28 | serial) << 32 | hop
        frame = [410, -390] * 11 + [400, -4000]
        for n in range(66):
            frame += [400, -820] if kc.bit(v, n) else [810, -400]
        frame[-1] = -15000
        raw += [int(d * rng.uniform(0.9, 1.1)) for d in frame]
    samples = "".join(("1" if d > 0 else "0") * round(abs(d) * RATE / 1e6) for d in raw)
    samples = list(samples)
    for i in rng.sample(range(len(samples)), 300):  # single-sample glitches
        samples[i] = "1" if samples[i] == "0" else "0"
    samples = "".join(samples)
    samples += "0" * (-len(samples) % 8)
    data = int(samples, 2).to_bytes(len(samples) // 8, "big")
    frames = [fr for b in bursts(to_durations(data)) for fr in kc.decode_raw(b)]
    assert [(f["serial"], f["button"]) for f in frames] == [(serial, button)] * 2, f"decode: {frames}"
    assert kc.check(frames, keys)[0], "no match"
    print(f"🟢 two noisy presses sampled at {RATE // 1000} kHz decode back and match '{name}'")
    noise = bytes(rng.getrandbits(8) for _ in range(CHUNK))
    t = time.monotonic()
    for b in bursts(to_durations(noise)):
        kc.decode_raw(b)
    took = time.monotonic() - t
    assert took < 2, f"decoding 2 s of noise took {took:.1f} s"
    print(f"🟢 2 s of pure noise decodes in {took:.2f} s, keeps up")
    print("\n🟢 SUCCESS")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--seconds", type=int, default=60)
    ap.add_argument("--url", default="ftdi://ftdi:232h/1")
    ap.add_argument("--log", type=Path, default=LOG)
    ap.add_argument("--self-test", action="store_true")
    args = ap.parse_args()
    if args.self_test:
        self_test()
        sys.exit(0)
    capture(args.url, args.seconds, args.log)
    if not kc.KEYFILE.exists():
        print(f"\n🟢 SUCCESS: presses captured; no {kc.KEYFILE.name} next to the script, so no key check")
        sys.exit(0)
    ok = kc.report(kc.read_log(args.log), kc.load_keys(kc.KEYFILE))
    print("\n🟢 SUCCESS: a known key fits" if ok else "\n🔴 FAILED: no known key fits")
    sys.exit(0 if ok else 1)
