# pylint: disable=E0602
Import("env")  # noqa

import os
import re
import subprocess

# ESPHome marks ISR code IRAM_ATTR, which on LibreTiny maps to a per-family
# section routed into RAM-executable memory (see esphome/core/hal.h). The
# input section name is always .sram.text; only the output section it lands
# in differs per family.
#
# This script is NOT loaded on BK72xx (IRAM_ATTR is a no-op there; the SDK
# masks FIQ+IRQ around flash writes). On the remaining families:
# - RTL8720C: stock linker consumes *(.sram.text*) into .ram.code_text.
# - RTL8710B: stock linker has KEEP(*(.image2.ram.text*)) in .ram_image2.text,
#   but ltchiptool's AmebaZ elf2bin (soc/ambz/binary.py) does NOT list
#   .ram_image2.text in sections_ram, so code there is silently dropped from
#   the flashed image. Inject KEEP(*(.sram.text*)) at the top of
#   .ram_image2.data (which IS extracted) instead.
# - LN882H: stock linker has no glob for ".sram.text", so we inject
#   KEEP(*(.sram.text*)) into ".flash_copysection" (> RAM0 AT> FLASH)
#   immediately after KEEP(*(.vectors)), so the vector table stays at
#   __copysection_ram0_start (0x20000000) for correct Cortex-M4 VTOR alignment.
#
# All families also get a post-link summary showing where IRAM_ATTR landed.


_MARKER = "/* esphome .sram.text */"
# Strong assignments (not PROVIDE) so the symbols are always emitted in the
# ELF; PROVIDE symbols with no references can be garbage-collected.
_KEEP_LINE = (
    "    __esphome_sram_text_start = .; "
    "KEEP(*(.sram.text*)) "
    "__esphome_sram_text_end = .; "
    + _MARKER + "\n"
)
# Inject after KEEP(*(.vectors)) so the vector table stays at
# __copysection_ram0_start (0x20000000). Cortex-M4 VTOR requires a 512-byte-
# aligned address; injecting before the vectors would push them to an
# unaligned offset and mis-route every IRQ handler.
_LN_COPY = re.compile(r"(KEEP\(\*\(\.vectors\)\)[^\n]*\n)")
# Inject at the top of .ram_image2.data, before __data_start__ so our code
# does not fall inside the data range markers. .ram_image2.data is one of the
# sections ltchiptool's AmebaZ elf2bin extracts; BD_RAM is rwx so the code is
# executable. AmbZ has no C runtime .data copy loop (the bootloader loads
# image2 into BD_RAM whole) so the inline code is not clobbered after boot.
#
# The regex is intentionally strict (no attribute / ALIGN between the section
# name and the opening brace, brace on its own line). If a future AmbZ SDK
# linker template changes this format, _pre_link raises RuntimeError on the
# unpatched .ld file(s), and the RTL8710B CI compile job in
# tests/test_build_components fails on the PR, surfacing the mismatch loudly
# rather than silently shipping a binary with IRAM_ATTR code dropped from
# one or both OTA slots.
_AMBZ_DATA = re.compile(r"(\.ram_image2\.data\s*:\s*\n?\s*\{\s*\n)")


def _detect(env):
    prefix = "USE_LIBRETINY_VARIANT_"
    # CPPDEFINES may hold strings or (name, value) tuples; BUILD_FLAGS holds
    # the raw "-DNAME" strings. PlatformIO populates both, but the exact order
    # vs. extra_scripts varies, so check both to be robust.
    for token in env.get("CPPDEFINES", []):
        if isinstance(token, (list, tuple)):
            token = token[0]
        if isinstance(token, str) and token.startswith(prefix):
            return token[len(prefix):]
    for flag in env.get("BUILD_FLAGS", []):
        if isinstance(flag, str) and "-D" + prefix in flag:
            name = flag.split("-D", 1)[1].split("=", 1)[0].strip()
            if name.startswith(prefix):
                return name[len(prefix):]
    return None


KNOWN_VARIANTS = frozenset({
    "LN882H",
    "RTL8710B",
    "RTL8720C",
})


def _inject_keep(host_section):
    """Return a patcher that injects _KEEP_LINE after `host_section` match."""
    def patch(content):
        if _MARKER in content:
            return content
        return host_section.sub(r"\1" + _KEEP_LINE, content, count=1)
    return patch


# Variants not listed here intentionally have no .ld patcher:
# - RTL8720C: stock linker already consumes *(.sram.text*) into .ram.code_text.
# - BK72xx (all): SDK masks FIQ+IRQ around flash writes, IRAM_ATTR is no-op.
_PATCHERS_BY_VARIANT = {
    "LN882H": (_inject_keep(_LN_COPY),),
    "RTL8710B": (_inject_keep(_AMBZ_DATA),),
}


def _patchers_for(variant):
    return _PATCHERS_BY_VARIANT.get(variant, ())


def _pre_link(target, source, env):
    build_dir = env.subst("$BUILD_DIR")
    ld_files = [f for f in os.listdir(build_dir) if f.endswith(".ld")]
    patched = []
    unpatched = []
    for name in ld_files:
        path = os.path.join(build_dir, name)
        with open(path, "r", encoding="utf-8") as fh:
            original = fh.read()
        if _MARKER in original:
            patched.append(name)
            continue
        content = original
        for fn in _patchers:
            content = fn(content)
        if content != original:
            with open(path, "w", encoding="utf-8") as fh:
                fh.write(content)
            print("ESPHome: patched {} for IRAM_ATTR placement".format(name))
            patched.append(name)
        else:
            unpatched.append(name)
    if not patched:
        raise RuntimeError(
            "ESPHome: no .ld in {} was patched for IRAM_ATTR. Update the "
            "regex in patch_linker.py.script (_PATCHERS_BY_VARIANT).".format(
                build_dir
            )
        )
    # Every .ld in the build must be patched. RTL8710B generates one .ld per
    # OTA slot (xip1, xip2); if only one matches, the unpatched slot would
    # ship with IRAM_ATTR code dropped to zeros and brick the device on the
    # boot after an OTA into that slot.
    if unpatched:
        raise RuntimeError(
            "ESPHome: {} of {} .ld file(s) in {} were not patched for "
            "IRAM_ATTR: {}. The regex in patch_linker.py.script "
            "(_PATCHERS_BY_VARIANT[{!r}]) matched the others but not "
            "these. Update the regex to cover all linker scripts.".format(
                len(unpatched), len(ld_files), build_dir,
                ", ".join(unpatched), _variant,
            )
        )


# Substrings matched against demangled names as a fallback on RTL8720C,
# where we cannot inject __esphome_sram_text_start/end markers.
_FALLBACK_SUBSTRINGS = ("wake_loop_any_context", "wake_loop_isrsafe",
                        "enable_loop_soon_any_context")


def _post_link(target, source, env):
    """Print where IRAM_ATTR ended up so users can confirm at a glance."""
    elf = env.subst("$BUILD_DIR/${PROGNAME}.elf")
    if not os.path.isfile(elf):
        return
    nm = env.subst("$NM")
    try:
        out = subprocess.check_output(
            [nm, "--defined-only", "--demangle", elf], text=True
        )
    except (OSError, subprocess.CalledProcessError) as exc:
        print("ESPHome: IRAM_ATTR summary unavailable (nm failed: {})".format(exc))
        return
    start = end = None
    fallback = []
    for line in out.splitlines():
        parts = line.split(maxsplit=2)
        if len(parts) != 3:
            continue
        addr_str, _kind, name = parts
        if name == "__esphome_sram_text_start":
            start = int(addr_str, 16)
        elif name == "__esphome_sram_text_end":
            end = int(addr_str, 16)
        elif "veneer" not in name and any(s in name for s in _FALLBACK_SUBSTRINGS):
            fallback.append(int(addr_str, 16))
    print("ESPHome: IRAM_ATTR placement summary ({}):".format(_variant))
    if start is not None and end is not None:
        print("  .sram.text:  {} bytes at 0x{:08x} - 0x{:08x}".format(end - start, start, end))
    elif fallback:
        lo, hi = min(fallback), max(fallback)
        print("  IRAM symbols at 0x{:08x} - 0x{:08x} (approx {} bytes)".format(lo, hi, hi - lo))
    else:
        print("  no IRAM_ATTR symbols found")


if (_variant := _detect(env)) is None:
    raise RuntimeError(
        "ESPHome: could not determine LibreTiny variant from build flags. "
        "patch_linker.py needs USE_LIBRETINY_VARIANT_* to route IRAM_ATTR "
        "into SRAM; without it, ISR handlers would silently end up in flash."
    )
if _variant not in KNOWN_VARIANTS:
    raise RuntimeError(
        "ESPHome: unknown LibreTiny variant {!r}; patch_linker.py does not "
        "know how to route IRAM_ATTR into SRAM for this family. Update "
        "patch_linker.py.script before shipping firmware.".format(_variant)
    )

if _patchers := _patchers_for(_variant):
    # LibreTiny writes the processed .ld templates into $BUILD_DIR during its
    # own builder setup, which may run after this script. Register the patch
    # as a pre-link action so it executes once the linker scripts exist.
    env.AddPreAction("$BUILD_DIR/${PROGNAME}.elf", _pre_link)

# Post-link summary for every family that reaches this script.
env.AddPostAction("$BUILD_DIR/${PROGNAME}.elf", _post_link)