OpenIris-ESPIDF/tools/switchBoardType.py

import os
import difflib
import shutil
import argparse
from typing import Dict, Optional, List

HEADER_COLOR = "\033[95m"
OKGREEN = "\033[92m"
WARNING = "\033[93m"
OKBLUE = "\033[94m"
ENDC = "\033[0m"

BOARDS_DIR_NAME = "boards"
SDKCONFIG_DEFAULTS_FILENAME = "sdkconfig.base_defaults"


# some components are super platform specific.
# to a point where building them with esp32 will fail every single time due to depndencies not supporting it
# with our own components, we can ship some shims to keep things clean
# but with those, unless we roll our own somehow, we're out of luck.
# So, to make things simpler, when selecting for which board to build, we're gonna reconfigure the components
# on the fly.
PLATFORM_SPECIFIC_COMPONENTS = {"esp32s3": ["usb_device_uvc"]}
PLATFORM_SPECIFIC_COMPONENTS_DIRS = {"esp32s3": "esp32s3"}


def get_root_path() -> str:
    return os.path.split(os.path.dirname(os.path.realpath(__file__)))[0]


def get_boards_root() -> str:
    return os.path.join(get_root_path(), BOARDS_DIR_NAME)


def get_config_platform(_parsed_config: dict) -> str:
    # 1:-1 to strip quotes
    return _parsed_config["CONFIG_IDF_TARGET"][1:-1]


def enumerate_board_configs() -> Dict[str, str]:
    """Walk the boards directory and build a mapping of board names to absolute file paths.

    Naming strategy:
      - Relative path from boards/ to file with path separators replaced by '_'.
      - If the last two path segments are identical (e.g. project_babble/project_babble) collapse to a single segment.
      - For facefocusvr eye boards we keep eye_L / eye_R suffix to distinguish configs even though WHO_AM_I is same.
    """
    boards_dir = get_boards_root()
    mapping: Dict[str, str] = {}
    if not os.path.isdir(boards_dir):
        return mapping
    for root, _dirs, files in os.walk(boards_dir):
        for f in files:
            if f == SDKCONFIG_DEFAULTS_FILENAME:
                continue
            rel_path = os.path.relpath(os.path.join(root, f), boards_dir)
            parts = rel_path.split(os.sep)
            if len(parts) >= 2 and parts[-1] == parts[-2]:  # collapse duplicate tail
                parts = parts[:-1]
            board_key = "_".join(parts)
            mapping[board_key] = os.path.join(root, f)
    return mapping


BOARD_CONFIGS = enumerate_board_configs()


def build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser()
    p.add_argument(
        "-b",
        "--board",
        help="Board name (run with --list to see options). Flexible: accepts path-like or partial if unique.",
    )
    p.add_argument(
        "--list", help="List discovered boards and exit", action="store_true"
    )
    p.add_argument("--dry-run", help="Dry run, won't modify files", action="store_true")
    p.add_argument(
        "--diff",
        help="Show the difference between base config and selected board",
        action="store_true",
    )
    p.add_argument("--ssid", help="Set the WiFi SSID", type=str, default="")
    p.add_argument("--password", help="Set the WiFi password", type=str, default="")
    p.add_argument("--clear-wifi", help="Clear WiFi credentials", action="store_true")
    return p


def list_boards():
    print("Discovered boards:")
    width = max((len(k) for k in BOARD_CONFIGS), default=0)
    for name, path in sorted(BOARD_CONFIGS.items()):
        print(f"  {name.ljust(width)}  ->  {os.path.relpath(path, get_root_path())}")


def _suggest_boards(partial: str) -> List[str]:
    if not partial:
        return []
    partial_low = partial.lower()
    contains = [b for b in BOARD_CONFIGS if partial_low in b.lower()]
    if contains:
        return contains[:10]
    # fallback to fuzzy matching using difflib
    choices = list(BOARD_CONFIGS.keys())
    return difflib.get_close_matches(partial, choices, n=5, cutoff=0.4)


def normalize_board_name(raw: Optional[str]) -> Optional[str]:
    if raw is None:
        return None
    candidate = raw.strip()
    if not candidate:
        return None
    candidate = candidate.replace("\\", "/").rstrip("/")
    # strip leading folders like tools/, boards/
    parts = [
        p
        for p in candidate.split("/")
        if p not in (".", "") and p not in ("tools", "boards")
    ]
    if parts:
        candidate = parts[-1] if len(parts) == 1 else "_".join(parts)
    candidate = candidate.replace("-", "_")
    # exact match
    if candidate in BOARD_CONFIGS:
        return candidate
    # try ending match
    endings = [b for b in BOARD_CONFIGS if b.endswith(candidate)]
    if len(endings) == 1:
        return endings[0]
    if len(endings) > 1:
        print(f"Ambiguous board '{raw}'. Could be: {', '.join(endings)}")
        return None
    # attempt case-insensitive
    lower_map = {b.lower(): b for b in BOARD_CONFIGS}
    if candidate.lower() in lower_map:
        return lower_map[candidate.lower()]
    return None


def get_main_config_path() -> str:
    return os.path.join(get_root_path(), "sdkconfig")


def get_board_config_path(board_key: str) -> str:
    return BOARD_CONFIGS[board_key]


def get_base_config_path() -> str:
    # base defaults moved under boards directory
    return os.path.join(get_boards_root(), SDKCONFIG_DEFAULTS_FILENAME)


def parse_config(config_file) -> dict:
    config = {}
    for line in config_file:
        line = line.strip().split("=")
        if len(line) == 2:
            config[line[0]] = line[1]
        else:
            # lines without value are usually comments, we're safe to store empty string there
            config[line[0]] = ""
    return config


def handle_wifi_config(_new_config: dict, _main_config: dict, _args) -> dict:
    if _args.ssid:
        _new_config["CONFIG_WIFI_SSID"] = f'"{_args.ssid}"'
        _new_config["CONFIG_WIFI_PASSWORD"] = f'"{_args.password}"'
    else:
        if "CONFIG_WIFI_SSID" in _main_config:
            _new_config["CONFIG_WIFI_SSID"] = _main_config["CONFIG_WIFI_SSID"]
        if "CONFIG_WIFI_PASSWORD" in _main_config:
            _new_config["CONFIG_WIFI_PASSWORD"] = _main_config["CONFIG_WIFI_PASSWORD"]

    if _args.clear_wifi:
        _new_config["CONFIG_WIFI_SSID"] = '""'
        _new_config["CONFIG_WIFI_PASSWORD"] = '""'
    return _new_config


def compute_diff(_parsed_base_config: dict, _parsed_board_config: dict) -> dict:
    _diff = {}
    for _key in _parsed_board_config:
        if _key not in _parsed_base_config:
            if _parsed_board_config[_key] != "":
                _diff[_key] = f"{OKGREEN}+{ENDC} {_parsed_board_config[_key]}"
        else:
            if _parsed_board_config[_key] != _parsed_base_config[_key]:
                _diff[_key] = (
                    f"{OKGREEN}{_parsed_base_config[_key]}{ENDC} -> {OKBLUE}{_parsed_board_config[_key]}{ENDC}"
                )
    return _diff


def _move_directories(component: str, destination_path: str):
    if os.path.exists(component):
        shutil.move(component, destination_path)


def handle_extra_components(old_platform: str, new_platform: str, dry_run: bool):
    print(
        f"{OKGREEN}Switching components configuration from platform:{ENDC} {OKBLUE}{old_platform}{ENDC} {OKGREEN}to platform:{ENDC} {OKBLUE}{new_platform}{ENDC}"
    )

    if old_platform == new_platform:
        print(f"{OKGREEN}The platform is the same. Nothing to do here.{ENDC}")
        return

    old_platform_components = PLATFORM_SPECIFIC_COMPONENTS.get(old_platform, [])
    new_platform_components = PLATFORM_SPECIFIC_COMPONENTS.get(new_platform, [])
    if dry_run:
        print(f"{OKGREEN}Would remove: {ENDC}")
        for component in old_platform_components:
            print(f"{OKBLUE}- {component} {ENDC}")

        print(f"{OKGREEN}Would add: {ENDC}")
        for component in new_platform_components:
            print(f"{OKBLUE}- {component} {ENDC}")

        return

    components_path = os.path.join(get_root_path(), "components")

    if old_base_dir := PLATFORM_SPECIFIC_COMPONENTS_DIRS.get(old_platform):
        old_extra_components_path = os.path.join(
            os.path.join(get_root_path(), "extra_components"), old_base_dir
        )
        for component in old_platform_components:
            component_path = os.path.join(components_path, component)
            print(
                f"{OKGREEN}Moving:{ENDC}{OKBLUE} {component}{ENDC} to {OKBLUE}{old_extra_components_path}{ENDC}"
            )
            _move_directories(component_path, old_extra_components_path)

    if new_base_dir := PLATFORM_SPECIFIC_COMPONENTS_DIRS.get(new_platform):
        new_extra_components_path = os.path.join(
            os.path.join(get_root_path(), "extra_components"), new_base_dir
        )
        for component in new_platform_components:
            component_path = os.path.join(new_extra_components_path, component)
            print(
                f"{OKGREEN}Moving:{ENDC}{OKBLUE} {component}{ENDC} to {OKBLUE}{components_path}{ENDC}"
            )
            _move_directories(component_path, components_path)


def main():
    parser = build_arg_parser()
    args = parser.parse_args()

    if args.list:
        list_boards()
        return

    board_input = args.board
    if not board_input:
        parser.error("--board is required (or use --list)")
    normalized = normalize_board_name(board_input)
    if not normalized:
        print(f"{WARNING}Unknown board '{board_input}'.")
        suggestions = _suggest_boards(board_input)
        if suggestions:
            print("Did you mean: " + ", ".join(suggestions))
        print("Use --list to see all boards.")
        raise SystemExit(2)

    if not os.path.isfile(get_base_config_path()):
        raise SystemExit(f"Base defaults file not found: {get_base_config_path()}")

    print(
        f"{OKGREEN}Switching configuration to board:{ENDC} {OKBLUE}{normalized}{ENDC}"
    )
    print(f"{OKGREEN}Using defaults from :{ENDC} {get_base_config_path()}")
    print(
        f"{OKGREEN}Using board config from :{ENDC} {get_board_config_path(normalized)}"
    )

    with open(get_main_config_path(), "r+") as main_config:
        parsed_main_config = parse_config(main_config)

    with (
        open(get_base_config_path(), "r") as base_config,
        open(get_board_config_path(normalized), "r") as board_config,
    ):
        parsed_base_config = parse_config(base_config)
        parsed_board_config = parse_config(board_config)

    new_board_config = {**parsed_base_config, **parsed_board_config}
    new_board_config = handle_wifi_config(new_board_config, parsed_main_config, args)

    if args.diff:
        print("---" * 5, f"{WARNING}DIFF{ENDC}", "---" * 5)
        diff = compute_diff(parsed_main_config, new_board_config)
        if not diff:
            print(
                f"{HEADER_COLOR}[DIFF]{ENDC} No changes between existing main config and {OKBLUE}{normalized}{ENDC}"
            )
        else:
            print(
                f"{HEADER_COLOR}[DIFF]{ENDC} Keys differing (main -> new {OKBLUE}{normalized}{ENDC}):"
            )
            for key in sorted(diff):
                print(f"{HEADER_COLOR}[DIFF]{ENDC} {key} : {diff[key]}")
        print("---" * 14)

    if not args.dry_run:
        print(f"{WARNING}Writing changes to main config file{ENDC}")
        with open(get_main_config_path(), "w") as main_config:
            for key, value in new_board_config.items():
                if value:
                    main_config.write(f"{key}={value}\n")
                else:
                    main_config.write(f"{key}\n")
    else:
        print(f"{WARNING}[DRY-RUN]{ENDC} Skipping writing to files")

    handle_extra_components(
        get_config_platform(parsed_main_config),
        get_config_platform(new_board_config),
        args.dry_run,
    )

    print(
        f"{OKGREEN}Done. ESP-IDF is setup to build for:{ENDC} {OKBLUE}{normalized}{ENDC}"
    )


if __name__ == "__main__":
    main()