#!/usr/bin/env python3
"""
smprox_healthcheck.py — StellarMate Pro X RP2350 boot health check

Fast pre-start check for the INDI SM Pro driver:

  - Ping the RP2350 over UART (short timeout, few retries)
  - Responding  → exit 0 immediately (fast path, <1s typical)
  - Silent      → auto-recovery flash from the installed firmware file

No pacman version checks — normal updates are handled by the pacman
post-transaction hook. This script only guarantees the RP is alive.

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:
  smprox_healthcheck.py [--port /dev/ttyACM0] [--firmware /path/to/file.uf2]

Exit codes:
  0 = RP responding (or recovery flash succeeded)
  1 = recovery flash failed / firmware missing
"""

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

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

# ================= Configuration =================
DEFAULT_UART_PORT   = '/dev/ttyACM0'
BAUD_RATE           = 115200

# Fast-path tuning: keep these tight — this runs on every boot.
PORT_WAIT_TIMEOUT   = 5.0    # max seconds to wait for /dev/ttyACM* to enumerate
PING_TIMEOUT        = 1.0    # seconds per ping attempt
PING_RETRIES        = 3      # attempts before declaring the RP dead

DEFAULT_FIRMWARE_PATH = '/usr/lib/firmware/stellarmate/stellarmateprox.uf2'

# GPIO
CHIP_PATH   = '/dev/gpiochip0'
PIN_GPP_D2  = 194   # RP2350_RUN
PIN_GPP_D3  = 195   # RP2350_USB_BOOT

# BOOTSEL block-device detection (recovery path only)
BOOTSEL_LABELS      = ('RP2350', 'RPI-RP2')
BLOCKDEV_WAIT       = 15     # seconds to wait for BOOTSEL device
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,
    matching the same USB descriptor as the C++ findSerialPort().

    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


# ------------------------------------------------------------------
# Fast path: ping
# ------------------------------------------------------------------

def wait_for_port(port: str, timeout: float) -> bool:
    """Wait briefly for the UART device node to enumerate."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if os.path.exists(port):
            return True
        time.sleep(0.1)
    return False


def ping_rp(port: str) -> bool:
    """
    Single ping attempt: send get_firmware_info, accept ANY valid JSON
    response as proof of life. Returns True if RP responded.
    """
    command = json.dumps({
        "subsystem": "general",
        "command": "get_firmware_info"
    }) + '\n'

    try:
        with serial.Serial(port, BAUD_RATE, timeout=PING_TIMEOUT) as ser:
            ser.reset_input_buffer()
            ser.write(command.encode('utf-8'))

            deadline = time.monotonic() + PING_TIMEOUT
            while time.monotonic() < deadline:
                line = ser.readline().decode('utf-8', errors='replace').strip()
                if not line:
                    continue
                try:
                    json.loads(line)
                    return True          # any well-formed JSON = alive
                except json.JSONDecodeError:
                    continue
    except serial.SerialException as exc:
        log(f"UART error: {exc}")

    return False


def rp_is_alive(port: str) -> bool:
    for attempt in range(1, PING_RETRIES + 1):
        if ping_rp(port):
            return True
        log(f"Ping attempt {attempt}/{PING_RETRIES} failed")
    return False


# ------------------------------------------------------------------
# Recovery path: GPIO / BOOTSEL
# ------------------------------------------------------------------

def enter_bootsel_mode():
    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)})

    line_d2.set_value(PIN_GPP_D2, Value.ACTIVE)      # hold in reset
    time.sleep(0.5)
    line_d3.set_value(PIN_GPP_D3, Value.ACTIVE)      # select BOOTSEL
    time.sleep(0.5)
    line_d2.set_value(PIN_GPP_D2, Value.INACTIVE)    # boot into BOOTSEL
    time.sleep(0.5)
    line_d3.set_value(PIN_GPP_D3, Value.INACTIVE)    # release (latched)
    log("RP2350 booted into USB BOOTSEL mode")
    return line_d2, line_d3


def reset_rp(line_d2) -> None:
    line_d2.set_value(PIN_GPP_D2, Value.ACTIVE)
    time.sleep(0.5)
    line_d2.set_value(PIN_GPP_D2, Value.INACTIVE)
    log("RP2350 reset — booting normally")


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


# ------------------------------------------------------------------
# Recovery path: block device + self-mount
# ------------------------------------------------------------------

def find_bootsel_blockdev(timeout: int = BLOCKDEV_WAIT) -> str | None:
    """
    Find the BOOTSEL mass-storage partition via lsblk (no automounter
    required on headless boots). Returns /dev/xxx or None.
    """
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            result = subprocess.run(
                ['lsblk', '-o', 'NAME,LABEL,PATH', '-J'],
                capture_output=True, text=True, check=True)
            data = json.loads(result.stdout)

            def scan(devices):
                for dev in devices:
                    label = (dev.get('label') or '').upper()
                    if any(pat in label for pat in BOOTSEL_LABELS):
                        return dev.get('path')
                    found = scan(dev.get('children', []))
                    if found:
                        return found
                return None

            path = scan(data.get('blockdevices', []))
            if path:
                log(f"BOOTSEL block device found: {path}")
                return path
        except Exception as exc:
            log(f"lsblk error: {exc}")
        time.sleep(0.5)
    return None


def flash_firmware(firmware_path: str, blockdev: str) -> bool:
    """Mount the BOOTSEL device ourselves, copy the UF2, sync, unmount."""
    mount_dir = tempfile.mkdtemp(prefix='smprox_bootsel_')
    mounted = False
    try:
        subprocess.run(['mount', blockdev, mount_dir], check=True)
        mounted = True
        dest = os.path.join(mount_dir, os.path.basename(firmware_path))
        log(f"Copying {os.path.basename(firmware_path)} → {blockdev} ...")
        shutil.copy2(firmware_path, dest)
        subprocess.run(['sync'], check=True)
        log("Firmware copy complete")
        return True
    except Exception as exc:
        log(f"ERROR flashing firmware: {exc}")
        return False
    finally:
        if mounted:
            # RP may yank the device the instant the UF2 lands; lazy unmount
            subprocess.run(['umount', '-l', mount_dir], check=False)
        try:
            os.rmdir(mount_dir)
        except OSError:
            pass


def recover(firmware_path: str, port: str) -> bool:
    line_d2 = line_d3 = None
    try:
        line_d2, line_d3 = enter_bootsel_mode()

        blockdev = find_bootsel_blockdev()
        if blockdev is None:
            log("ERROR: RP2350 did not enumerate as USB mass storage")
            return False

        # Give the kernel a moment to settle the partition
        time.sleep(0.5)

        if not flash_firmware(firmware_path, blockdev):
            return False

        time.sleep(POST_FLASH_WAIT)
        reset_rp(line_d2)
        time.sleep(2)

        # Verify: port should re-enumerate and answer a ping
        if wait_for_port(port, PORT_WAIT_TIMEOUT) and rp_is_alive(port):
            log("Recovery successful — RP2350 responding")
            return True

        log("WARNING: RP2350 flashed but not yet responding")
        return False
    finally:
        release_gpio(line_d2, line_d3)


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

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Fast RP2350 health check + auto-recovery for SM Pro X")
    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)
    args = parser.parse_args()

    # Auto-detect the serial port unless the user explicitly overrides it.
    # If --port was given on the command line, use it verbatim.
    if args.port:
        port = args.port
    else:
        port = find_serial_port(DEFAULT_UART_PORT)

    start = time.monotonic()

    # ── Fast path ──────────────────────────────────────────────
    if not wait_for_port(port, PORT_WAIT_TIMEOUT):
        log(f"WARNING: {port} never enumerated — attempting recovery")
    elif rp_is_alive(port):
        log(f"RP2350 alive ({time.monotonic() - start:.2f}s). OK.")
        sys.exit(0)
    else:
        log("RP2350 not responding — attempting recovery")

    # ── Recovery path ──────────────────────────────────────────
    if not os.path.exists(args.firmware):
        log(f"ERROR: Firmware file not found: {args.firmware}")
        sys.exit(1)

    if recover(args.firmware, port):
        log(f"Done in {time.monotonic() - start:.1f}s")
        sys.exit(0)

    sys.exit(1)


if __name__ == "__main__":
    main()