#!/usr/bin/env python3
"""
XBLK pattern library serializer for xblink.

Converts pattern definitions to the XBLK binary format and writes them
to NTAG5 EEPROM blocks 256+ (upper 1K) via ntag5sensor ISO15693 commands.

Usage:
    # Serialize built-in patterns to binary file:
    python xblk_serialize.py --output patterns.bin

    # Write to NTAG5 via PCSC reader:
    python xblk_serialize.py --write

    # Load from JSON file:
    python xblk_serialize.py --json patterns.json --write
"""

import argparse
import json
import struct
import sys
import os

# XBLK format constants (must match src/pattern/mod.rs)
XBLK_MAGIC = b"XBLK"
XBLK_VERSION = 0x01
HEADER_SIZE = 16
PATTERN_ENTRY_SIZE = 112
MAX_PATTERNS = 9
MAX_COMMANDS_PER_ENGINE = 16

# EEPROM block offset for pattern library (upper 1K)
LIBRARY_BASE_BLOCK = 256

# LED_MAP encoding helpers
LED_MAP_LOOKUP = {
    "direct": 0b00, "i2c": 0b00,
    "engine1": 0b01, "e1": 0b01,
    "engine2": 0b10, "e2": 0b10,
    "engine3": 0b11, "e3": 0b11,
}

LED_MODE_LOOKUP = {"rgbw": 0x00, "mono3": 0x01}


def crc16(data: bytes) -> int:
    """CRC-16/CCITT-FALSE."""
    crc = 0xFFFF
    for b in data:
        crc ^= b << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = (crc << 1) ^ 0x1021
            else:
                crc <<= 1
            crc &= 0xFFFF
    return crc


def encode_led_map(mapping: dict) -> int:
    """Encode {"b": "engine1", "g": "engine1", ...} to LP5562 LED_MAP register byte."""
    b = LED_MAP_LOOKUP.get(mapping.get("b", "direct"), 0)
    g = LED_MAP_LOOKUP.get(mapping.get("g", "direct"), 0)
    r = LED_MAP_LOOKUP.get(mapping.get("r", "direct"), 0)
    w = LED_MAP_LOOKUP.get(mapping.get("w", "direct"), 0)
    return b | (g << 2) | (r << 4) | (w << 6)


def encode_pattern(pat: dict) -> bytes:
    """Encode a single pattern dict to PATTERN_ENTRY_SIZE bytes."""
    engines = pat.get("engines", [[], [], []])
    while len(engines) < 3:
        engines.append([])

    engine_count = sum(1 for e in engines if len(e) > 0)
    led_map_reg = encode_led_map(pat.get("led_map", {}))
    direct_pwm = pat.get("direct_pwm", [0, 0, 0, 0])
    while len(direct_pwm) < 4:
        direct_pwm.append(0)

    buf = bytearray(PATTERN_ENTRY_SIZE)
    buf[0] = engine_count
    buf[1] = led_map_reg
    buf[2:6] = bytes(direct_pwm[:4])

    for eng_idx, cmds in enumerate(engines[:3]):
        if len(cmds) > MAX_COMMANDS_PER_ENGINE:
            raise ValueError(f"Engine {eng_idx+1} has {len(cmds)} commands (max {MAX_COMMANDS_PER_ENGINE})")
        base = 6 + eng_idx * 34  # 2 bytes count + 32 bytes commands
        struct.pack_into(">H", buf, base, len(cmds))
        for i, cmd in enumerate(cmds):
            struct.pack_into(">H", buf, base + 2 + i * 2, cmd & 0xFFFF)

    return bytes(buf)


def encode_library(config: dict) -> bytes:
    """Encode a full XBLK library (header + patterns) to bytes."""
    patterns = config.get("patterns", [])
    if len(patterns) == 0:
        raise ValueError("No patterns defined")
    if len(patterns) > MAX_PATTERNS:
        raise ValueError(f"Too many patterns: {len(patterns)} (max {MAX_PATTERNS})")

    current = config.get("current", 20)
    mode = LED_MODE_LOOKUP.get(config.get("mode", "rgbw"), 0x00)
    active = config.get("active", 0) % len(patterns)

    # Build header (without CRC)
    header = bytearray(HEADER_SIZE)
    header[0:4] = XBLK_MAGIC
    header[4] = XBLK_VERSION
    header[5] = len(patterns)
    header[6] = active
    header[7] = current & 0xFF
    header[8] = mode
    # bytes 9-13 reserved
    # bytes 14-15 CRC (filled below)

    # Build pattern data
    pat_data = b""
    for pat in patterns:
        pat_data += encode_pattern(pat)

    # Compute CRC over header (bytes 0-13) + all pattern data
    crc = crc16(bytes(header[:14]) + pat_data)
    struct.pack_into(">H", header, 14, crc)

    return bytes(header) + pat_data


# ---------------------------------------------------------------------------
# Built-in patterns (matching src/pattern/mod.rs)
# ---------------------------------------------------------------------------

# LP5562 EngineCommand helpers (matching lp5562.rs encoding)
def ramp_wait(prescale_slow: bool, step_time: int, up: bool, increment: int) -> int:
    prescale_bit = 0x4000 if prescale_slow else 0
    sign_bit = 0 if up else 0x0080
    return prescale_bit | ((step_time & 0x3F) << 8) | sign_bit | (increment & 0x7F)

def wait(prescale_slow: bool, step_time: int) -> int:
    prescale_bit = 0x4000 if prescale_slow else 0
    return prescale_bit | ((step_time & 0x3F) << 8)

def set_pwm(value: int) -> int:
    return 0x4000 | value  # Actually: 0x40xx format
    # Wait, let me check the actual encoding...

def branch(step: int, loop_count: int) -> int:
    return 0xA000 | ((loop_count & 0x3F) << 7) | (step & 0x7F)

def trigger(wait_mask: int, send_mask: int) -> int:
    return 0xE000 | ((wait_mask & 0x07) << 7) | (send_mask & 0x07)


def builtin_patterns() -> dict:
    """Return the 5 built-in patterns as a config dict."""
    # Breathe: 1 engine, all RGB channels
    breathe_cmds = [
        ramp_wait(True, 1, True, 127),   # 0→128
        ramp_wait(True, 1, True, 127),   # 128→255
        ramp_wait(True, 1, False, 127),  # 255→127
        ramp_wait(True, 1, False, 127),  # 127→0
        wait(True, 48),                   # pause
        branch(0, 0),                     # loop
    ]

    # Heartbeat: 1 engine, double-pulse
    heartbeat_cmds = [
        0x40FF,                           # set_pwm(255)
        wait(False, 20),                  # hold ~10ms
        0x4000,                           # set_pwm(0)
        wait(False, 40),                  # gap ~20ms
        0x40FF,                           # set_pwm(255)
        wait(False, 20),                  # hold ~10ms
        0x4000,                           # set_pwm(0)
        wait(True, 63),                   # rest ~1.0s
        wait(True, 32),                   # rest ~0.5s
        branch(0, 0),
    ]

    # Slow pulse: 1 engine, very gentle
    slow_pulse_cmds = [
        ramp_wait(True, 4, True, 127),   # 0→128
        ramp_wait(True, 4, True, 127),   # 128→255
        wait(True, 32),                   # hold
        ramp_wait(True, 4, False, 127),  # 255→127
        ramp_wait(True, 4, False, 127),  # 127→0
        wait(True, 63),                   # pause
        branch(0, 0),
    ]

    # RGB cycle: 3 engines with trigger sync
    rgb_e1 = [
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, True, 127),
        trigger(0, 0b010),                # send to E2
        ramp_wait(True, 1, False, 127),
        ramp_wait(True, 1, False, 127),
        wait(True, 63),
        branch(0, 0),
    ]
    rgb_e2 = [
        trigger(0b001, 0),                # wait for E1
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, True, 127),
        trigger(0, 0b100),                # send to E3
        ramp_wait(True, 1, False, 127),
        ramp_wait(True, 1, False, 127),
        wait(True, 32),
        branch(0, 0),
    ]
    rgb_e3 = [
        trigger(0b010, 0),                # wait for E2
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, False, 127),
        ramp_wait(True, 1, False, 127),
        wait(True, 32),
        branch(0, 0),
    ]

    # Color wash: 3 engines, free-running with different periods
    wash_e1 = [
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, False, 127),
        ramp_wait(True, 1, False, 127),
        branch(0, 0),
    ]
    wash_e2 = [
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, False, 127),
        ramp_wait(True, 1, False, 127),
        wait(True, 32),
        branch(0, 0),
    ]
    wash_e3 = [
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, True, 127),
        ramp_wait(True, 1, False, 127),
        ramp_wait(True, 1, False, 127),
        wait(True, 63),
        branch(0, 0),
    ]

    single_rgb_map = {"b": "engine1", "g": "engine1", "r": "engine1", "w": "direct"}
    triple_map = {"b": "engine1", "g": "engine2", "r": "engine3", "w": "direct"}

    return {
        "current": 20,
        "mode": "rgbw",
        "active": 0,
        "patterns": [
            {"name": "breathe",    "led_map": single_rgb_map, "engines": [breathe_cmds, [], []]},
            {"name": "heartbeat",  "led_map": single_rgb_map, "engines": [heartbeat_cmds, [], []]},
            {"name": "slow_pulse", "led_map": single_rgb_map, "engines": [slow_pulse_cmds, [], []]},
            {"name": "rgb_cycle",  "led_map": triple_map,     "engines": [rgb_e1, rgb_e2, rgb_e3]},
            {"name": "color_wash", "led_map": triple_map,     "engines": [wash_e1, wash_e2, wash_e3]},
        ],
    }


def write_to_ntag5(data: bytes):
    """Write binary data to NTAG5 EEPROM blocks 256+ via ntag5sensor."""
    # Add ntag5sensor to path
    ntag5sensor_path = os.path.join(os.path.dirname(__file__), "..", "..", "ntag5sensor")
    sys.path.insert(0, ntag5sensor_path)

    from vicinity.iso15693 import ISO15693, ISO_FLAG_DATA_RATE
    from vicinity.ntag5link import Ntag5Link

    reader = ISO15693()
    chip = Ntag5Link(reader)

    print(f"Writing {len(data)} bytes to EEPROM blocks {LIBRARY_BASE_BLOCK}-{LIBRARY_BASE_BLOCK + len(data)//4 - 1}")

    # Write in 4-byte blocks
    for i in range(0, len(data), 4):
        block = LIBRARY_BASE_BLOCK + i // 4
        chunk = data[i:i+4]
        if len(chunk) < 4:
            chunk = chunk + b'\x00' * (4 - len(chunk))

        # Use ISO15693 WRITE SINGLE BLOCK (unaddressed)
        # Block address needs protocol extension for blocks > 255
        flags = ISO_FLAG_DATA_RATE | 0x08  # data rate + protocol extension
        cmd = bytes([flags, 0x21]) + struct.pack("<H", block) + chunk
        reader.transmit_iso15693(cmd, True)

        # EEPROM write cycle delay
        import time
        time.sleep(0.006)

        if (i // 4) % 10 == 0:
            print(f"  Block {block}...")

    print("Done!")


def main():
    parser = argparse.ArgumentParser(description="XBLK pattern library serializer")
    parser.add_argument("--json", help="JSON pattern definition file")
    parser.add_argument("--output", "-o", help="Output binary file")
    parser.add_argument("--write", action="store_true", help="Write to NTAG5 via PCSC")
    parser.add_argument("--builtin", action="store_true", help="Use built-in patterns (default if no --json)")
    parser.add_argument("--dump", action="store_true", help="Hex dump the binary")
    args = parser.parse_args()

    if args.json:
        with open(args.json) as f:
            config = json.load(f)
    else:
        config = builtin_patterns()

    data = encode_library(config)
    print(f"Library: {len(config['patterns'])} patterns, {len(data)} bytes")
    print(f"CRC-16: 0x{struct.unpack_from('>H', data, 14)[0]:04X}")

    if args.dump:
        for i in range(0, len(data), 16):
            hex_str = " ".join(f"{b:02X}" for b in data[i:i+16])
            ascii_str = "".join(chr(b) if 32 <= b < 127 else "." for b in data[i:i+16])
            print(f"  {i:04X}: {hex_str:<48s} {ascii_str}")

    if args.output:
        with open(args.output, "wb") as f:
            f.write(data)
        print(f"Written to {args.output}")

    if args.write:
        write_to_ntag5(data)


if __name__ == "__main__":
    main()