#!/usr/bin/env python3
"""
Genpix SkyWalker-1 EEPROM firmware flash tool.

Writes C2-format firmware images to the Cypress FX2 boot EEPROM via
the I2C_WRITE vendor command.

Protocol:
  I2C_WRITE (0x83): wValue=0x51, wIndex=offset, data=bytes
  I2C_READ  (0x84): wValue=0x51, wIndex=offset, length=chunk_size

The EEPROM uses Cypress C2 IIC boot format:
  - Header: C2 VID_L VID_H PID_L PID_H DID_L DID_H CONFIG
  - Records: LEN_H LEN_L ADDR_H ADDR_L DATA[LEN]
  - End: 80 01 ENTRY_H ENTRY_L (reset vector)

WARNING: Flashing incorrect firmware can brick the device. The FX2
boots from this EEPROM on power-up -- a corrupted image means the
device will not enumerate on USB until the EEPROM is reprogrammed
with an external programmer or the FX2 boot ROM's A0 vendor request.
"""
import usb.core, usb.util, sys, struct, time, os

VENDOR_ID = 0x09C0
PRODUCT_ID = 0x0203
I2C_WRITE = 0x83
I2C_READ = 0x84
EEPROM_SLAVE = 0x51

# EEPROM page write parameters
PAGE_SIZE = 16           # Conservative page size for 24Cxx EEPROMs
WRITE_CYCLE_MS = 10      # Max internal write cycle time per page
MAX_EEPROM_SIZE = 16384  # 128Kbit / 16KB max addressable via wIndex


def find_device():
    dev = usb.core.find(idVendor=VENDOR_ID, idProduct=PRODUCT_ID)
    if dev is None:
        print("SkyWalker-1 not found")
        sys.exit(1)
    return dev


def detach_driver(dev):
    intf_num = None
    for cfg in dev:
        for intf in cfg:
            if dev.is_kernel_driver_active(intf.bInterfaceNumber):
                try:
                    dev.detach_kernel_driver(intf.bInterfaceNumber)
                    intf_num = intf.bInterfaceNumber
                except usb.core.USBError as e:
                    print(f"Cannot detach driver: {e}")
                    print("Try: sudo modprobe -r dvb_usb_gp8psk")
                    sys.exit(1)
    try:
        dev.set_configuration()
    except:
        pass
    return intf_num


def eeprom_read(dev, offset, length=64):
    """Read from EEPROM at given offset."""
    return dev.ctrl_transfer(
        usb.util.CTRL_TYPE_VENDOR | usb.util.CTRL_IN,
        I2C_READ, EEPROM_SLAVE, offset, length, 2000)


def eeprom_write(dev, offset, data):
    """Write data to EEPROM at given offset. Caller handles page alignment."""
    return dev.ctrl_transfer(
        usb.util.CTRL_TYPE_VENDOR | usb.util.CTRL_OUT,
        I2C_WRITE, EEPROM_SLAVE, offset, data, 2000)


def eeprom_read_all(dev, size, label="Reading"):
    """Read entire EEPROM contents up to size bytes."""
    chunk_size = 64
    data = bytearray()
    for offset in range(0, size, chunk_size):
        remaining = min(chunk_size, size - offset)
        chunk = eeprom_read(dev, offset, remaining)
        if chunk is None:
            print(f"\n  Read failed at offset 0x{offset:04X}")
            return None
        data.extend(bytes(chunk))
        if offset % 1024 == 0:
            pct = offset * 100 // size
            print(f"\r  {label}: 0x{offset:04X} / 0x{size:04X}  [{pct:3d}%]",
                  end="", flush=True)
    print(f"\r  {label}: 0x{size:04X} / 0x{size:04X}  [100%]   ")
    return data


def parse_c2_header(data):
    """Parse Cypress C2 boot EEPROM header. Returns dict or None."""
    if len(data) < 8:
        return None
    if data[0] != 0xC2:
        return None

    vid = data[2] << 8 | data[1]
    pid = data[4] << 8 | data[3]
    did = data[6] << 8 | data[5]
    config = data[7]

    return {"vid": vid, "pid": pid, "did": did, "config": config}


def parse_records(data, offset=8):
    """Parse C2 load records from EEPROM data."""
    records = []
    while offset < len(data) - 4:
        rec_len = (data[offset] << 8) | data[offset + 1]
        rec_addr = (data[offset + 2] << 8) | data[offset + 3]

        if rec_len == 0x8001:
            records.append({
                "type": "end",
                "entry_point": rec_addr,
                "offset": offset
            })
            break
        elif rec_len == 0 or rec_len > 0x4000:
            records.append({
                "type": "invalid",
                "raw_len": rec_len,
                "offset": offset
            })
            break

        rec_data = data[offset + 4:offset + 4 + rec_len]
        records.append({
            "type": "data",
            "length": rec_len,
            "load_addr": rec_addr,
            "data": bytes(rec_data),
            "offset": offset
        })
        offset += 4 + rec_len

    return records


def print_c2_header(header, prefix="  "):
    """Display parsed C2 header fields."""
    print(f"{prefix}Format:  C2 (Large EEPROM, code loads to internal RAM)")
    print(f"{prefix}VID:     0x{header['vid']:04X}"
          f" {'(Genpix)' if header['vid'] == 0x09C0 else ''}")
    print(f"{prefix}PID:     0x{header['pid']:04X}"
          f" {'(SkyWalker-1)' if header['pid'] == 0x0203 else ''}")
    print(f"{prefix}DID:     0x{header['did']:04X}")
    print(f"{prefix}Config:  0x{header['config']:02X}", end="")

    config_flags = []
    if header["config"] & 0x40:
        config_flags.append("400kHz I2C")
    if header["config"] & 0x04:
        config_flags.append("disconnect")
    if config_flags:
        print(f" ({', '.join(config_flags)})")
    else:
        print()


def print_c2_records(records, prefix="  "):
    """Display parsed C2 load records."""
    total_code = 0
    for i, rec in enumerate(records):
        if rec["type"] == "data":
            end_addr = rec["load_addr"] + rec["length"] - 1
            preview = rec["data"][:8].hex(' ')
            print(f"{prefix}[{i}] {rec['length']:5d} bytes -> "
                  f"0x{rec['load_addr']:04X}-0x{end_addr:04X}  "
                  f"[{preview}...]")
            total_code += rec["length"]
        elif rec["type"] == "end":
            print(f"{prefix}[{i}] END MARKER -> entry point: "
                  f"0x{rec['entry_point']:04X}")
        else:
            print(f"{prefix}[{i}] INVALID (raw_len=0x{rec['raw_len']:04X}) "
                  f"at EEPROM offset 0x{rec['offset']:04X}")

    data_recs = [r for r in records if r["type"] == "data"]
    print(f"\n{prefix}Total firmware: {total_code} bytes in "
          f"{len(data_recs)} segments")
    end_recs = [r for r in records if r["type"] == "end"]
    if end_recs:
        print(f"{prefix}Entry point:   0x{end_recs[0]['entry_point']:04X} "
              f"(LJMP target after boot)")


def validate_c2_image(data, label="image"):
    """Validate a C2 firmware image. Returns (header, records) or exits."""
    if len(data) < 12:
        print(f"  {label}: too small ({len(data)} bytes, need at least 12)")
        return None, None

    if data[0] != 0xC2:
        print(f"  {label}: not a C2 image (first byte: 0x{data[0]:02X}, "
              f"expected 0xC2)")
        return None, None

    header = parse_c2_header(data)
    if header is None:
        print(f"  {label}: failed to parse C2 header")
        return None, None

    records = parse_records(data)
    if not records:
        print(f"  {label}: no load records found")
        return None, None

    end_recs = [r for r in records if r["type"] == "end"]
    invalid_recs = [r for r in records if r["type"] == "invalid"]
    if not end_recs:
        print(f"  {label}: WARNING -- no end marker found")
    if invalid_recs:
        print(f"  {label}: WARNING -- {len(invalid_recs)} invalid record(s)")

    return header, records


def cmd_info(args):
    """Parse and display C2 header info from a .bin file."""
    if not os.path.exists(args.file):
        print(f"File not found: {args.file}")
        sys.exit(1)

    with open(args.file, 'rb') as f:
        data = f.read()

    print(f"C2 Image: {args.file}")
    print(f"File size: {len(data)} bytes")
    print("=" * 40)

    header, records = validate_c2_image(data, args.file)
    if header is None:
        sys.exit(1)

    print("\nHeader:")
    print_c2_header(header)

    print("\nLoad Records:")
    print_c2_records(records)

    # Compute EEPROM usage (header + record headers + data + end marker)
    if records:
        last = records[-1]
        if last["type"] == "end":
            eeprom_end = last["offset"] + 4
        elif last["type"] == "data":
            eeprom_end = last["offset"] + 4 + last["length"]
        else:
            eeprom_end = last["offset"]
        print(f"\n  EEPROM footprint: {eeprom_end} bytes "
              f"(0x{eeprom_end:04X})")


def cmd_backup(args):
    """Dump current EEPROM contents to a file."""
    print("Genpix SkyWalker-1 EEPROM Backup")
    print("=" * 40)

    dev = find_device()
    print(f"Found device: Bus {dev.bus} Addr {dev.address}")
    intf = detach_driver(dev)

    try:
        size = args.max_size
        print(f"\nReading EEPROM ({size} bytes)...")
        data = eeprom_read_all(dev, size)
        if data is None:
            print("Backup failed: read error")
            sys.exit(1)

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

        # Show header info
        header = parse_c2_header(data)
        if header:
            print("\nHeader:")
            print_c2_header(header)
    finally:
        if intf is not None:
            try:
                usb.util.release_interface(dev, intf)
                dev.attach_kernel_driver(intf)
                print("\nRe-attached kernel driver")
            except:
                print("\nNote: run 'sudo modprobe dvb_usb_gp8psk' to reload")


def cmd_verify(args):
    """Compare a .bin file against current EEPROM contents."""
    if not os.path.exists(args.file):
        print(f"File not found: {args.file}")
        sys.exit(1)

    with open(args.file, 'rb') as f:
        image = f.read()

    print("Genpix SkyWalker-1 EEPROM Verify")
    print("=" * 40)

    # Validate the image first
    header, records = validate_c2_image(image, args.file)
    if header is None:
        sys.exit(1)

    print(f"\nImage: {args.file} ({len(image)} bytes)")
    print_c2_header(header)

    dev = find_device()
    print(f"\nFound device: Bus {dev.bus} Addr {dev.address}")
    intf = detach_driver(dev)

    try:
        print(f"\nReading EEPROM ({len(image)} bytes)...")
        eeprom = eeprom_read_all(dev, len(image), label="Verify")
        if eeprom is None:
            print("Verify failed: read error")
            sys.exit(1)

        # Compare byte-by-byte
        mismatches = []
        for i in range(len(image)):
            if i < len(eeprom) and image[i] != eeprom[i]:
                mismatches.append(i)

        if not mismatches:
            print(f"\n  MATCH -- EEPROM contents match {args.file}")
        else:
            print(f"\n  MISMATCH -- {len(mismatches)} byte(s) differ:")
            for off in mismatches[:32]:
                exp = image[off]
                got = eeprom[off] if off < len(eeprom) else 0xFF
                print(f"    0x{off:04X}: expected 0x{exp:02X}, "
                      f"got 0x{got:02X}")
            if len(mismatches) > 32:
                print(f"    ... and {len(mismatches) - 32} more")
            sys.exit(1)
    finally:
        if intf is not None:
            try:
                usb.util.release_interface(dev, intf)
                dev.attach_kernel_driver(intf)
                print("\nRe-attached kernel driver")
            except:
                print("\nNote: run 'sudo modprobe dvb_usb_gp8psk' to reload")


def cmd_flash(args):
    """Write a C2-format .bin file to the EEPROM."""
    if not os.path.exists(args.file):
        print(f"File not found: {args.file}")
        sys.exit(1)

    with open(args.file, 'rb') as f:
        image = f.read()

    print("Genpix SkyWalker-1 EEPROM Flash")
    print("=" * 40)
    print()
    print("  *** FIRMWARE FLASH -- READ CAREFULLY ***")
    print("  Writing bad firmware will brick the device.")
    print("  The SkyWalker-1 boots from this EEPROM on power-up.")
    print("  A corrupted image = no USB enumeration.")
    print()

    # Validate input image
    img_header, img_records = validate_c2_image(image, args.file)
    if img_header is None:
        sys.exit(1)

    print(f"Image: {args.file} ({len(image)} bytes)")
    print_c2_header(img_header)

    # Size sanity check
    if len(image) > MAX_EEPROM_SIZE:
        print(f"\n  Image too large: {len(image)} bytes "
              f"(max {MAX_EEPROM_SIZE})")
        sys.exit(1)

    if len(image) < 12:
        print(f"\n  Image too small: {len(image)} bytes")
        sys.exit(1)

    # Connect to device
    dev = find_device()
    print(f"\nFound device: Bus {dev.bus} Addr {dev.address}")
    intf = detach_driver(dev)

    try:
        # Check VID/PID against the connected device
        if not args.force:
            if img_header["vid"] != VENDOR_ID:
                print(f"\n  VID mismatch: image has 0x{img_header['vid']:04X},"
                      f" device is 0x{VENDOR_ID:04X}")
                print("  Use --force to override")
                sys.exit(1)
            if img_header["pid"] != PRODUCT_ID:
                print(f"\n  PID mismatch: image has 0x{img_header['pid']:04X},"
                      f" device is 0x{PRODUCT_ID:04X}")
                print("  Use --force to override")
                sys.exit(1)
        elif img_header["vid"] != VENDOR_ID or img_header["pid"] != PRODUCT_ID:
            print(f"\n  WARNING: VID/PID mismatch (--force active)")
            print(f"  Image:  VID=0x{img_header['vid']:04X} "
                  f"PID=0x{img_header['pid']:04X}")
            print(f"  Device: VID=0x{VENDOR_ID:04X} PID=0x{PRODUCT_ID:04X}")

        # Backup current EEPROM
        if not args.no_backup:
            ts = time.strftime("%Y%m%d_%H%M%S")
            backup_file = f"eeprom_backup_{ts}.bin"
            print(f"\nBacking up current EEPROM to {backup_file}...")
            backup = eeprom_read_all(dev, MAX_EEPROM_SIZE, label="Backup")
            if backup is None:
                print("  Backup failed: read error. Aborting.")
                sys.exit(1)
            with open(backup_file, 'wb') as f:
                f.write(backup)
            print(f"  Backup saved: {backup_file} ({len(backup)} bytes)")

            # Show what's currently on the EEPROM
            cur_header = parse_c2_header(backup)
            if cur_header:
                print("\n  Current EEPROM:")
                print_c2_header(cur_header, prefix="    ")
        else:
            print("\n  Skipping backup (--no-backup)")

        # Dry-run stops here
        if args.dry_run:
            print("\n  DRY RUN -- would write {0} bytes in {1} pages".format(
                len(image), (len(image) + PAGE_SIZE - 1) // PAGE_SIZE))
            print("  No changes made.")
            return

        # Final confirmation
        print(f"\nAbout to write {len(image)} bytes to EEPROM...")
        print("  Press Ctrl+C within 3 seconds to abort.")
        try:
            for i in range(3, 0, -1):
                print(f"\r  Writing in {i}... ", end="", flush=True)
                time.sleep(1)
            print("\r  Writing now...       ")
        except KeyboardInterrupt:
            print("\n  Aborted.")
            return

        # Write in page-sized chunks
        total_pages = (len(image) + PAGE_SIZE - 1) // PAGE_SIZE
        write_errors = 0

        for page_num in range(total_pages):
            offset = page_num * PAGE_SIZE
            end = min(offset + PAGE_SIZE, len(image))
            chunk = image[offset:end]

            pct = (page_num + 1) * 100 // total_pages
            print(f"\r  Write: 0x{offset:04X} / 0x{len(image):04X}  "
                  f"[{pct:3d}%]", end="", flush=True)

            try:
                written = eeprom_write(dev, offset, chunk)
                if written != len(chunk):
                    print(f"\n  Short write at 0x{offset:04X}: "
                          f"sent {len(chunk)}, wrote {written}")
                    write_errors += 1
            except usb.core.USBError as e:
                print(f"\n  Write error at 0x{offset:04X}: {e}")
                write_errors += 1

            # Wait for EEPROM internal write cycle
            time.sleep(WRITE_CYCLE_MS / 1000.0)

        print(f"\r  Write: 0x{len(image):04X} / 0x{len(image):04X}  "
              f"[100%]   ")

        if write_errors:
            print(f"\n  WARNING: {write_errors} write error(s) occurred")

        # Verify by reading back
        print(f"\nVerifying ({len(image)} bytes)...")
        verify = eeprom_read_all(dev, len(image), label="Verify")
        if verify is None:
            print("  Verify failed: read error")
            print("  *** EEPROM STATE UNKNOWN -- check before power cycling ***")
            sys.exit(1)

        mismatches = []
        for i in range(len(image)):
            if i < len(verify) and image[i] != verify[i]:
                mismatches.append(i)

        if not mismatches:
            print(f"\n  VERIFIED -- all {len(image)} bytes match")
            print("  Flash complete. Power cycle the device to boot new firmware.")
        else:
            print(f"\n  VERIFY FAILED -- {len(mismatches)} byte(s) differ:")
            for off in mismatches[:16]:
                exp = image[off]
                got = verify[off] if off < len(verify) else 0xFF
                print(f"    0x{off:04X}: wrote 0x{exp:02X}, "
                      f"read 0x{got:02X}")
            if len(mismatches) > 16:
                print(f"    ... and {len(mismatches) - 16} more")
            print("\n  *** EEPROM CONTENTS DO NOT MATCH IMAGE ***")
            print("  Do NOT power cycle until this is resolved.")
            sys.exit(1)

    finally:
        if intf is not None:
            try:
                usb.util.release_interface(dev, intf)
                dev.attach_kernel_driver(intf)
                print("\nRe-attached kernel driver")
            except:
                print("\nNote: run 'sudo modprobe dvb_usb_gp8psk' to reload")


def main():
    import argparse
    parser = argparse.ArgumentParser(
        description="SkyWalker-1 EEPROM firmware flash tool")
    sub = parser.add_subparsers(dest='command', required=True)

    # info
    p_info = sub.add_parser('info',
                            help='Parse and display C2 header from a .bin file')
    p_info.add_argument('file', help='C2 firmware image (.bin)')

    # backup
    p_backup = sub.add_parser('backup',
                              help='Dump current EEPROM to a file')
    p_backup.add_argument('-o', '--output', default='skywalker1_eeprom.bin',
                          help='Output file (default: skywalker1_eeprom.bin)')
    p_backup.add_argument('--max-size', type=int, default=MAX_EEPROM_SIZE,
                          help=f'Bytes to read (default: {MAX_EEPROM_SIZE})')

    # verify
    p_verify = sub.add_parser('verify',
                              help='Compare .bin file against EEPROM')
    p_verify.add_argument('file', help='C2 firmware image (.bin)')

    # flash
    p_flash = sub.add_parser('flash',
                             help='Write C2 firmware image to EEPROM')
    p_flash.add_argument('file', help='C2 firmware image (.bin)')
    p_flash.add_argument('--dry-run', action='store_true',
                         help='Show what would happen without writing')
    p_flash.add_argument('--no-backup', action='store_true',
                         help='Skip pre-flash EEPROM backup')
    p_flash.add_argument('--force', action='store_true',
                         help='Override VID/PID mismatch check')

    args = parser.parse_args()

    if args.command == 'info':
        cmd_info(args)
    elif args.command == 'backup':
        cmd_backup(args)
    elif args.command == 'verify':
        cmd_verify(args)
    elif args.command == 'flash':
        cmd_flash(args)


if __name__ == '__main__':
    main()