#!/usr/bin/env python3
"""
Append an include-tree comment after the contents of each `memmap*.ld` under
`src/*/pico_platform/`.

The comment matches the output of `memmap_include_tree.py` (two columns, no color).
Re-running replaces a previous block marked `BEGIN_MEMMAP_INCLUDE_TREE` /
`END_MEMMAP_INCLUDE_TREE` (at the end of the file, or at the beginning for older runs).

A companion `memmap_xxx.cmake` file is written next to each `memmap_xxx.ld`,
containing `pico_add_link_depend(pico_standard_link ...)` for every transitively
included linker-script fragment of that memmap, with paths expressed as
`${PICO_SDK_PATH}/...` relative to the SDK root.

Examples:

    ./tools/memmap_annotate.py # annotates all memmap*.ld files in the SDK

    ./tools/memmap_annotate.py --ci # does not write files; exits with failure if any memmap would need updating
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

_TOOLS = Path(__file__).resolve().parent
if str(_TOOLS) not in sys.path:
    sys.path.insert(0, str(_TOOLS))

import memmap_include_tree as inc_tree

_ANNOT_BEGIN = "/* BEGIN_MEMMAP_INCLUDE_TREE\n"
_ANNOT_END = " * END_MEMMAP_INCLUDE_TREE */\n"
_ANNOT_BLOCK_CORE = (
    r"/\* BEGIN_MEMMAP_INCLUDE_TREE\n.*?\n \* END_MEMMAP_INCLUDE_TREE \*/\s*"
)


def strip_annotation_block(text: str) -> str:
    """Remove a previously inserted annotation block (trailing, or legacy leading)."""
    t = text
    # Preferred: block at end of file
    m = re.search(_ANNOT_BLOCK_CORE + r"\Z", t, re.DOTALL)
    if m:
        head = t[: m.start()].rstrip("\n")
        if head:
            head += "\n"
        return head
    # Legacy: block was prepended at the beginning
    m2 = re.match(_ANNOT_BLOCK_CORE + r"\n?", t, re.DOTALL)
    if m2:
        return t[m2.end() :].lstrip("\n")
    return t


def build_annotation_block(tree_lines: list[str]) -> str:
    parts = [
        _ANNOT_BEGIN,
        " * This is the default include tree for this file (generated by tools/memmap_annotate.py\n"
        " * and tools/memmap_include_tree.py). To override any files in this tree, see the\n"
        " * pico_add_linker_script_override_path CMake function under pico_standard_link.\n",
        " *\n",
    ]
    for ln in tree_lines:
        parts.append(" * " + ln + "\n")
    parts.append(_ANNOT_END)
    return "".join(parts)


def collect_included_files(nodes: list[inc_tree.ChildRow], acc: set[Path] | None = None) -> set[Path]:
    """Recursively collect all resolved file paths from an include tree (excludes the root)."""
    if acc is None:
        acc = set()
    for path, subs, _ in nodes:
        if path.exists():
            acc.add(path)
        collect_included_files(subs, acc)
    return acc


def build_cmake_block(included_files: set[Path], sdk_root: Path, memmap_stem: str) -> str:
    func_name = "pico_add_memmap_link_depends_internal"
    lines = [
        "# Auto-generated by tools/memmap_annotate.py — do not edit by hand\n",
        f"function({func_name} TARGET)\n",
    ]
    for p in sorted(included_files):
        try:
            rel = p.relative_to(sdk_root)
        except ValueError:
            rel = p
        lines.append(
            f"    pico_add_link_depend(${{TARGET}} ${{PICO_SDK_PATH}}/{rel.as_posix()})\n"
        )
    lines.append("endfunction()\n")
    return "".join(lines)


def annotate_one(
    path: Path,
    sdk_root: Path,
    platform: str | None,
    *,
    ci: bool,
) -> tuple[bool, str]:
    """Return (ok, status message). ok is False on unresolved includes."""
    path = path.resolve()
    search_dirs = inc_tree.default_search_dirs(sdk_root, path, platform)
    visited: set[Path] = set()
    user_l_set: frozenset[Path] = frozenset()
    user_l_labels: dict[Path, str] = {}
    root_node, root_children = inc_tree.walk_includes(
        path, search_dirs, user_l_set, user_l_labels, visited
    )
    missing = inc_tree.collect_unresolved_include_paths(root_children)
    if missing:
        names = ", ".join(m.name for m in missing)
        return False, f"unresolved INCLUDE: {names}"

    included_files = collect_included_files(root_children)

    lines = inc_tree.format_tree_lines(root_node, root_children, color=False)
    block = build_annotation_block(lines)
    old_ld = path.read_text(encoding="utf-8", errors="replace")
    body = strip_annotation_block(old_ld)
    if body and not body.endswith("\n"):
        body += "\n"
    new_ld = body + "\n" + block if body else block

    cmake_path = path.with_suffix(".cmake")
    new_cmake = build_cmake_block(included_files, sdk_root, path.stem)
    old_cmake = cmake_path.read_text(encoding="utf-8", errors="replace") if cmake_path.exists() else ""

    if new_ld == old_ld and new_cmake == old_cmake:
        return True, "unchanged"
    if ci:
        return True, "needs updating"
    if new_ld != old_ld:
        path.write_text(new_ld, encoding="utf-8")
    if new_cmake != old_cmake:
        cmake_path.write_text(new_cmake, encoding="utf-8")
    return True, "updated"


def parse_args(argv: list[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=inc_tree.fold_docstring_for_argparse(__doc__ or ""),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument(
        "--sdk-root",
        type=Path,
        default=None,
        help=inc_tree.CLI_HELP_SDK_ROOT,
    )
    p.add_argument(
        "--platform",
        default=None,
        metavar="CHIP",
        help=inc_tree.CLI_HELP_PLATFORM,
    )
    p.add_argument(
        "--ci",
        action="store_true",
        help="Do not write files; exit with failure if any memmap would need updating",
    )
    p.add_argument(
        "paths",
        nargs="*",
        type=Path,
        metavar="MEMMAP.LD",
        help="Linker scripts to annotate (default: all src/*/pico_platform/memmap*.ld)",
    )
    return p.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    sdk_root = args.sdk_root
    if args.paths:
        files = [p.resolve() for p in args.paths]
        for f in files:
            if not f.is_file():
                print(f"error: not a file: {f}", file=sys.stderr)
                return 1
        if sdk_root is None:
            sdk_root = inc_tree.infer_sdk_root(files[0])
        if sdk_root is None:
            print(inc_tree.CLI_MSG_SDK_ROOT_INFER_FAILED, file=sys.stderr)
            return 1
    else:
        if sdk_root is None:
            sdk_root = inc_tree.infer_sdk_root(Path.cwd())
        if sdk_root is None:
            print(inc_tree.CLI_MSG_SDK_ROOT_INFER_FAILED, file=sys.stderr)
            return 1
        files = sorted(sdk_root.glob("src/*/pico_platform/memmap*.ld"))

    if not files:
        print("error: no memmap*.ld files to process", file=sys.stderr)
        return 1

    sdk_root = sdk_root.resolve()
    if args.platform is not None:
        perr = inc_tree.validate_explicit_platform(sdk_root, args.platform)
        if perr:
            print(f"error: {perr}", file=sys.stderr)
            return 1

    err = 0
    for f in files:
        plat = args.platform or inc_tree.infer_platform_chip(f)
        ok, msg = annotate_one(f, sdk_root, plat, ci=args.ci)
        try:
            tag = str(f.relative_to(sdk_root))
        except ValueError:
            tag = str(f)
        if ok:
            print(inc_tree.format_memmap_annotate_status_line(tag, msg))
            if args.ci and msg == "needs updating":
                err = 1
        else:
            print(f"{tag}: {msg}", file=sys.stderr)
            err = 1
    return err


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))