#!/usr/bin/env python3
"""Reboot G2 firmware and capture boot output with timestamps.

Pass 1: Observe boot sequence timing.
Pass 2: Try interrupt sequences during bootloader phase.

Usage:
    uv run scripts/boot_capture.py [--interrupt]
"""

import argparse
import time

import serial

PORT = "/dev/ttyUSB2"
BAUD = 115200
BOOT_CAPTURE_SECS = 30  # how long to capture after reboot


def timestamp() -> str:
    return f"{time.monotonic():.4f}"


def capture_boot(ser: serial.Serial, duration: float) -> list[tuple[float, bytes]]:
    """Read all data for `duration` seconds, returning (time, data) pairs."""
    chunks = []
    start = time.monotonic()
    ser.timeout = 0.05  # 50ms polling
    while time.monotonic() - start < duration:
        data = ser.read(4096)
        if data:
            chunks.append((time.monotonic() - start, data))
    return chunks


def enter_os_and_reboot(ser: serial.Serial) -> None:
    """Navigate to OS menu and send reboot command."""
    ser.reset_input_buffer()

    # We left the port in the OS> menu before closing mcserial.
    # But the port close/reopen may have reset state. Try sending
    # a bare CR first to see where we are.
    ser.write(b"\r")
    time.sleep(0.3)
    response = ser.read(4096)
    decoded = response.decode("utf-8", errors="replace")
    print(f"  Prompt check: {decoded.strip()!r}")

    if "OS>" in decoded:
        print("  Already in OS menu.")
    elif "TRK>" in decoded:
        print("  At root prompt, entering OS menu...")
        ser.write(b"os\r")
        time.sleep(0.3)
        ser.read(4096)  # consume echo
    elif "MOT>" in decoded or "DVB>" in decoded:
        print("  In a submenu, exiting to root first...")
        ser.write(b"q\r")
        time.sleep(0.3)
        ser.read(4096)
        ser.write(b"os\r")
        time.sleep(0.3)
        ser.read(4096)
    else:
        # Unknown state — try brute force: q to root, then os
        print("  Unknown state, trying q -> os...")
        ser.write(b"q\r")
        time.sleep(0.3)
        ser.read(4096)
        ser.write(b"os\r")
        time.sleep(0.3)
        ser.read(4096)

    # Flush buffer
    ser.reset_input_buffer()

    print("  Sending 'reboot' command...")
    ser.write(b"reboot\r")


def pass1_observe(ser: serial.Serial) -> list[tuple[float, bytes]]:
    """Reboot and passively observe boot output."""
    print("\n=== PASS 1: OBSERVE BOOT SEQUENCE ===\n")
    enter_os_and_reboot(ser)

    print(f"  Capturing for {BOOT_CAPTURE_SECS}s...\n")
    chunks = capture_boot(ser, BOOT_CAPTURE_SECS)

    print("--- Boot Output (with timestamps) ---\n")
    full_output = bytearray()
    for t, data in chunks:
        full_output.extend(data)
        text = (
            data.decode("utf-8", errors="replace")
            .replace("\r\n", "\n")
            .replace("\r", "\n")
        )
        for line in text.split("\n"):
            if line.strip():
                print(f"  [{t:7.3f}s] {line}")

    print(f"\n--- Total: {len(full_output)} bytes in {len(chunks)} chunks ---")
    return chunks


def pass2_interrupt(ser: serial.Serial) -> None:
    """Reboot and try interrupt sequences during bootloader phase."""
    print("\n=== PASS 2: INTERRUPT BOOT ===\n")

    # Interrupt sequences to try, with timing (seconds after reboot)
    interrupts = [
        # (delay_after_reboot, description, bytes_to_send)
        (0.05, "immediate CR", b"\r"),
        (0.05, "immediate ESC", b"\x1b"),
        (0.05, "immediate BREAK", "BREAK"),  # special handling
        (0.05, "immediate 0x55 (autobaud)", b"\x55\x55\x55\x55\x55"),
        (0.05, "immediate space", b" "),
        (0.1, "100ms CR", b"\r"),
        (0.1, "100ms ESC", b"\x1b"),
        (0.2, "200ms CR", b"\r"),
        (0.5, "500ms CR+ESC", b"\r\x1b\r"),
        (1.0, "1s CR", b"\r"),
        (1.0, "1s 'bl'", b"bl\r"),
        (1.0, "1s 'boot'", b"boot\r"),
        (2.0, "2s CR", b"\r"),
        (2.0, "2s '?'", b"?\r"),
    ]

    for delay, desc, payload in interrupts:
        print(f"\n--- Trying: {desc} (at +{delay}s) ---")
        enter_os_and_reboot(ser)

        # Wait for the specified delay
        time.sleep(delay)

        # Send the interrupt
        if payload == "BREAK":
            ser.send_break(duration=0.25)
            print("  Sent BREAK signal (250ms)")
        else:
            ser.write(payload)
            print(f"  Sent: {payload!r}")

        # Capture response for 5 seconds
        chunks = capture_boot(ser, 8)

        full = bytearray()
        for _, data in chunks:
            full.extend(data)
        text = full.decode("utf-8", errors="replace")

        # Check for interesting responses
        interesting = False
        for keyword in [
            "boot",
            "loader",
            "BL>",
            "CMD>",
            ">>",
            "ready",
            "download",
            "upload",
            "flash",
            "update",
            "xmodem",
            "ymodem",
            "zmodem",
            "srec",
            "hex",
            "binary",
        ]:
            if keyword.lower() in text.lower():
                interesting = True
                break

        if interesting:
            print("  *** INTERESTING RESPONSE ***")
            for t, data in chunks:
                line = data.decode("utf-8", errors="replace").strip()
                if line:
                    print(f"  [{t:7.3f}s] {line}")
        else:
            # Print condensed summary
            lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
            if lines:
                print(f"  Response: {len(full)} bytes, first: {lines[0][:80]!r}")
                if len(lines) > 1:
                    print(f"           last:  {lines[-1][:80]!r}")
            else:
                print(f"  No response ({len(full)} bytes)")

        # Wait for boot to complete before next attempt
        print("  Waiting for full boot...")
        time.sleep(max(0, BOOT_CAPTURE_SECS - 8 - delay))
        ser.reset_input_buffer()


def main():
    parser = argparse.ArgumentParser(description="G2 bootloader capture")
    parser.add_argument(
        "--interrupt", action="store_true", help="Pass 2: try interrupt sequences"
    )
    parser.add_argument("--port", default=PORT)
    parser.add_argument("--baud", type=int, default=BAUD)
    args = parser.parse_args()

    print(f"Opening {args.port} @ {args.baud}...")
    ser = serial.Serial(args.port, args.baud, timeout=1)
    ser.reset_input_buffer()

    try:
        pass1_observe(ser)

        if args.interrupt:
            # Wait for boot to fully complete
            print("\n  Waiting for boot to settle...")
            time.sleep(5)
            ser.reset_input_buffer()
            pass2_interrupt(ser)
    except KeyboardInterrupt:
        print("\n\nInterrupted by user.")
    finally:
        ser.close()
        print("\nPort closed.")


if __name__ == "__main__":
    main()