#!/usr/bin/env python3
"""Build and OTA-update the full RFP Infinity installation.

This script intentionally delegates flashing to tools/rfp_network_flash.py so
the actual OTA path stays the existing WLED /update workflow.
"""

from __future__ import annotations

import argparse
import ipaddress
import os
import subprocess
import sys
from pathlib import Path


NODE_ENV = "rfp_esp32s3_wroom1_n16r8_3x106"
MASTER_ENV = "rfp_esp32s3_wroom1_n16r8_master"
NODE_RELEASE = "RFP_N16R8_NODE3x106_V20260511E"
MASTER_RELEASE = "RFP_N16R8_MASTER_V20260511E"
NODE_FIRMWARE_FALLBACK = Path(".pio/build") / NODE_ENV / "firmware.bin"
MASTER_FIRMWARE_FALLBACK = Path(".pio/build") / MASTER_ENV / "firmware.bin"
DEFAULT_SUBNET = "192.168.178.0/24"
NODE_HOSTS = range(11, 17)
MASTER_HOST = 10


def repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def python_executable(root: Path) -> str:
    local_python = root / ".venv/bin/python"
    return str(local_python) if local_python.exists() else sys.executable


def platformio_env(root: Path) -> dict[str, str]:
    env = os.environ.copy()
    node_bin = "/home/jan/Documents/RFP/Finanz_App/node/current/bin"
    env["PATH"] = f"{node_bin}:{env.get('PATH', '')}"
    env["NPM_CONFIG_CACHE"] = str(root / ".npm-cache")
    env["PLATFORMIO_CORE_DIR"] = str(root / ".piohome")
    env["PLATFORMIO_PACKAGES_DIR"] = str(root / ".piohome/packages")
    env["PLATFORMIO_PLATFORMS_DIR"] = str(root / ".piohome/platforms")
    env["PLATFORMIO_CACHE_DIR"] = str(root / ".piohome/.cache")
    env["PLATFORMIO_BUILD_CACHE_DIR"] = str(root / ".piohome/buildcache")
    return env


def targets_from_subnet(subnet: str) -> tuple[list[str], str]:
    network = ipaddress.ip_network(subnet, strict=False)
    if network.version != 4:
        raise ValueError("Only IPv4 subnets are supported.")
    octets = str(network.network_address).split(".")
    if len(octets) != 4:
        raise ValueError(f"Invalid IPv4 subnet: {subnet}")
    prefix = ".".join(octets[:3])
    return [f"{prefix}.{host}" for host in NODE_HOSTS], f"{prefix}.{MASTER_HOST}"


def run(cmd: list[str], root: Path, env: dict[str, str], dry_run: bool) -> None:
    print("+ " + " ".join(cmd))
    if dry_run:
        return
    subprocess.run(cmd, cwd=root, env=env, check=True)


def build(root: Path, env_name: str, env: dict[str, str], dry_run: bool) -> None:
    run([python_executable(root), "-m", "platformio", "run", "-e", env_name], root, env, dry_run)


def release_firmware(root: Path, expected_release: str, fallback: Path) -> Path:
    release_dir = root / "build_output/release"
    matches = sorted(release_dir.glob(f"WLEDMM_*_{expected_release}.bin"))
    if len(matches) == 1:
        return matches[0].relative_to(root)
    if len(matches) > 1:
        names = ", ".join(str(path.relative_to(root)) for path in matches)
        raise RuntimeError(f"Multiple release firmware files match {expected_release}: {names}")
    return fallback


def flash_one(
    root: Path,
    target: str,
    firmware: Path,
    expected_release: str,
    env: dict[str, str],
    args: argparse.Namespace,
) -> None:
    firmware_path = root / firmware
    if not args.dry_run and not firmware_path.exists():
        raise FileNotFoundError(f"Firmware does not exist: {firmware_path}")

    cmd = [
        python_executable(root),
        "tools/rfp_network_flash.py",
        "flash",
        "--targets",
        target,
        "--firmware",
        str(firmware),
        "--expect-release",
        expected_release,
    ]
    if args.preflight_only:
        cmd.append("--preflight-only")
    if args.skip_validation and not args.preflight_only:
        cmd.append("--skip-validation")
    if args.force_reboot_after_upload:
        cmd.append("--force-reboot-after-upload")
    if args.ignore_preflight_warnings:
        cmd.append("--ignore-preflight-warnings")
    if args.skip_current_release and not args.preflight_only:
        cmd.append("--skip-if-release-matches")
    if args.upload_timeout is not None:
        cmd += ["--upload-timeout", str(args.upload_timeout)]
    run(cmd, root, env, args.dry_run)


def flash(
    root: Path,
    targets: list[str],
    firmware: Path,
    expected_release: str,
    env: dict[str, str],
    args: argparse.Namespace,
) -> None:
    if not targets:
        return
    for target in targets:
        try:
            flash_one(root, target, firmware, expected_release, env, args)
        except subprocess.CalledProcessError as exc:
            print()
            print(f"OTA failed at {target}; stopping before continuing to the remaining devices.")
            print("After fixing the cause, resume with:")
            print(f"  {python_executable(root)} tools/rfp_update_all_ota.py --no-build --start-from {target}")
            raise exc


def flash_nodes_via_master_usb(
    root: Path,
    targets: list[str],
    firmware: Path,
    expected_release: str,
    env: dict[str, str],
    args: argparse.Namespace,
) -> None:
    if not targets:
        return
    cmd = [
        python_executable(root),
        "tools/rfp_master_usb_relay.py",
        "--port",
        args.port,
        "--firmware",
        str(firmware),
        "--expect-release",
        expected_release,
        "--targets",
        ",".join(targets),
        "--baud",
        str(args.relay_baud),
        "--startup-delay",
        str(args.relay_startup_delay),
    ]
    if args.start_from:
        cmd += ["--start-from", args.start_from]
    if not args.skip_current_release:
        cmd.append("--force-current-release")
    run(cmd, root, env, args.dry_run)


def filtered_plan(nodes: list[str], master: str, args: argparse.Namespace) -> list[tuple[str, str]]:
    plan: list[tuple[str, str]] = []
    if not args.master_only:
        plan.extend((ip, "node") for ip in nodes)
    if not args.nodes_only:
        plan.append((master, "master"))

    if not args.start_from:
        return plan

    for index, (ip, _role) in enumerate(plan):
        if ip == args.start_from:
            return plan[index:]
    raise ValueError(f"--start-from {args.start_from} is not in the selected update plan.")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Build and OTA-update RFP Infinity nodes, then master.")
    role = parser.add_mutually_exclusive_group()
    role.add_argument("--nodes-only", action="store_true", help="Only build/flash nodes .11-.16")
    role.add_argument("--master-only", action="store_true", help="Only build/flash master .10")
    parser.add_argument("--start-from", help="Resume at this IP, for example 192.168.178.14 or 192.168.178.10")
    parser.add_argument("--subnet", default=DEFAULT_SUBNET, help=f"Show subnet used to derive .10-.16 targets (default: {DEFAULT_SUBNET})")
    parser.add_argument("--no-build", action="store_true", help="Skip PlatformIO builds and flash existing firmware.bin files")
    parser.add_argument("--dry-run", action="store_true", help="Print build and flash steps without executing them")
    parser.add_argument("--preflight-only", action="store_true", help="Only run per-device OTA diagnostics, do not upload")
    parser.add_argument("--via-master-usb", action="store_true", help="Update nodes through the USB-connected master instead of direct PC-to-node HTTP")
    parser.add_argument("--port", default="/dev/ttyACM0", help="Master serial port for --via-master-usb")
    parser.add_argument("--relay-baud", type=int, default=921600, help="Master USB relay baud rate (default: 921600)")
    parser.add_argument("--relay-startup-delay", type=float, default=4.0, help="Seconds to wait after opening master USB serial (default: 4)")
    parser.add_argument("--ignore-preflight-warnings", action="store_true", help="Try OTA even if /update preflight looks blocked")
    parser.add_argument("--force-reboot-after-upload", action="store_true", help="Request a reboot if upload response is uncertain and no reboot is observed")
    parser.add_argument("--force-current-release", dest="skip_current_release", action="store_false", help="Reflash devices even when their current release already matches the selected firmware")
    parser.add_argument("--upload-timeout", type=float, default=180.0, help="Per-device OTA upload timeout in seconds (default: 180)")
    parser.add_argument("--no-skip-validation", dest="skip_validation", action="store_false", help="Do not send WLED skipValidation=1 during upload")
    parser.set_defaults(skip_validation=True, skip_current_release=True)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    root = repo_root()
    env = platformio_env(root)
    nodes, master = targets_from_subnet(args.subnet)
    plan = filtered_plan(nodes, master, args)
    node_targets = [ip for ip, role in plan if role == "node"]
    master_targets = [ip for ip, role in plan if role == "master"]

    if not plan:
        print("Nothing selected.")
        return 0

    print("RFP Infinity OTA update plan:")
    for ip, role in plan:
        print(f"- {ip} ({role})")
    print()

    if not args.no_build:
        if node_targets:
            print(f"Building node firmware: {NODE_ENV}")
            build(root, NODE_ENV, env, args.dry_run)
        if master_targets:
            print(f"Building master firmware: {MASTER_ENV}")
            build(root, MASTER_ENV, env, args.dry_run)

    node_firmware = release_firmware(root, NODE_RELEASE, NODE_FIRMWARE_FALLBACK)
    master_firmware = release_firmware(root, MASTER_RELEASE, MASTER_FIRMWARE_FALLBACK)
    print(f"Node firmware:   {node_firmware} (expect release {NODE_RELEASE})")
    print(f"Master firmware: {master_firmware} (expect release {MASTER_RELEASE})")
    if args.skip_validation and not args.preflight_only:
        print("WLED OTA release validation: skipValidation=1 for controlled RFP release-name migration")
    if args.skip_current_release and not args.preflight_only:
        print("Already matching devices: skipped by default; use --force-current-release to reflash them")
    print()

    if node_targets:
        if args.via_master_usb:
            if args.preflight_only:
                print("USB master relay preflight is not separate; run without --preflight-only to query through the master.")
            else:
                print("Flashing nodes sequentially via USB-connected master...")
                flash_nodes_via_master_usb(root, node_targets, node_firmware, NODE_RELEASE, env, args)
        else:
            print(("Preflighting" if args.preflight_only else "Flashing") + " nodes sequentially...")
            flash(root, node_targets, node_firmware, NODE_RELEASE, env, args)
    if master_targets:
        if args.via_master_usb:
            print(
                "Skipping master firmware in --via-master-usb mode. "
                "For full updates, use tools/rfp_update_master_usb_then_nodes.py "
                "to flash the master by USB first and then relay node OTA."
            )
        else:
            print(("Preflighting" if args.preflight_only else "Flashing") + " master last...")
            flash(root, master_targets, master_firmware, MASTER_RELEASE, env, args)

    print("OTA update plan completed.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())