#!/usr/bin/env python3
"""
update_firmware.py — StellarMate Pro X RP2350 Firmware Updater

Flashes the RP2350 firmware unconditionally by entering BOOTSEL mode,
copying the UF2 file over USB mass storage, then resetting the RP.

GPIO mapping (LattePanda Mu / gpiochip0):
  SOC_GPP_D2 (offset 194) = RP2350_RUN      (active HIGH = hold in reset)
  SOC_GPP_D3 (offset 195) = RP2350_USB_BOOT (active HIGH = enter BOOTSEL)

Usage:
  sudo python3 update_firmware.py [--port /dev/ttyACM0] [--firmware /path/to/file.uf2]
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import time

import gpiod
import serial
from gpiod.line import Direction, Value

# ================= Configuration =================
DEFAULT_UART_PORT     = '/dev/ttyACM0'
DEFAULT_BAUD_RATE     = 115200
DEFAULT_UART_TIMEOUT  = 3          # seconds to wait for RP response
DEFAULT_FIRMWARE_PATH = '/usr/lib/firmware/stellarmate/stellarmateprox.uf2'

# GPIO
CHIP_PATH    = '/dev/gpiochip0'
PIN_GPP_D2   = 194   # RP2350_RUN      (active: hold RP in reset)
PIN_GPP_D3   = 195   # RP2350_USB_BOOT (active: enter USB BOOTSEL mode)

# USB mass-storage detection
PICO_PATTERNS        = ['RPI-RP2', 'PICO', 'RP2350']
MOUNT_WAIT_TIMEOUT   = 15   # seconds to wait for mass-storage to appear
POST_FLASH_WAIT      = 4    # seconds after copy before reset

# /dev/serial/by-id/ matching for stable RP2350 port discovery
BY_ID_DIR   = '/dev/serial/by-id'
NAME_PREFIX = 'usb-Raspberry_Pi_Pico_'
NAME_SUFFIX = '-if00'
# =================================================


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


# ------------------------------------------------------------------
# Serial port auto-detection
# ------------------------------------------------------------------

def find_serial_port(fallback: str) -> str:
    """
    Locate the RP2350 USB-CDC device via /dev/serial/by-id/ stable symlinks.
    The kernel creates entries like:
      /dev/serial/by-id/usb-Raspberry_Pi_Pico_F97E9EC4508F6499-if00 -> ../../ttyACM0
    Returns the resolved device path if found, otherwise the fallback.
    """
    if not os.path.isdir(BY_ID_DIR):
        return fallback

    prefix_len = len(NAME_PREFIX)
    suffix_len = len(NAME_SUFFIX)

    try:
        for entry in os.listdir(BY_ID_DIR):
            if not entry.startswith(NAME_PREFIX):
                continue
            if not entry.endswith(NAME_SUFFIX):
                continue
            # Ensure there's a serial number between prefix and suffix
            if len(entry) <= prefix_len + suffix_len:
                continue

            link_path = os.path.join(BY_ID_DIR, entry)
            resolved = os.path.realpath(link_path)
            log(f"RP2350 serial port auto-detected: {resolved} (from {entry})")
            return resolved
    except OSError as exc:
        log(f"Error scanning {BY_ID_DIR}: {exc}")

    return fallback


# ------------------------------------------------------------------
# UART handshake
# ------------------------------------------------------------------

def query_rp_version(port: str, baud: int, timeout: int) -> str | None:
    """
    Send a get_firmware_info command over UART and return the version string.
    Returns None if the RP does not respond within timeout.
    """
    command = json.dumps({
        "subsystem": "general",
        "command":   "get_firmware_info"
    }) + '\n'

    try:
        with serial.Serial(port, baud, timeout=timeout) as ser:
            ser.reset_input_buffer()
            ser.write(command.encode('utf-8'))
            log(f"Sent get_firmware_info to {port}, waiting for response...")

            deadline = time.time() + timeout
            while time.time() < deadline:
                line = ser.readline().decode('utf-8', errors='replace').strip()
                if not line:
                    continue
                try:
                    response = json.loads(line)
                    version = response.get('data', {}).get('version') or response.get('version')
                    if version:
                        return version
                except json.JSONDecodeError:
                    continue   # ignore non-JSON lines

    except serial.SerialException as exc:
        log(f"UART error: {exc}")

    return None


# ------------------------------------------------------------------
# GPIO / BOOTSEL helpers
# ------------------------------------------------------------------

def enter_bootsel_mode() -> tuple:
    """
    Assert RP2350_RUN + RP2350_USB_BOOT to put the RP into USB BOOTSEL mode.
    Returns (line_d2, line_d3) handles that must be released by the caller.
    """
    line_d2 = gpiod.request_lines(CHIP_PATH,
        consumer="GPIO_D2_RP2350_RUN",
        config={PIN_GPP_D2: gpiod.LineSettings(
            direction=Direction.OUTPUT,
            output_value=Value.INACTIVE)}
    )
    line_d3 = gpiod.request_lines(CHIP_PATH,
        consumer="GPIO_D3_RP2350_USB_BOOT",
        config={PIN_GPP_D3: gpiod.LineSettings(
            direction=Direction.OUTPUT,
            output_value=Value.INACTIVE)}
    )

    # Step 1: Assert RUN (hold RP in reset)
    line_d2.set_value(PIN_GPP_D2, Value.ACTIVE)
    log("RP2350 held in RESET (RP2350_RUN asserted)")
    time.sleep(0.5)

    # Step 2: Assert USB_BOOT (select BOOTSEL mode)
    line_d3.set_value(PIN_GPP_D3, Value.ACTIVE)
    log("RP2350_USB_BOOT asserted (BOOTSEL selected)")
    time.sleep(0.5)

    # Step 3: Release RUN → RP boots into BOOTSEL
    line_d2.set_value(PIN_GPP_D2, Value.INACTIVE)
    log("RP2350_RUN released — RP booting into USB Mass Storage mode")
    time.sleep(0.5)

    # Step 4: Release USB_BOOT (RP latched BOOTSEL on boot edge)
    line_d3.set_value(PIN_GPP_D3, Value.INACTIVE)
    log("RP2350_USB_BOOT released")

    return line_d2, line_d3


def reset_rp(line_d2) -> None:
    """Pulse RP2350_RUN to perform a clean restart."""
    line_d2.set_value(PIN_GPP_D2, Value.ACTIVE)
    log("RP2350 held in RESET for restart")
    time.sleep(0.5)
    line_d2.set_value(PIN_GPP_D2, Value.INACTIVE)
    log("RP2350 released — booting normally")


def release_gpio(*lines) -> None:
    for line in lines:
        if line:
            try:
                line.release()
            except Exception:
                pass


# ------------------------------------------------------------------
# Mass-storage detection
# ------------------------------------------------------------------

def find_pico_mount_point(timeout: int = MOUNT_WAIT_TIMEOUT) -> str | None:
    """
    Poll /proc/mounts until the RP appears as a USB mass-storage device.
    Returns the mount point string or None on timeout.
    """
    start = time.time()
    while time.time() - start < timeout:
        try:
            with open('/proc/mounts', 'r') as f:
                for line in f:
                    parts = line.split()
                    if len(parts) < 2:
                        continue
                    device, mountpoint = parts[0], parts[1]
                    for pattern in PICO_PATTERNS:
                        if (pattern.lower() in device.lower() or
                                pattern.lower() in mountpoint.lower()):
                            if os.path.ismount(mountpoint):
                                log(f"RP2350 mass storage mounted at: {mountpoint}")
                                return mountpoint
        except Exception as exc:
            log(f"Error reading /proc/mounts: {exc}")
        time.sleep(0.5)
    return None


# ------------------------------------------------------------------
# Flash
# ------------------------------------------------------------------

def flash_firmware(firmware_path: str, mount_point: str) -> bool:
    """Copy the firmware binary to the RP2350 mass-storage mount point."""
    if not os.path.exists(firmware_path):
        log(f"ERROR: Firmware file not found: {firmware_path}")
        return False
    if not os.path.ismount(mount_point):
        log(f"ERROR: {mount_point} is not a valid mount point")
        return False

    filename = os.path.basename(firmware_path)
    dest = os.path.join(mount_point, filename)
    try:
        log(f"Copying '{filename}' → {mount_point} ...")
        shutil.copy2(firmware_path, dest)
        subprocess.run(['sync'], check=True)
        log("Firmware copy complete.")
        return True
    except Exception as exc:
        log(f"ERROR copying firmware: {exc}")
        return False


# ------------------------------------------------------------------
# Main
# ------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="StellarMate Pro X RP2350 firmware updater"
    )
    parser.add_argument(
        '--port', default=None,
        help=f'UART port (default: auto-detect via {BY_ID_DIR}, '
             f'fallback {DEFAULT_UART_PORT})'
    )
    parser.add_argument(
        '--firmware', default=DEFAULT_FIRMWARE_PATH,
        help=f"Path to firmware .uf2 (default: {DEFAULT_FIRMWARE_PATH})"
    )
    args = parser.parse_args()

    # Auto-detect the serial port unless the user explicitly overrides it
    if args.port:
        port = args.port
    else:
        port = find_serial_port(DEFAULT_UART_PORT)

    log(f"Using UART port: {port}")
    log(f"Using firmware:  {args.firmware}")

    # ── 1. Check firmware file exists before touching GPIO ────────
    if not os.path.exists(args.firmware):
        log(f"ERROR: Firmware file not found: {args.firmware}")
        log("Cannot proceed with flash. Aborting.")
        sys.exit(1)

    # ── 2. Enter BOOTSEL mode ─────────────────────────────────────
    line_d2 = line_d3 = None
    try:
        log("Entering RP2350 USB BOOTSEL mode...")
        line_d2, line_d3 = enter_bootsel_mode()

        # ── 3. Wait for mass-storage mount ────────────────────────
        log(f"Waiting up to {MOUNT_WAIT_TIMEOUT}s for RP2350 mass storage...")
        mount_point = find_pico_mount_point(timeout=MOUNT_WAIT_TIMEOUT)

        if mount_point is None:
            log("ERROR: RP2350 did not appear as USB mass storage within timeout.")
            sys.exit(1)

        # ── 4. Copy firmware ──────────────────────────────────────
        success = flash_firmware(args.firmware, mount_point)
        if not success:
            log("Flash failed.")
            sys.exit(1)

        # ── 5. Wait for RP to process the new image ───────────────
        log(f"Waiting {POST_FLASH_WAIT}s for RP2350 to process firmware...")
        time.sleep(POST_FLASH_WAIT)

        # ── 6. Reset RP into normal operation ─────────────────────
        reset_rp(line_d2)
        time.sleep(2)

        # ── 7. Optional: verify new version over UART ────────────
        log("Verifying new firmware version via UART...")
        new_ver_str = query_rp_version(port, DEFAULT_BAUD_RATE, DEFAULT_UART_TIMEOUT * 2)
        if new_ver_str:
            log(f"RP2350 now reports version: {new_ver_str}")
        else:
            log("Note: Could not confirm version via UART (RP may still be booting).")

        log("Firmware update complete!")

    except KeyboardInterrupt:
        log("\nUpdate interrupted by user.")
        sys.exit(1)

    except Exception as exc:
        log(f"Unexpected error: {exc}")
        sys.exit(1)

    finally:
        release_gpio(line_d2, line_d3)


if __name__ == "__main__":
    main()