#!/usr/bin/env python3
"""Probe G2 bootloader at different baud rates and try interrupt sequences.

The bootloader phase is <50ms. This script:
1. Shows the raw hex of the garbage bytes after "Application is running..."
2. Tries rebooting at different baud rates to decode the bootloader's native rate
3. Tries sending interrupt sequences during the bootloader window

Usage:
    uv run scripts/boot_baud_probe.py
"""

import time

import serial

PORT = "/dev/ttyUSB2"
APP_BAUD = 115200


def reboot_and_capture(
    port: str,
    capture_baud: int,
    duration: float = 5.0,
    pre_interrupt: bytes | None = None,
    interrupt_delay: float = 0.01,
) -> bytes:
    """Reboot at app baud, optionally switch to capture_baud, read everything."""

    # Open at application baud rate to send reboot command
    ser = serial.Serial(port, APP_BAUD, timeout=0.5)
    ser.reset_input_buffer()

    # Make sure we're at a prompt — send CR
    ser.write(b"\r")
    time.sleep(0.3)
    resp = ser.read(4096).decode("utf-8", errors="replace")

    if "OS>" not in resp:
        if "TRK>" in resp:
            ser.write(b"os\r")
            time.sleep(0.3)
            ser.read(4096)
        else:
            ser.write(b"q\r")
            time.sleep(0.3)
            ser.read(4096)
            ser.write(b"os\r")
            time.sleep(0.3)
            ser.read(4096)

    ser.reset_input_buffer()
    ser.write(b"reboot\r")

    # If we need to switch baud rate, do it immediately
    if capture_baud != APP_BAUD:
        time.sleep(0.01)  # tiny delay for reboot cmd to be sent
        ser.baudrate = capture_baud

    # Send interrupt if requested
    if pre_interrupt:
        time.sleep(interrupt_delay)
        if pre_interrupt == b"BREAK":
            ser.send_break(duration=0.25)
        else:
            ser.write(pre_interrupt)

    # Capture everything
    buf = bytearray()
    start = time.monotonic()
    ser.timeout = 0.05
    while time.monotonic() - start < duration:
        data = ser.read(4096)
        if data:
            buf.extend(data)

    ser.close()
    return bytes(buf)


def hex_dump(data: bytes, prefix: str = "  ") -> None:
    """Print hex dump with ASCII sidebar."""
    for i in range(0, len(data), 16):
        chunk = data[i : i + 16]
        hex_part = " ".join(f"{b:02x}" for b in chunk)
        ascii_part = "".join(chr(b) if 32 <= b < 127 else "." for b in chunk)
        print(f"{prefix}{i:04x}: {hex_part:<48s}  {ascii_part}")


def main():
    print("=" * 70)
    print("G2 BOOTLOADER BAUD RATE PROBE")
    print("=" * 70)

    # --- Phase 1: Capture raw boot at 115200 and show hex ---
    print("\n--- Phase 1: Raw boot capture at 115200 (hex dump) ---\n")
    data = reboot_and_capture(PORT, APP_BAUD, duration=2.0)
    print(f"Captured {len(data)} bytes:\n")
    hex_dump(data)

    # Find the garbage bytes between "running..." and "Application Starting"
    idx1 = data.find(b"running...")
    idx2 = data.find(b"Application Starting")
    if idx1 >= 0 and idx2 >= 0:
        between = data[idx1 + len(b"running...") : idx2]
        print("\nBytes between 'running...' and 'Application Starting':")
        print(f"  Raw: {between!r}")
        print(f"  Hex: {between.hex(' ')}")
        print(f"  Len: {len(between)}")

    # Wait for boot to complete
    print("\n  Waiting for boot to complete...")
    time.sleep(12)

    # --- Phase 2: Try different baud rates during bootloader ---
    test_bauds = [9600, 19200, 38400, 57600, 115200, 230400, 460800]

    print("\n--- Phase 2: Bootloader at different baud rates ---\n")
    for baud in test_bauds:
        print(f"\n  Trying {baud} baud...")
        data = reboot_and_capture(PORT, baud, duration=2.0)

        # Filter to just the first ~200 bytes (bootloader phase)
        preview = data[:200]
        # Check if it looks like readable ASCII
        ascii_count = sum(1 for b in preview if 32 <= b < 127 or b in (10, 13))
        ratio = ascii_count / max(len(preview), 1)

        print(f"    Captured {len(data)} bytes, ASCII ratio: {ratio:.0%}")
        if ratio > 0.5:
            text = preview.decode("utf-8", errors="replace")
            print(f"    Preview: {text[:120]!r}")
        else:
            if preview:
                print(f"    Hex: {preview[:32].hex(' ')}")

        # Wait for boot
        time.sleep(12)

    # --- Phase 3: Try interrupt sequences at 115200 ---
    print("\n--- Phase 3: Interrupt sequences at 115200 ---\n")

    interrupts = [
        (0.005, "5ms CR", b"\r\r\r\r\r"),
        (0.005, "5ms ESC", b"\x1b\x1b\x1b"),
        (0.005, "5ms BREAK", b"BREAK"),
        (0.005, "5ms 0x55 autobaud", b"\x55\x55\x55\x55\x55"),
        (0.005, "5ms 0x7F (DEL)", b"\x7f\x7f\x7f"),
        (0.005, "5ms 'bl'", b"bl\r"),
        (0.01, "10ms CR burst", b"\r" * 20),
        (0.02, "20ms '?' burst", b"?\r?\r?\r"),
        (0.03, "30ms CR", b"\r\r\r"),
    ]

    for delay, desc, payload in interrupts:
        print(f"\n  [{desc}] (delay={delay}s)")
        data = reboot_and_capture(
            PORT, APP_BAUD, duration=3.0, pre_interrupt=payload, interrupt_delay=delay
        )

        text = data.decode("utf-8", errors="replace")

        # Look for anything unusual (not normal boot)
        normal_markers = ["Application Starting", "MotorInit", "BCM4515"]
        unusual = False

        # Check if we got a response BEFORE normal boot messages
        bl_idx = text.find("Bootloader version")
        app_idx = text.find("Application Starting")

        if bl_idx >= 0 and app_idx >= 0:
            between = text[bl_idx:app_idx]
            if len(between) > 80:  # more than just "Bootloader version: 1.01\r\n..."
                unusual = True
                print("    *** EXTRA DATA in bootloader phase! ***")
                print(f"    Between BL and App: {between!r}")

        if "BL>" in text or "CMD>" in text or "download" in text.lower():
            unusual = True
            print("    *** BOOTLOADER PROMPT DETECTED ***")

        if not unusual:
            # Check if boot proceeded normally
            if any(m in text for m in normal_markers):
                print("    Normal boot (no intercept)")
            else:
                print(f"    Unexpected: {text[:150]!r}")

        time.sleep(12)

    print("\n" + "=" * 70)
    print("PROBE COMPLETE")
    print("=" * 70)


if __name__ == "__main__":
    main()