#!/usr/bin/env python3
"""
ArduPilot custom build tool.
Interactive menu to configure and build firmware for selected boards.

Note: build/ is symlinked to a tmpfs RAM disk (/mnt/ramdisk) to reduce
SSD wear.  Firmware artifacts (.bin, .hex, .apj) are exported to
FIRMWARE_EXPORT_DIR on the real filesystem after each board build.
"""

import subprocess
import sys
import os
import re
import shutil
import datetime
from pathlib import Path

# --- Configuration ---
BOARDS_F4_FILE = "my_boards_f4.txt"     # flash-constrained boards
BOARDS_H7_FILE = "my_boards_h7.txt"     # large-flash boards
EXTRA_HWDEF_F4 = "extra_hwdef_f4.dat"      # strips features for F4
EXTRA_HWDEF_H7 = "extra_hwdef_h7.dat"   # adds features for H7/F7
LOG_DIR = "build_logs"
FIRMWARE_EXPORT_DIR = "firmware_export"

# --- Colors ---
class C:
    BOLD    = "\033[1m"
    RED     = "\033[91m"
    GREEN   = "\033[92m"
    YELLOW  = "\033[93m"
    CYAN    = "\033[96m"
    DIM     = "\033[2m"
    MAGENTA = "\033[95m"
    RESET   = "\033[0m"


def run_live(cmd: list[str], board_log=None) -> tuple[int, str]:
    """Run a command, streaming output to screen and board log file.
    Returns (returncode, combined_output)."""
    output_lines = []
    proc = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
    )
    for line in proc.stdout:
        sys.stdout.write(line)
        sys.stdout.flush()
        output_lines.append(line)
        if board_log:
            board_log.write(line)
            board_log.flush()
    proc.wait()
    return proc.returncode, "".join(output_lines)


def get_all_boards() -> list[str]:
    """Get list of all available boards from waf."""
    proc = subprocess.run(["./waf", "list_boards"], capture_output=True, text=True)
    if proc.returncode != 0:
        print(f"{C.RED}Failed to list boards{C.RESET}")
        sys.exit(1)
    boards = []
    for line in proc.stdout.splitlines():
        if line.startswith("'") or line.startswith("Waf:"):
            continue
        boards.extend(line.split())
    return sorted(boards, key=str.lower)


def _read_board_file(path: str) -> list[tuple[str, str | None]]:
    """Read boards from a file, deduplicating.
    Supports format: BoardName or BoardName:custom_hwdef.dat
    Returns list of (board_name, custom_hwdef_path_or_None) tuples."""
    p = Path(path)
    if not p.exists():
        return []
    boards = []
    seen = set()
    for line in p.read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#") and line not in seen:
            seen.add(line)
            if ":" in line:
                board, custom = line.split(":", 1)
                boards.append((board.strip(), custom.strip()))
            else:
                boards.append((line, None))
    return boards


def get_favourite_boards() -> list[tuple[str, str | None, str | None]]:
    """Read favourite boards from F4 and H7 files.
    Returns list of (board_name, extra_hwdef_path, custom_hwdef_path) tuples."""
    result = []
    f4 = _read_board_file(BOARDS_F4_FILE)
    h7 = _read_board_file(BOARDS_H7_FILE)
    if not f4 and not h7:
        print(f"{C.YELLOW}Warning: no boards found in {BOARDS_F4_FILE} or {BOARDS_H7_FILE}{C.RESET}")
        return []
    for board, custom in f4:
        result.append((board, EXTRA_HWDEF_F4, custom))
    for board, custom in h7:
        result.append((board, EXTRA_HWDEF_H7, custom))
    return result


def get_favourite_board_names() -> list[str]:
    """Just the names, for listing."""
    return [b for b, _ in get_favourite_boards()]


def search_boards(all_boards: list[str]) -> list[str]:
    """Interactive search/filter of boards."""
    while True:
        query = input(f"\n{C.CYAN}Search pattern (regex ok, empty=show all): {C.RESET}").strip()
        if not query:
            for i, b in enumerate(all_boards, 1):
                print(f"  {i:3d}. {b}")
            print(f"\n  {C.DIM}{len(all_boards)} boards total{C.RESET}")
            continue

        try:
            pattern = re.compile(query, re.IGNORECASE)
        except re.error as e:
            print(f"{C.RED}Bad regex: {e}{C.RESET}")
            continue

        matches = [b for b in all_boards if pattern.search(b)]
        if not matches:
            print(f"{C.YELLOW}No matches for '{query}'{C.RESET}")
            continue

        for i, b in enumerate(matches, 1):
            print(f"  {i:3d}. {b}")
        print(f"\n  {C.DIM}{len(matches)} matches{C.RESET}")

        sel = input(f"{C.CYAN}Enter numbers to build (comma-sep), [a]ll matches, or [b]ack: {C.RESET}").strip()
        if sel.lower() == "b":
            continue
        if sel.lower() == "a":
            return matches
        try:
            indices = [int(x.strip()) - 1 for x in sel.split(",")]
            selected = [matches[i] for i in indices if 0 <= i < len(matches)]
            if selected:
                return selected
            print(f"{C.YELLOW}No valid selection{C.RESET}")
        except (ValueError, IndexError):
            print(f"{C.RED}Invalid input{C.RESET}")


def pick_one(all_boards: list[str]) -> list[str]:
    """Let user pick a single board by search."""
    while True:
        query = input(f"\n{C.CYAN}Board name (partial, regex ok): {C.RESET}").strip()
        if not query:
            return []
        try:
            pattern = re.compile(query, re.IGNORECASE)
        except re.error as e:
            print(f"{C.RED}Bad regex: {e}{C.RESET}")
            continue

        matches = [b for b in all_boards if pattern.search(b)]
        if not matches:
            print(f"{C.YELLOW}No matches{C.RESET}")
            continue
        if len(matches) == 1:
            print(f"  -> {matches[0]}")
            return matches

        for i, b in enumerate(matches, 1):
            print(f"  {i:3d}. {b}")
        sel = input(f"{C.CYAN}Pick number (or Enter to re-search): {C.RESET}").strip()
        if not sel:
            continue
        try:
            idx = int(sel) - 1
            if 0 <= idx < len(matches):
                return [matches[idx]]
        except ValueError:
            pass
        print(f"{C.RED}Invalid{C.RESET}")


def pick_vehicles(bootloader_mode: bool = False) -> tuple[list[str] | None, bool]:
    """Let user choose which vehicles to build.
    Returns (vehicle_list, build_bootloader).
    vehicle_list=None means 'build all' (./waf with no target).
    vehicle_list=[] means 'no firmware' (bootloader only)."""
    if bootloader_mode:
        return [], True

    print(f"\n{C.BOLD}Vehicles to build:{C.RESET}")
    print(f"  1. plane + copter        {C.DIM}(default){C.RESET}")
    print(f"  2. plane only")
    print(f"  3. copter only")
    print(f"  4. all (waf default)")
    print(f"  5. bootloader only")
    print(f"  6. plane + copter + bootloader")
    sel = input(f"{C.CYAN}Choice [1]: {C.RESET}").strip()
    if sel == "2":
        return ["plane"], False
    if sel == "3":
        return ["copter"], False
    if sel == "4":
        return None, False
    if sel == "5":
        return [], True
    if sel == "6":
        return ["plane", "copter"], True
    return ["plane", "copter"], False


def build_boards(boards: list[tuple[str, str | None, str | None]], vehicles: list[str] | None,
                 build_bootloader: bool = False):
    """Configure and build each board with live output.
    boards: list of (board_name, extra_hwdef_path_or_None, custom_hwdef_path_or_None) tuples.
    Logs: batch_dir/summary.log  + batch_dir/<board>.log per board.
    vehicles=None  -> ./waf (all targets)
    vehicles=[]    -> no firmware (bootloader only)
    vehicles=[...] -> specific targets"""
    if not boards:
        print(f"{C.YELLOW}No boards to build.{C.RESET}")
        return

    build_firmware = vehicles is None or len(vehicles) > 0

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    batch_dir = Path(LOG_DIR) / timestamp
    batch_dir.mkdir(parents=True, exist_ok=True)
    summary_path = batch_dir / "summary.log"

    export_dir = Path(FIRMWARE_EXPORT_DIR) / timestamp

    failed: list[tuple[str, str, str]] = []  # (board, stage, reason)
    succeeded: list[str] = []
    all_exported: list[str] = []

    total = len(boards)
    parts = []
    if vehicles is None:
        parts.append("all")
    elif vehicles:
        parts.append("+".join(vehicles))
    if build_bootloader:
        parts.append("bootloader")
    vehicle_label = "+".join(parts) if parts else "bootloader"

    print(f"\n{C.BOLD}Building {total} board(s) [{vehicle_label}]{C.RESET}")
    print(f"{C.DIM}Logs: {batch_dir}/{C.RESET}\n")

    with open(summary_path, "w") as summary:
        summary.write(f"ArduPilot Build - {timestamp}\n")
        summary.write(f"Vehicles: {vehicle_label}\n")
        summary.write(f"Boards: {total}\n")
        summary.write("=" * 70 + "\n\n")
        summary.write(f"{'Board':<40} {'Stage':<15} {'Result'}\n")
        summary.write(f"{'-'*40} {'-'*15} {'-'*40}\n")

        for idx, (board, extra_hwdef, custom_hwdef) in enumerate(boards, 1):
            header = f"[{idx}/{total}] {board}"
            if extra_hwdef:
                header += f"  ({Path(extra_hwdef).stem})"
            if custom_hwdef:
                header += f"  +custom({Path(custom_hwdef).stem})"
            separator = "=" * 70
            print(f"\n{C.BOLD}{separator}{C.RESET}")
            print(f"{C.BOLD}{header}{C.RESET}")
            print(f"{C.BOLD}{separator}{C.RESET}")

            board_log_path = batch_dir / f"{board}.log"
            board_ok = True
            temp_hwdef = None  # Track temp hwdef file for cleanup

            with open(board_log_path, "w") as blog:
                blog.write(f"Board: {board}\n")
                blog.write(f"Date: {timestamp}\n")
                blog.write(f"Vehicles: {vehicle_label}\n")
                blog.write(f"Extra hwdef: {extra_hwdef or 'none'}\n")
                blog.write(f"Custom hwdef: {custom_hwdef or 'none'}\n")
                blog.write("=" * 70 + "\n\n")

                # --- Build firmware ---
                if build_firmware:
                    cfg_cmd = ["./waf", "configure", "--board", board]
                    
                    # Merge extra_hwdef and custom_hwdef if both exist
                    # (waf only uses the last --extra-hwdef, so we must combine them)
                    if extra_hwdef and custom_hwdef and Path(extra_hwdef).exists() and Path(custom_hwdef).exists():
                        temp_hwdef = Path(f".temp_hwdef_{board}.dat")
                        with open(temp_hwdef, "w") as tf:
                            tf.write(f"# Merged hwdef: {extra_hwdef} + {custom_hwdef}\n\n")
                            tf.write(f"# ===== From {extra_hwdef} =====\n")
                            tf.write(Path(extra_hwdef).read_text())
                            tf.write(f"\n\n# ===== From {custom_hwdef} =====\n")
                            tf.write(Path(custom_hwdef).read_text())
                        cfg_cmd += [f"--extra-hwdef={temp_hwdef}"]
                        blog.write(f"Created merged hwdef: {temp_hwdef}\n")
                    elif extra_hwdef and Path(extra_hwdef).exists():
                        cfg_cmd += [f"--extra-hwdef={extra_hwdef}"]
                    elif custom_hwdef and Path(custom_hwdef).exists():
                        cfg_cmd += [f"--extra-hwdef={custom_hwdef}"]

                    print(f"\n{C.CYAN}$ {' '.join(cfg_cmd)}{C.RESET}")
                    blog.write(f"$ {' '.join(cfg_cmd)}\n")

                    rc, output = run_live(cfg_cmd, blog)

                    if rc != 0:
                        reason = _extract_error(output)
                        failed.append((board, "configure", reason))
                        print(f"\n  {C.RED}>>> CONFIGURE FAILED{C.RESET}")
                        blog.write(f"\n>>> CONFIGURE FAILED: {reason}\n")
                        summary.write(f"{board:<40} {'configure':<15} FAILED: {reason}\n")
                        board_ok = False
                    else:
                        print(f"\n  {C.GREEN}>>> configured OK{C.RESET}")
                        blog.write(f"\n>>> CONFIGURE OK\n")
                        summary.write(f"{board:<40} {'configure':<15} OK\n")

                        # Build each vehicle target
                        targets = vehicles if vehicles else [None]
                        for vehicle in targets:
                            build_cmd = ["./waf"]
                            if vehicle:
                                build_cmd.append(vehicle)
                            vlabel = vehicle or "build"

                            print(f"\n{C.CYAN}$ {' '.join(build_cmd)}{C.RESET}")
                            blog.write(f"\n$ {' '.join(build_cmd)}\n")

                            rc, output = run_live(build_cmd, blog)

                            if rc != 0:
                                reason = _extract_error(output)
                                failed.append((board, vlabel, reason))
                                print(f"\n  {C.RED}>>> {vlabel} FAILED{C.RESET}")
                                blog.write(f"\n>>> {vlabel.upper()} FAILED: {reason}\n")
                                summary.write(f"{board:<40} {vlabel:<15} FAILED: {reason}\n")
                                board_ok = False
                            else:
                                print(f"\n  {C.GREEN}>>> {vlabel} OK{C.RESET}")
                                blog.write(f"\n>>> {vlabel.upper()} OK\n")
                                summary.write(f"{board:<40} {vlabel:<15} OK\n")

                # --- Export firmware artifacts ---
                if build_firmware:
                    exported = _export_firmware(board, export_dir, custom_hwdef, blog)
                    if exported:
                        all_exported.extend(f"{board}/{n}" for n in exported)
                        summary.write(f"{board:<40} {'export':<15} {len(exported)} file(s)\n")

                # --- Build bootloader ---
                if build_bootloader:
                    bl_ok = _build_bootloader(board, blog)
                    if bl_ok:
                        summary.write(f"{board:<40} {'bootloader':<15} OK\n")
                    else:
                        reason = "see board log"
                        failed.append((board, "bootloader", reason))
                        summary.write(f"{board:<40} {'bootloader':<15} FAILED\n")
                        board_ok = False

                if board_ok:
                    succeeded.append(board)
                
                # --- Cleanup temp hwdef file ---
                if temp_hwdef and temp_hwdef.exists():
                    temp_hwdef.unlink()
                    blog.write(f"Cleaned up temp hwdef: {temp_hwdef}\n")

                # --- Wipe build/ contents to free ramdisk & avoid stale state ---
                build_dir = Path("build")
                if build_dir.is_dir():
                    freed = 0
                    for item in list(build_dir.iterdir()):
                        if item.is_dir():
                            freed += sum(f.stat().st_size for f in item.rglob("*") if f.is_file())
                            shutil.rmtree(item)
                        elif item.is_file():
                            freed += item.stat().st_size
                            item.unlink()
                    freed_mb = freed / (1024 * 1024)
                    msg = f"Wiped build/ ({freed_mb:.0f} MB freed)"
                    print(f"  {C.DIM}{msg}{C.RESET}")
                    blog.write(f">>> {msg}\n")

            summary.flush()

        # Final summary in log
        summary.write(f"\n{'=' * 70}\n")
        summary.write(f"TOTAL: {len(succeeded)} OK, {len(failed)} FAILED out of {total} boards\n")
        if all_exported:
            summary.write(f"Exported: {len(all_exported)} files to {export_dir}/\n")
        summary.write(f"{'=' * 70}\n")
        if failed:
            summary.write(f"\nFailed details:\n")
            for board, stage, reason in failed:
                summary.write(f"  {board}:{stage} -> {reason}\n")

    # Print final summary to screen
    _print_summary(succeeded, failed, total, batch_dir, export_dir if all_exported else None)


def _export_firmware(board: str, export_dir: Path, custom_hwdef: str | None, blog) -> list[str]:
    """Copy .bin/.hex/.apj artifacts from build/{board}/bin/ to export_dir/{board}/.
    If custom_hwdef is provided, rename files to include the aircraft identifier.
    Returns list of copied filenames."""
    src = Path(f"build/{board}/bin")
    if not src.is_dir():
        return []
    extensions = {".bin", ".hex", ".apj"}
    files = [f for f in src.iterdir() if f.is_file() and f.suffix in extensions]
    if not files:
        return []
    dst = export_dir / board
    dst.mkdir(parents=True, exist_ok=True)
    
    # Extract custom identifier from hwdef filename
    custom_id = None
    if custom_hwdef:
        # Extract from pattern: custom_hwdef_<identifier>_<board>.dat or custom_hwdef_<identifier>.dat
        match = re.match(r"custom_hwdef_([^_]+)(?:_.*?)?\.dat$", Path(custom_hwdef).name)
        if match:
            custom_id = match.group(1)
    
    copied = []
    for f in sorted(files):
        # Rename file to include custom identifier if present
        if custom_id:
            # Insert custom_id before extension
            # e.g., arduplane.apj -> arduplane-minitalon.apj
            new_name = f"{f.stem}-{custom_id}{f.suffix}"
        else:
            new_name = f.name
        
        shutil.copy2(f, dst / new_name)
        copied.append(new_name)
    msg = f"Exported {len(copied)} file(s) to {dst}/"
    print(f"  {C.DIM}{msg}{C.RESET}")
    blog.write(f">>> {msg}\n")
    for name in copied:
        blog.write(f"    {name}\n")
    return copied


def _build_bootloader(board: str, blog) -> bool:
    """Configure and build bootloader for a board, copy artifacts.
    Note: extra_hwdef is NOT used for bootloader builds."""
    cfg_cmd = ["./waf", "configure", "--board", board, "--bootloader"]

    print(f"\n{C.CYAN}$ {' '.join(cfg_cmd)}{C.RESET}")
    blog.write(f"\n$ {' '.join(cfg_cmd)}\n")

    rc, output = run_live(cfg_cmd, blog)
    if rc != 0:
        reason = _extract_error(output)
        print(f"\n  {C.RED}>>> bootloader configure FAILED{C.RESET}")
        blog.write(f"\n>>> BOOTLOADER CONFIGURE FAILED: {reason}\n")
        return False

    print(f"\n  {C.GREEN}>>> bootloader configured{C.RESET}")
    blog.write(f"\n>>> BOOTLOADER CONFIGURE OK\n")

    build_cmd = ["./waf", "bootloader"]
    print(f"\n{C.CYAN}$ {' '.join(build_cmd)}{C.RESET}")
    blog.write(f"\n$ {' '.join(build_cmd)}\n")

    rc, output = run_live(build_cmd, blog)
    if rc != 0:
        reason = _extract_error(output)
        print(f"\n  {C.RED}>>> bootloader FAILED{C.RESET}")
        blog.write(f"\n>>> BOOTLOADER BUILD FAILED: {reason}\n")
        return False

    # Copy artifacts to Tools/bootloaders/
    bl_dir = Path("Tools/bootloaders")
    bl_dir.mkdir(parents=True, exist_ok=True)
    build_dir = Path(f"build/{board}")

    copied = []
    elf_src = build_dir / "bootloader" / "AP_Bootloader"
    bin_dst = bl_dir / f"{board}_bl.bin"
    hex_dst = bl_dir / f"{board}_bl.hex"

    if elf_src.exists():
        try:
            subprocess.run(["arm-none-eabi-objcopy", "-O", "binary",
                           str(elf_src), str(bin_dst)], check=True,
                          capture_output=True, text=True)
            copied.append(bin_dst.name)
        except (subprocess.CalledProcessError, FileNotFoundError):
            for f in build_dir.rglob("*_bl.bin"):
                shutil.copy2(f, bl_dir / f.name)
                copied.append(f.name)
        try:
            subprocess.run(["arm-none-eabi-objcopy", "-O", "ihex",
                           str(elf_src), str(hex_dst)], check=True,
                          capture_output=True, text=True)
            copied.append(hex_dst.name)
        except (subprocess.CalledProcessError, FileNotFoundError):
            pass
    else:
        for f in build_dir.rglob("*.bin"):
            if "bootloader" in str(f).lower() or "_bl" in f.name:
                shutil.copy2(f, bl_dir / f.name)
                copied.append(f.name)

    if copied:
        print(f"\n  {C.GREEN}>>> bootloader OK{C.RESET} {C.DIM}-> {', '.join(copied)}{C.RESET}")
        blog.write(f"\n>>> BOOTLOADER OK: {', '.join(copied)}\n")
    else:
        print(f"\n  {C.GREEN}>>> bootloader OK{C.RESET} {C.DIM}(no artifacts copied){C.RESET}")
        blog.write(f"\n>>> BOOTLOADER OK (no artifacts found to copy)\n")

    return True


def _extract_error(output: str) -> str:
    """Extract a concise error reason from build output."""
    lines = output.splitlines()

    error_line = ""
    for line in lines:
        if "region" in line and "overflowed" in line:
            return line.strip()
        if ": error:" in line:
            error_line = line.strip()
        if "undefined reference" in line:
            error_line = line.strip()
        if "Build failed" in line or "Configuration failed" in line:
            if not error_line:
                error_line = line.strip()

    if error_line:
        if len(error_line) > 120:
            return error_line[:117] + "..."
        return error_line

    for line in reversed(lines):
        if line.strip():
            s = line.strip()
            return s[:120] if len(s) > 120 else s
    return "unknown error"


def _print_summary(succeeded: list[str], failed: list[tuple[str, str, str]],
                   total: int, batch_dir: Path, export_dir: Path | None = None):
    """Print final build summary to screen."""
    print(f"\n{'=' * 70}")
    if len(failed) == 0:
        print(f"{C.GREEN}{C.BOLD}SUMMARY: {len(succeeded)} OK, 0 FAILED out of {total} boards{C.RESET}")
    else:
        print(f"{C.RED}{C.BOLD}SUMMARY: {len(succeeded)} OK, {len(failed)} FAILED out of {total} boards{C.RESET}")
    print(f"{'=' * 70}")

    if succeeded:
        print(f"\n  {C.GREEN}Succeeded ({len(succeeded)}):{C.RESET}")
        for b in succeeded:
            print(f"    {b}")

    if failed:
        print(f"\n  {C.RED}Failed ({len(failed)}):{C.RESET}")
        for board, stage, reason in failed:
            print(f"    {board}:{stage}")
            print(f"      {C.DIM}-> {reason}{C.RESET}")

    print(f"\n  {C.DIM}Logs: {batch_dir}/{C.RESET}")
    print(f"  {C.DIM}Summary: {batch_dir}/summary.log{C.RESET}")
    if export_dir:
        print(f"  {C.DIM}Firmware: {export_dir}/{C.RESET}")


def _get_custom_hwdef_files() -> list[str]:
    """Get list of custom_hwdef_*.dat files in current directory."""
    return sorted([f.name for f in Path(".").glob("custom_hwdef_*.dat")])


def _pick_hwdef() -> str | None:
    """Let user choose which extra hwdef to use for ad-hoc board selection."""
    print(f"\n{C.BOLD}Extra hwdef:{C.RESET}")
    print(f"  1. {EXTRA_HWDEF_F4}          {C.DIM}(F4 - strip features){C.RESET}")
    print(f"  2. {EXTRA_HWDEF_H7}      {C.DIM}(H7/F7 - add features){C.RESET}")
    print(f"  3. none")
    sel = input(f"{C.CYAN}Choice [1]: {C.RESET}").strip()
    if sel == "2":
        return EXTRA_HWDEF_H7
    if sel == "3":
        return None
    return EXTRA_HWDEF_F4


def _pick_custom_hwdef() -> str | None:
    """Let user choose aircraft-specific custom hwdef for ad-hoc board selection."""
    custom_files = _get_custom_hwdef_files()
    if not custom_files:
        return None
    
    print(f"\n{C.BOLD}Custom hwdef (aircraft-specific):{C.RESET}")
    for i, f in enumerate(custom_files, 1):
        print(f"  {i}. {f}")
    print(f"  {len(custom_files)+1}. none")
    
    sel = input(f"{C.CYAN}Choice [{len(custom_files)+1}]: {C.RESET}").strip()
    if not sel:
        return None
    
    try:
        idx = int(sel) - 1
        if 0 <= idx < len(custom_files):
            return custom_files[idx]
    except ValueError:
        pass
    
    return None


# =============================================================================
# Git Management Functions
# =============================================================================

def _git_cmd(args: list[str]) -> tuple[int, str]:
    """Run git command and return (returncode, output)."""
    proc = subprocess.run(["git"] + args, capture_output=True, text=True)
    return proc.returncode, proc.stdout.strip()


def get_git_status() -> dict:
    """Get current git status information."""
    status = {}
    
    # Current branch or tag
    rc, output = _git_cmd(["symbolic-ref", "--short", "HEAD"])
    if rc == 0:
        status["ref"] = output
        status["type"] = "branch"
    else:
        rc, output = _git_cmd(["describe", "--tags", "--exact-match"])
        if rc == 0:
            status["ref"] = output
            status["type"] = "tag"
        else:
            rc, output = _git_cmd(["rev-parse", "--short", "HEAD"])
            status["ref"] = output if rc == 0 else "unknown"
            status["type"] = "detached"
    
    # Dirty status
    rc, output = _git_cmd(["status", "--porcelain"])
    status["dirty"] = len(output) > 0 if rc == 0 else False
    
    return status


def get_latest_stable_tag() -> str | None:
    """Get the latest stable release tag (e.g., Copter-4.5.6, Plane-4.5.6)."""
    rc, output = _git_cmd(["tag", "-l", "--sort=-version:refname"])
    if rc != 0:
        return None
    
    # Look for tags like Copter-X.Y.Z, Plane-X.Y.Z, ArduPilot-X.Y.Z (not beta/rc)
    stable_pattern = re.compile(r"^(Copter|Plane|Rover|Sub|ArduPilot)-\d+\.\d+\.\d+$")
    for tag in output.splitlines():
        if stable_pattern.match(tag):
            return tag
    return None


def get_latest_beta_tag() -> str | None:
    """Get the latest beta release tag."""
    rc, output = _git_cmd(["tag", "-l", "--sort=-version:refname"])
    if rc != 0:
        return None
    
    # Look for tags with -beta or -rc
    beta_pattern = re.compile(r"^(Copter|Plane|Rover|Sub|ArduPilot)-\d+\.\d+\.\d+-(beta|rc)")
    for tag in output.splitlines():
        if beta_pattern.match(tag):
            return tag
    return None


def list_tags(pattern: str = None, limit: int = 20) -> list[str]:
    """List tags matching pattern, sorted by version (newest first)."""
    args = ["tag", "-l", "--sort=-version:refname"]
    if pattern:
        args.append(pattern)
    
    rc, output = _git_cmd(args)
    if rc != 0:
        return []
    
    tags = output.splitlines()
    return tags[:limit]


def switch_to_tag(tag: str) -> bool:
    """Switch to a specific tag. Returns True on success."""
    status = get_git_status()
    if status.get("dirty"):
        print(f"\n{C.YELLOW}Warning: You have uncommitted changes!{C.RESET}")
        confirm = input(f"{C.CYAN}Discard changes and switch? [y/N]: {C.RESET}").strip().lower()
        if confirm not in ("y", "yes"):
            return False
        # Stash changes
        _git_cmd(["stash", "push", "-u", "-m", f"Auto-stash before switching to {tag}"])
    
    print(f"\n{C.DIM}Fetching latest tags...{C.RESET}")
    rc, output = _git_cmd(["fetch", "--tags"])
    if rc != 0:
        print(f"{C.YELLOW}Warning: fetch failed (continuing anyway){C.RESET}")
        if output:
            print(f"{C.DIM}{output}{C.RESET}")
    
    print(f"{C.DIM}Switching to {tag}...{C.RESET}")
    rc, output = _git_cmd(["checkout", tag])
    if rc != 0:
        print(f"{C.RED}Failed to checkout {tag}:{C.RESET}")
        print(output)
        return False
    
    print(f"{C.GREEN}Successfully switched to {tag}{C.RESET}")
    
    # Update submodules
    print(f"{C.DIM}Updating submodules...{C.RESET}")
    rc, _ = _git_cmd(["submodule", "update", "--init", "--recursive"])
    if rc != 0:
        print(f"{C.YELLOW}Warning: submodule update had issues{C.RESET}")
    
    return True


def git_management_menu():
    """Interactive git management submenu."""
    while True:
        status = get_git_status()
        ref_type = status.get("type", "unknown")
        ref_name = status.get("ref", "unknown")
        is_dirty = status.get("dirty", False)
        
        latest_stable = get_latest_stable_tag()
        latest_beta = get_latest_beta_tag()
        
        print(f"\n{C.BOLD}=== Git Management ==={C.RESET}")
        print(f"\n{C.BOLD}Current:{C.RESET}")
        if ref_type == "tag":
            print(f"  Tag:    {C.GREEN}{ref_name}{C.RESET}")
        elif ref_type == "branch":
            print(f"  Branch: {C.CYAN}{ref_name}{C.RESET}")
        else:
            print(f"  Detached HEAD: {C.YELLOW}{ref_name}{C.RESET}")
        
        if is_dirty:
            print(f"  Status: {C.RED}DIRTY (uncommitted changes){C.RESET}")
        else:
            print(f"  Status: {C.GREEN}clean{C.RESET}")
        
        print(f"\n{C.BOLD}Latest Available:{C.RESET}")
        if latest_stable:
            print(f"  Stable: {C.GREEN}{latest_stable}{C.RESET}")
        if latest_beta:
            print(f"  Beta:   {C.YELLOW}{latest_beta}{C.RESET}")
        
        print(f"\n{C.BOLD}Options:{C.RESET}")
        print(f"  {C.CYAN}1{C.RESET}  Switch to latest stable")
        print(f"  {C.CYAN}2{C.RESET}  Switch to latest beta")
        print(f"  {C.CYAN}3{C.RESET}  List stable releases")
        print(f"  {C.CYAN}4{C.RESET}  List beta releases")
        print(f"  {C.CYAN}5{C.RESET}  Search & switch to specific tag")
        print(f"  {C.CYAN}6{C.RESET}  Fetch latest tags (no switch)")
        if is_dirty:
            print(f"  {C.CYAN}7{C.RESET}  Show uncommitted changes")
        print(f"  {C.CYAN}b{C.RESET}  Back to main menu")
        
        choice = input(f"\n{C.CYAN}Choice: {C.RESET}").strip().lower()
        
        if choice == "b":
            break
        elif choice == "1":
            if not latest_stable:
                print(f"{C.YELLOW}No stable release found{C.RESET}")
                continue
            if switch_to_tag(latest_stable):
                input(f"\n{C.DIM}Press Enter to continue...{C.RESET}")
        elif choice == "2":
            if not latest_beta:
                print(f"{C.YELLOW}No beta release found{C.RESET}")
                continue
            if switch_to_tag(latest_beta):
                input(f"\n{C.DIM}Press Enter to continue...{C.RESET}")
        elif choice == "3":
            print(f"\n{C.BOLD}Stable Releases (latest 20):{C.RESET}")
            # Get all tags, then filter for stable ones
            rc, output = _git_cmd(["tag", "-l", "--sort=-version:refname"])
            if rc == 0:
                stable_pattern = re.compile(r"^(Copter|Plane|Rover|Sub|ArduPilot)-\d+\.\d+\.\d+$")
                stable_tags = [t for t in output.splitlines() if stable_pattern.match(t)][:20]
                if stable_tags:
                    for i, tag in enumerate(stable_tags, 1):
                        current = f" {C.GREEN}(current){C.RESET}" if tag == ref_name else ""
                        print(f"  {i:2d}. {tag}{current}")
                else:
                    print(f"  {C.DIM}(none found){C.RESET}")
            else:
                print(f"  {C.RED}Failed to list tags{C.RESET}")
            input(f"\n{C.DIM}Press Enter to continue...{C.RESET}")
        elif choice == "4":
            print(f"\n{C.BOLD}Beta/RC Releases (latest 20):{C.RESET}")
            # Get all tags, then filter for beta/rc
            rc, output = _git_cmd(["tag", "-l", "--sort=-version:refname"])
            if rc == 0:
                beta_pattern = re.compile(r"^(Copter|Plane|Rover|Sub|ArduPilot)-\d+\.\d+\.\d+-(beta|rc)")
                beta_tags = [t for t in output.splitlines() if beta_pattern.match(t)][:20]
                if beta_tags:
                    for i, tag in enumerate(beta_tags, 1):
                        current = f" {C.GREEN}(current){C.RESET}" if tag == ref_name else ""
                        print(f"  {i:2d}. {tag}{current}")
                else:
                    print(f"  {C.DIM}(none found){C.RESET}")
            else:
                print(f"  {C.RED}Failed to list tags{C.RESET}")
            input(f"\n{C.DIM}Press Enter to continue...{C.RESET}")
        elif choice == "5":
            search = input(f"\n{C.CYAN}Tag pattern (e.g., 'Copter-4.5', '*4.5*'): {C.RESET}").strip()
            if not search:
                continue
            if not search.startswith("*") and not search.endswith("*"):
                search = f"*{search}*"
            tags = list_tags(search, limit=30)
            if not tags:
                print(f"{C.YELLOW}No matching tags found{C.RESET}")
                continue
            print(f"\n{C.BOLD}Matching tags:{C.RESET}")
            for i, tag in enumerate(tags, 1):
                print(f"  {i:2d}. {tag}")
            sel = input(f"\n{C.CYAN}Pick number (or Enter to cancel): {C.RESET}").strip()
            if sel:
                try:
                    idx = int(sel) - 1
                    if 0 <= idx < len(tags):
                        if switch_to_tag(tags[idx]):
                            input(f"\n{C.DIM}Press Enter to continue...{C.RESET}")
                except ValueError:
                    print(f"{C.YELLOW}Invalid number{C.RESET}")
        elif choice == "6":
            print(f"\n{C.DIM}Fetching latest tags...{C.RESET}")
            rc, output = _git_cmd(["fetch", "--tags"])
            if rc == 0:
                print(f"{C.GREEN}Tags updated{C.RESET}")
            else:
                print(f"{C.YELLOW}Fetch had issues (check network/permissions){C.RESET}")
                if output:
                    print(f"{C.DIM}{output[:200]}{C.RESET}")
            input(f"\n{C.DIM}Press Enter to continue...{C.RESET}")
        elif choice == "7" and is_dirty:
            print(f"\n{C.BOLD}Uncommitted changes:{C.RESET}")
            rc, output = _git_cmd(["status", "--short"])
            if rc == 0:
                print(output)
            input(f"\n{C.DIM}Press Enter to continue...{C.RESET}")
        else:
            print(f"{C.YELLOW}Invalid choice{C.RESET}")


def main():
    os.chdir(Path(__file__).parent)

    if not Path("waf").exists():
        print(f"{C.RED}Error: must be in ArduPilot root directory{C.RESET}")
        sys.exit(1)

    all_boards = None

    while True:
        f4_count = len(_read_board_file(BOARDS_F4_FILE))
        h7_count = len(_read_board_file(BOARDS_H7_FILE))
        total_fav = f4_count + h7_count
        
        # Get git status for header
        git_status = get_git_status()
        git_ref = git_status.get("ref", "unknown")
        git_dirty = git_status.get("dirty", False)
        git_type = git_status.get("type", "unknown")
        
        # Format git info
        if git_type == "tag":
            git_info = f"{C.GREEN}{git_ref}{C.RESET}"
        elif git_type == "branch":
            git_info = f"{C.CYAN}{git_ref}{C.RESET}"
        else:
            git_info = f"{C.YELLOW}{git_ref}{C.RESET}"
        
        if git_dirty:
            git_info += f" {C.RED}(dirty){C.RESET}"

        print(f"\n{C.BOLD}=== ArduPilot Build Tool ==={C.RESET}")
        print(f"{C.DIM}Git: {git_info}{C.RESET}")
        print(f"\n  {C.CYAN}1{C.RESET}  Build favourite boards       {C.DIM}({f4_count} F4 + {h7_count} H7 = {total_fav}){C.RESET}")
        print(f"  {C.CYAN}2{C.RESET}  Build one board")
        print(f"  {C.CYAN}3{C.RESET}  Search boards & build")
        print(f"  {C.CYAN}4{C.RESET}  Build ALL boards             {C.DIM}(takes very long){C.RESET}")
        print(f"  {C.CYAN}5{C.RESET}  Build bootloaders (favourites)")
        print(f"  {C.CYAN}6{C.RESET}  List favourite boards")
        print(f"  {C.CYAN}7{C.RESET}  Git management               {C.DIM}(switch versions){C.RESET}")
        print(f"  {C.CYAN}q{C.RESET}  Quit")

        choice = input(f"\n{C.CYAN}Choice: {C.RESET}").strip().lower()

        if choice == "q":
            break

        if choice == "6":
            favs = get_favourite_boards()
            if favs:
                for i, (b, hwdef, custom) in enumerate(favs, 1):
                    tag = Path(hwdef).stem if hwdef else "none"
                    if custom:
                        tag += f"+{Path(custom).stem}"
                    print(f"  {i:3d}. {b:<40} {C.DIM}({tag}){C.RESET}")
                print(f"\n  {C.DIM}{len(favs)} boards ({f4_count} F4 + {h7_count} H7){C.RESET}")
            else:
                print(f"  {C.YELLOW}(empty){C.RESET}")
            continue
        
        if choice == "7":
            git_management_menu()
            continue

        # boards is list of (name, hwdef, custom_hwdef) tuples
        boards: list[tuple[str, str | None, str | None]] = []
        bootloader_only = False

        if choice == "5":
            boards = get_favourite_boards()
            bootloader_only = True
            if not boards:
                continue
        elif choice == "1":
            boards = get_favourite_boards()
            if not boards:
                continue
        elif choice in ("2", "3", "4"):
            if all_boards is None:
                print(f"{C.DIM}Loading board list...{C.RESET}")
                all_boards = get_all_boards()
            if choice == "2":
                names = pick_one(all_boards)
            elif choice == "3":
                names = search_boards(all_boards)
            else:
                names = all_boards
            if not names:
                continue
            # For ad-hoc selections, ask which hwdef to use
            hwdef = _pick_hwdef()
            custom_hwdef = _pick_custom_hwdef()
            boards = [(b, hwdef, custom_hwdef) for b in names]
        else:
            print(f"{C.YELLOW}Invalid choice{C.RESET}")
            continue

        print(f"\n{C.BOLD}Selected boards ({len(boards)}):{C.RESET}")
        for b, hwdef, custom in boards:
            tag = Path(hwdef).stem if hwdef else "none"
            if custom:
                tag += f"+{Path(custom).stem}"
            print(f"  - {b:<40} {C.DIM}({tag}){C.RESET}")

        vehicles, build_bl = pick_vehicles(bootloader_mode=bootloader_only)

        confirm = input(f"\n{C.CYAN}Start build? [Y/n]: {C.RESET}").strip().lower()
        if confirm in ("", "y", "yes"):
            build_boards(boards, vehicles, build_bootloader=build_bl)
        else:
            print(f"{C.YELLOW}Cancelled{C.RESET}")


if __name__ == "__main__":
    main()
