#!/usr/bin/env python3
"""Compare BCM4500 register behavior: stock EEPROM firmware vs custom firmware.

USAGE:
1. Power cycle the SkyWalker-1 (unplug/replug USB)
2. Run this script IMMEDIATELY (before loading custom FW)
3. The script tests registers under stock firmware first,
   then loads custom firmware and tests again.

If stock firmware shows different register behavior, the issue is
in our custom firmware's boot sequence.
"""
import sys
import time
import usb.core
sys.path.insert(0, 'tools')

VID = 0x09C0
PID = 0x0203
BCM4500_ADDR = 0x08

# ============================================================
# Raw USB helpers (work with any firmware)
# ============================================================
def vendor_in(dev, cmd, value=0, index=0, length=1):
    """Send a vendor IN request and return the response bytes."""
    return dev.ctrl_transfer(
        0xC0,     # bmRequestType: vendor, device-to-host
        cmd,      # bRequest
        value,    # wValue
        index,    # wIndex
        length    # wLength
    )

def vendor_out(dev, cmd, value=0, index=0, data=None):
    """Send a vendor OUT request."""
    dev.ctrl_transfer(
        0x40,     # bmRequestType: vendor, host-to-device
        cmd,      # bRequest
        value,    # wValue
        index,    # wIndex
        data if data else b''
    )

def read_bcm_reg(dev, reg):
    """Read one BCM4500 register via stock-compatible I2C read (0xB5).
    This might not exist on stock firmware, so we use the
    stock READ_8PSK_REG (0x81) as a fallback."""
    try:
        data = vendor_in(dev, 0xB5, value=BCM4500_ADDR, index=reg, length=1)
        return data[0]
    except Exception:
        return None

def read_bcm_reg_stock(dev, reg):
    """Read BCM4500 register via stock READ_8PSK_REG (0x81).
    wValue = register address, returns 1 byte."""
    try:
        data = vendor_in(dev, 0x81, value=reg, index=0, length=1)
        return data[0]
    except Exception:
        return None

def read_all_key_regs(dev, method='0xB5'):
    """Read key BCM4500 registers."""
    regs = [0xA0, 0xA2, 0xA4, 0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xAB]
    results = {}
    for reg in regs:
        if method == '0xB5':
            val = read_bcm_reg(dev, reg)
        else:
            val = read_bcm_reg_stock(dev, reg)
        results[reg] = val
        v_str = f'0x{val:02X}' if val is not None else 'FAIL'
        print(f'    0x{reg:02X}: {v_str}')
    return results

# ============================================================
# MAIN
# ============================================================
print('=== Stock vs Custom Firmware BCM4500 Comparison ===')
print()

# Find the device
dev = usb.core.find(idVendor=VID, idProduct=PID)
if dev is None:
    print('ERROR: SkyWalker-1 not found!')
    print('Make sure to power cycle the device first (stock firmware must be running)')
    sys.exit(1)

print(f'Found SkyWalker-1: Bus {dev.bus} Addr {dev.address}')
print(f'  VID=0x{dev.idVendor:04X} PID=0x{dev.idProduct:04X}')

# Try to get firmware version (our custom command)
try:
    fw = vendor_in(dev, 0x80, value=0, index=0, length=3)
    print(f'  Firmware response (0x80): {list(fw)}')
except Exception as e:
    print(f'  Firmware version (0x80): {e}')

# ============================================================
# Phase 1: Test under stock firmware (before boot command)
# ============================================================
print('\n' + '='*60)
print('PHASE 1: Stock firmware — BEFORE boot command')
print('='*60)

print('\n  Registers via 0x81 (READ_8PSK_REG):')
regs_stock_pre_81 = {}
for reg in [0xA0, 0xA2, 0xA4, 0xA6, 0xA7, 0xA8]:
    val = read_bcm_reg_stock(dev, reg)
    regs_stock_pre_81[reg] = val
    v_str = f'0x{val:02X}' if val is not None else 'FAIL'
    print(f'    0x{reg:02X}: {v_str}')

print('\n  Registers via 0xB5 (I2C_RAW_READ) — may fail on stock FW:')
regs_stock_pre_b5 = {}
for reg in [0xA0, 0xA2, 0xA4, 0xA6, 0xA7, 0xA8]:
    val = read_bcm_reg(dev, reg)
    regs_stock_pre_b5[reg] = val
    v_str = f'0x{val:02X}' if val is not None else 'N/A (stock FW lacks 0xB5)'
    print(f'    0x{reg:02X}: {v_str}')

# ============================================================
# Phase 2: Boot BCM4500 under stock firmware
# ============================================================
print('\n' + '='*60)
print('PHASE 2: Stock firmware — boot BCM4500 (0x89)')
print('='*60)

try:
    result = vendor_in(dev, 0x89, value=1, index=0, length=3)
    print(f'  Boot result: [{result[0]:02X}, {result[1]:02X}, {result[2]:02X}]')
    cfg = result[0]
    bits = []
    if cfg & 0x01: bits.append('Started')
    if cfg & 0x02: bits.append('FW_Loaded')
    if cfg & 0x04: bits.append('Intersil')
    if cfg & 0x08: bits.append('DVBmode')
    print(f'  Config: 0x{cfg:02X} ({" | ".join(bits) if bits else "none"})')
    print(f'  Stage: 0x{result[1]:02X}')
except Exception as e:
    print(f'  Boot failed: {e}')

time.sleep(0.5)

print('\n  Registers via 0x81 after boot:')
regs_stock_post_81 = {}
for reg in [0xA0, 0xA2, 0xA4, 0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xAB]:
    val = read_bcm_reg_stock(dev, reg)
    regs_stock_post_81[reg] = val
    v_str = f'0x{val:02X}' if val is not None else 'FAIL'
    marker = ''
    if val is not None and reg in regs_stock_pre_81:
        if regs_stock_pre_81.get(reg) != val:
            marker = f' (was 0x{regs_stock_pre_81[reg]:02X})' if regs_stock_pre_81[reg] is not None else ''
    print(f'    0x{reg:02X}: {v_str}{marker}')

# Try indirect read via stock firmware's 0x81 with indirect addressing
# (0x81 might support indirect reads differently)
print('\n  Testing indirect register reads via 0x81:')
for page in [0x00, 0x06, 0x07, 0x0F]:
    # Stock firmware might use a different convention for indirect reads
    # Try reading page register values
    val = read_bcm_reg_stock(dev, page)
    v_str = f'0x{val:02X}' if val is not None else 'FAIL'
    print(f'    Page 0x{page:02X} via 0x81: {v_str}')

# Try signal monitor
print('\n  Signal status:')
try:
    sig = vendor_in(dev, 0x82, value=0, index=0, length=4)
    print(f'    GET_8PSK_SIGNAL (0x82): {list(sig)}')
except Exception as e:
    print(f'    0x82: {e}')

try:
    lock = vendor_in(dev, 0x83, value=0, index=0, length=1)
    print(f'    GET_8PSK_LOCK (0x83): 0x{lock[0]:02X}')
except Exception as e:
    print(f'    0x83: {e}')

# ============================================================
# Summary
# ============================================================
print('\n' + '='*60)
print('SUMMARY')
print('='*60)

def compare_regs(label, regs):
    vals = set(v for v in regs.values() if v is not None)
    if len(vals) == 1:
        print(f'  {label}: ALL = 0x{list(vals)[0]:02X}')
    elif len(vals) == 0:
        print(f'  {label}: ALL FAILED')
    else:
        print(f'  {label}: Mixed: {", ".join(f"0x{v:02X}" for v in sorted(vals))}')

compare_regs('Stock FW before boot (0x81)', regs_stock_pre_81)
compare_regs('Stock FW after boot (0x81)', regs_stock_post_81)

print()
print('If stock FW shows DIFFERENT register values (not all 0x02),')
print('then the BCM4500 is truly functional under stock FW and our')
print('custom boot sequence is missing something.')
print()
print('If stock FW also shows all 0x02, then the register behavior')
print('is normal and the 0x02 IS the legitimate power-on value.')

print('\n=== Done ===')