#!/usr/bin/env python3
# trunk-ignore-all(ruff/F821)
# trunk-ignore-all(flake8/F821): For SConstruct imports
import sys
from os.path import basename, isabs, isfile, join, realpath
import shutil
import subprocess
import json
import re
from datetime import datetime
from typing import Dict

from readprops import readProps

Import("env")
platform = env.PioPlatform()
progname = env.get("PROGNAME")
lfsbin = f"{progname.replace('firmware-', 'littlefs-')}.bin"
manifest_ran = False

def infer_architecture(board_cfg):
    try:
        mcu = board_cfg.get("build.mcu") if board_cfg else None
    except KeyError:
        mcu = None
    except Exception:
        mcu = None
    if not mcu:
        return None
    mcu_l = str(mcu).lower()
    if "esp32s3" in mcu_l:
        return "esp32-s3"
    if "esp32c6" in mcu_l:
        return "esp32-c6"
    if "esp32c3" in mcu_l:
        return "esp32-c3"
    if "esp32" in mcu_l:
        return "esp32"
    if "rp2040" in mcu_l:
        return "rp2040"
    if "rp2350" in mcu_l:
        return "rp2350"
    if "nrf52" in mcu_l or "nrf52840" in mcu_l:
        return "nrf52840"
    if "nrf54l15" in mcu_l:
        return "nrf54l15"
    if "stm32" in mcu_l:
        return "stm32"
    print(f"mtjson: could not infer architecture from MCU '{mcu_l}'")
    return None

def run_size_tool(env, flag, purpose):
    """Run the toolchain size tool against the built ELF and return its output.

    Shared plumbing for compute_ram_bytes / compute_flash_bytes. Returns None
    when the ELF is missing or the tool fails; the manifest then simply omits
    the metric and downstream size reports show "n/a".
    """
    elf = env.File(env.subst("$BUILD_DIR/${PROGNAME}.elf"))
    if not elf.exists():
        return None
    size_tool = env.subst("$SIZETOOL") or "size"
    try:
        return subprocess.check_output(
            [size_tool, flag, elf.get_abspath()],
            env=env["ENV"],
            universal_newlines=True,
        )
    except Exception as exc:
        print(f"mtjson: skipping {purpose} ({size_tool} failed: {exc})")
        return None

def compute_ram_bytes(env):
    """Static RAM usage (.data + .bss) of the ELF, via the toolchain size tool.

    Deliberately excludes heap/stack placeholder sections (e.g. the nRF52 .heap
    section): on nRF52840 the heap arena is the linker gap after .bss, so static
    RAM growth shrinks the usable heap 1:1 - which is exactly why we track it.
    Returns None when the value cannot be determined; the manifest then simply
    omits ram_bytes and downstream size reports show "n/a".
    """
    output = run_size_tool(env, "-A", "ram_bytes")
    if output is None:
        return None
    ram = 0
    found = False
    for line in output.splitlines():
        parts = line.split()
        if len(parts) < 2:
            continue
        name = parts[0]
        # Main-SRAM static sections: .data/.bss, platform-prefixed variants (e.g.
        # ESP32 .dram0.data/.dram0.bss), and RISC-V small-data .sdata/.sbss.
        # ESP-IDF .rtc.* sections live outside the heap-competing SRAM; .heap and
        # .tdata never match.
        if name.startswith(".rtc"):
            continue
        if name in (".data", ".bss", ".sdata", ".sbss") or name.endswith(".data") or name.endswith(".bss"):
            try:
                ram += int(parts[1])
                found = True
            except ValueError:
                continue
    return ram if found else None

def compute_flash_bytes(env):
    """Flash footprint (text + data) of the ELF via the toolchain size tool.

    Approximates the flashed app image for targets whose packaged artifacts do
    not include a raw .bin (e.g. nRF52: hex/uf2/DFU zip only); consumed by
    bin/collect_sizes.py as a fallback flash measurement for the size report
    and the bin/ram_budgets.json budget gate. Returns None when the value
    cannot be determined; the manifest then simply omits flash_bytes.
    """
    output = run_size_tool(env, "-B", "flash_bytes")
    if output is None:
        return None
    for line in output.splitlines():
        parts = line.split()
        if len(parts) >= 3 and parts[0].isdigit() and parts[1].isdigit():
            return int(parts[0]) + int(parts[1])
    return None

def export_link_files(source, target, env):
    """Copy the linker scripts that define MEMORY regions, and the linker map, next to the ELF.

    CI uploads them with the ELF so MemBrowse can report usage per memory region and
    attribute symbols to archives. Scripts come from LDSCRIPT_PATH and -T link flags,
    resolved like the linker does: relative to the project dir, then each LIBPATH entry.
    --defsym values a script uses (e.g. STM32's LD_MAX_SIZE) are prepended to its copy so
    its MEMORY block resolves standalone. The map is the last -Map in LINKFLAGS, the one
    the linker writes: STM32's framework adds its own after the global output.map, and ESP-IDF
    builds use --Map instead.
    """
    names = []
    if env.get("LDSCRIPT_PATH"):
        names.append(env.subst("$LDSCRIPT_PATH"))
    flags = [env.subst(str(f)) for f in env.Flatten(env.get("LINKFLAGS", []))]
    for i, flag in enumerate(flags):
        if flag == "-T" and i + 1 < len(flags):
            names.append(flags[i + 1])
        elif flag.startswith("-Wl,-T"):
            names.append(flag[len("-Wl,-T"):].lstrip(","))
        elif flag.startswith("-T"):
            names.append(flag[2:])
    defsyms = [m.groups() for m in (re.match(r"-Wl,--defsym[,=]([A-Za-z_]\w*)=(\S+)", f) for f in flags) if m]
    search_dirs = [env.subst("$PROJECT_DIR")] + [
        env.Dir(env.subst(str(p))).get_abspath() for p in env.Flatten(env.get("LIBPATH", []))
    ]
    exported = set()
    for name in names:
        name = name.strip('"')
        candidates = [name] if isabs(name) else [join(d, name) for d in search_dirs]
        path = next((c for c in candidates if isfile(c)), None)
        if path is None or realpath(path) in exported:
            continue
        with open(path, errors="replace") as f:
            script = f.read()
        if not re.search(r"\bMEMORY\s*\{", script):
            continue
        exported.add(realpath(path))
        dest = env.subst(f"$BUILD_DIR/${{PROGNAME}}.{basename(path)}")
        with open(dest, "w") as f:
            used = [(n, v) for n, v in defsyms if re.search(rf"\b{n}\b", script)]
            if used:
                f.write("/* --defsym values from the link command */\n")
                f.writelines(f"{n} = {v};\n" for n, v in used)
            f.write(script)
        print(f"Exported linker script {path} -> {dest}")

    map_path = None
    for flag in flags:
        m = re.match(r"-Wl,--?Map[,=](.+)", flag)
        if m:
            map_path = m.group(1).replace('"', "")
    if map_path:
        map_path = join(env.subst("$PROJECT_DIR"), map_path)
        dest = env.subst("$BUILD_DIR/${PROGNAME}.map")
        if isfile(map_path) and realpath(map_path) != realpath(dest):
            shutil.copyfile(map_path, dest)
            print(f"Exported linker map {map_path} -> {dest}")

def manifest_gather(source, target, env):
    global manifest_ran
    if manifest_ran:
        return
    # Skip manifest generation if we cannot determine architecture (host/native builds)
    board_arch = infer_architecture(env.BoardConfig())
    if not board_arch:
        print(f"Skipping mtjson generation for unknown architecture (env={env.get('PIOENV')})")
        manifest_ran = True
        return
    manifest_ran = True
    out = []
    board_platform = env.BoardConfig().get("platform")
    board_mcu = env.BoardConfig().get("build.mcu").lower()
    needs_ota_suffix = board_platform == "nordicnrf52"
    
    # Mapping of bin files to their target partition names
    # Maps the filename pattern to the partition name where it should be flashed
    partition_map = {
        f"{progname}.bin": "app0",              # primary application slot (app0 / OTA_0)
        lfsbin: "spiffs",                        # filesystem image flashed to spiffs
    }
    
    check_paths = [
        progname,
        f"{progname}.elf",
        f"{progname}.bin",
        f"{progname}.factory.bin",
        f"{progname}.hex",
        f"{progname}.merged.hex",
        f"{progname}.uf2",
        f"{progname}.factory.uf2",
        f"{progname}.zip",
        lfsbin,
        f"mt-{board_mcu}-ota.bin",
        "bleota-c3.bin"
    ]
    for p in check_paths:
        f = env.File(env.subst(f"$BUILD_DIR/{p}"))
        if f.exists():
            manifest_name = p
            if needs_ota_suffix and p == f"{progname}.zip":
                manifest_name = f"{progname}-ota.zip"
            d = {
                "name": manifest_name,
                "md5": f.get_content_hash(), # Returns MD5 hash
                "bytes": f.get_size() # Returns file size in bytes
            }
            # Add part_name if this file represents a partition that should be flashed
            if p in partition_map:
                d["part_name"] = partition_map[p]
            out.append(d)
            print(d)
    manifest_write(out, env, compute_ram_bytes(env), compute_flash_bytes(env))

def manifest_write(files, env, ram_bytes=None, flash_bytes=None):
    # Defensive: also skip manifest writing if we cannot determine architecture
    def get_project_option(name):
        try:
            return env.GetProjectOption(name)
        except Exception:
            return None

    def get_project_option_any(names):
        for name in names:
            val = get_project_option(name)
            if val is not None:
                return val
        return None

    def as_bool(val):
        return str(val).strip().lower() in ("1", "true", "yes", "on")

    def as_int(val):
        try:
            return int(str(val), 10)
        except (TypeError, ValueError):
            return None

    def as_list(val):
        return [item.strip() for item in str(val).split(",") if item.strip()]

    manifest = {
        "version": verObj["long"],
        "build_epoch": build_epoch,
        "platformioTarget": env.get("PIOENV"),
        "mcu": env.get("BOARD_MCU"),
        "repo": repo_owner,
        "files": files,
        "has_mui": False,
        "has_inkhud": False,
    }
    # Static RAM footprint (.data + .bss); consumed by bin/collect_sizes.py for
    # the CI size report and the bin/ram_budgets.json budget gate.
    if ram_bytes is not None:
        manifest["ram_bytes"] = ram_bytes
    # Flash footprint (text + data); fallback flash measurement for targets
    # whose packaged artifacts include no raw .bin (e.g. nRF52).
    if flash_bytes is not None:
        manifest["flash_bytes"] = flash_bytes
    # Get partition table (generated in esp32_pre.py) if it exists
    if env.get("custom_mtjson_part"):
        # custom_mtjson_part is a JSON string, convert it back to a dict
        pj = json.loads(env.get("custom_mtjson_part"))
        manifest["part"] = pj
    # Enable has_mui for TFT builds
    if ("HAS_TFT", 1) in env.get("CPPDEFINES", []):
        manifest["has_mui"] = True
    if "MESHTASTIC_INCLUDE_INKHUD" in env.get("CPPDEFINES", []):
        manifest["has_inkhud"] = True

    pioenv = env.get("PIOENV")
    device_meta = {}
    device_meta_fields = [
        ("hwModel", ["custom_meshtastic_hw_model"], as_int),
        ("hwModelSlug", ["custom_meshtastic_hw_model_slug"], str),
        ("architecture", ["custom_meshtastic_architecture"], str),
        ("activelySupported", ["custom_meshtastic_actively_supported"], as_bool),
        ("displayName", ["custom_meshtastic_display_name"], str),
        ("supportLevel", ["custom_meshtastic_support_level"], as_int),
        ("images", ["custom_meshtastic_images"], as_list),
        ("tags", ["custom_meshtastic_tags"], as_list),
        ("requiresDfu", ["custom_meshtastic_requires_dfu"], as_bool),
        ("partitionScheme", ["custom_meshtastic_partition_scheme"], str),
        ("url", ["custom_meshtastic_url"], str),
        ("key", ["custom_meshtastic_key"], str),
        ("variant", ["custom_meshtastic_variant"], str),
    ]


    for manifest_key, option_keys, caster in device_meta_fields:
        raw_val = get_project_option_any(option_keys)
        if raw_val is None:
            continue
        parsed = caster(raw_val) if callable(caster) else raw_val
        if parsed is not None and parsed != "":
            device_meta[manifest_key] = parsed

    # Board MCU wins over a hand-typed custom_meshtastic_architecture: only the
    # spellings infer_architecture emits are recognized downstream.
    declared = device_meta.get("architecture")
    board_arch = infer_architecture(env.BoardConfig()) or declared
    if declared and declared != board_arch:
        print(f"{pioenv}: architecture '{declared}' overridden with '{board_arch}'")
    if not board_arch:
        print(f"Skipping mtjson write for unknown architecture (env={env.get('PIOENV')})")
        return

    device_meta["architecture"] = board_arch

    # Always set requiresDfu: true for nrf52840 targets
    if board_arch == "nrf52840":
        device_meta["requiresDfu"] = True

    device_meta.setdefault("displayName", pioenv)
    device_meta.setdefault("activelySupported", False)

    if device_meta:
        manifest.update(device_meta)

    # Write the manifest to the build directory
    with open(env.subst("$BUILD_DIR/${PROGNAME}.mt.json"), "w") as f:
        json.dump(manifest, f, indent=2)

Import("projenv")

prefsLoc = projenv["PROJECT_DIR"] + "/version.properties"
verObj = readProps(prefsLoc)
print(f"Using meshtastic platformio-custom.py, firmware version {verObj['long']} on {env.get('PIOENV')}")

# get repository owner if git is installed
try:
    r_owner = (
        subprocess.check_output(["git", "config", "--get", "remote.origin.url"])
        .decode("utf-8")
        .strip().split("/")
    )
    repo_owner = r_owner[-2] + "/" + r_owner[-1].replace(".git", "")
except subprocess.CalledProcessError:
    repo_owner = "unknown"

jsonLoc = env["PROJECT_DIR"] + "/userPrefs.jsonc"
with open(jsonLoc) as f:
    jsonStr = re.sub("//.*","", f.read(), flags=re.MULTILINE)
    userPrefs = json.loads(jsonStr)

# Channels::initDefaultChannel() applies a configured index as a whole, so resolve per-field
# optionality here: any field the vendor left out gets the value that function would have kept.
MAX_NUM_CHANNELS = 8
CHANNEL_FIELD_DEFAULTS = {
    "PSK": "{ 0x01 }",  # short-form index into the well-known default PSK
    "NAME": "",
    "PRECISION": "0",
    "IS_MUTED": "false",
    "UPLINK_ENABLED": "false",
    "DOWNLINK_ENABLED": "false",
}
channelsToWriteRaw = userPrefs.get("USERPREFS_CHANNELS_TO_WRITE", "1")
channelsToWrite = int(channelsToWriteRaw, 16 if channelsToWriteRaw.lower().startswith("0x") else 10)
if channelsToWrite > MAX_NUM_CHANNELS:
    sys.exit(
        f"userPrefs.jsonc: USERPREFS_CHANNELS_TO_WRITE is {channelsToWrite}, "
        f"the channel table holds {MAX_NUM_CHANNELS}"
    )
for i in range(MAX_NUM_CHANNELS):
    prefix = f"USERPREFS_CHANNEL_{i}_"
    if any(k.startswith(prefix) for k in list(userPrefs)):
        for field, default in CHANNEL_FIELD_DEFAULTS.items():
            userPrefs.setdefault(prefix + field, default)

pref_flags = []
# Pre-process the userPrefs
for pref in userPrefs:
    if userPrefs[pref].startswith("{"):
        pref_flags.append("-D" + pref + "=" + userPrefs[pref])
    elif userPrefs[pref].lstrip("-").replace(".", "").isdigit():
        pref_flags.append("-D" + pref + "=" + userPrefs[pref])
    elif re.fullmatch(r"0[xX][0-9a-fA-F]+", userPrefs[pref]):
        pref_flags.append("-D" + pref + "=" + userPrefs[pref])
    elif userPrefs[pref] == "true" or userPrefs[pref] == "false":
        pref_flags.append("-D" + pref + "=" + userPrefs[pref])
    elif userPrefs[pref].startswith("meshtastic_"):
        pref_flags.append("-D" + pref + "=" + userPrefs[pref])
    # If the value is a string, we need to wrap it in quotes
    else:
        pref_flags.append("-D" + pref + "=" + env.StringifyMacro(userPrefs[pref]) + "")

# General options that are passed to the C and C++ compilers
# Calculate unix epoch for current day (midnight)
current_date = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
build_epoch = int(current_date.timestamp())

flags = [
        "-DAPP_VERSION=" + verObj["long"],
        "-DAPP_VERSION_SHORT=" + verObj["short"],
        "-DAPP_ENV=" + env.get("PIOENV"),
        "-DAPP_REPO=" + repo_owner,
        "-DBUILD_EPOCH=" + str(build_epoch),
    ] + pref_flags

print("Using flags:")
for flag in flags:
    print(flag)

projenv.Append(
    CCFLAGS=flags,
)

for lb in env.GetLibBuilders():
    if lb.name == "meshtastic-device-ui":
        lb.env.Append(CPPDEFINES=[("APP_VERSION", verObj["long"])])
        break

# DEBUG_MUTE already compiles every LOG_* call out entirely; nanopb's own error-message
# strings are a separate mechanism it doesn't touch, so mirror the same intent into nanopb.
debug_mute_defined = any(
    d == "DEBUG_MUTE" or (isinstance(d, tuple) and d[0] == "DEBUG_MUTE") for d in env.get("CPPDEFINES", [])
)
if debug_mute_defined:
    projenv.Append(CPPDEFINES=[("PB_NO_ERRMSG", 1)])
    for lb in env.GetLibBuilders():
        if lb.name == "Nanopb":
            lb.env.Append(CPPDEFINES=[("PB_NO_ERRMSG", 1)])
            break

# Get the display resolution from macros
def get_display_resolution(build_flags):
    # Check "DISPLAY_SIZE" to determine the screen resolution
    for flag in build_flags:
        if isinstance(flag, tuple) and flag[0] == "DISPLAY_SIZE":
            screen_width, screen_height = map(int, flag[1].split("x"))
            return screen_width, screen_height
    print("No screen resolution defined in build_flags. Please define DISPLAY_SIZE.")
    exit(1)

def load_boot_logo(source, target, env):
    build_flags = env.get("CPPDEFINES", [])
    logo_w, logo_h = get_display_resolution(build_flags)
    print(f"TFT build with {logo_w}x{logo_h} resolution detected")

    # Load the boot logo from `branding/logo_<width>x<height>.png` if it exists
    source_path = join(env["PROJECT_DIR"], "branding", f"logo_{logo_w}x{logo_h}.png")
    dest_dir = join(env["PROJECT_DIR"], "data", "boot")
    dest_path = join(dest_dir, "logo.png")
    if env.File(source_path).exists():
        print(f"Loading boot logo from {source_path}")
        # Prepare the destination
        env.Execute(f"mkdir -p {dest_dir} && rm -f {dest_path}")
        # Copy the logo to the `data/boot` directory
        env.Execute(f"cp {source_path} {dest_path}")

# Load the boot logo on TFT builds
if ("HAS_TFT", 1) in env.get("CPPDEFINES", []):
    env.AddPreAction(f"$BUILD_DIR/{lfsbin}", load_boot_logo)

board_arch = infer_architecture(env.BoardConfig())
should_skip_manifest = board_arch is None

# Most platforms can generate the manifest as part of the default 'buildprog' target.
# Typically this passes success/failure properly.
mtjson_deps = ["buildprog"]
if platform.name == "espressif32":
    # On ESP32, we need to explicitly depend upon the binary to prevent fake-success upon failure.
    mtjson_deps = ["$BUILD_DIR/${PROGNAME}.bin"]
    # Build littlefs image as part of mtjson target
    # Equivalent to `pio run -t buildfs`
    target_lfs = env.DataToBin(
        join("$BUILD_DIR", "${ESP32_FS_IMAGE_NAME}"), "$PROJECT_DATA_DIR"
    )
    mtjson_deps.append(target_lfs)

if should_skip_manifest:
    def skip_manifest(source, target, env):
        print(f"mtjson: skipped for native environment: {env.get('PIOENV')}")

    env.AddCustomTarget(
        name="mtjson",
        # For host/native envs, avoid depending on 'buildprog' (some targets don't define it)
        dependencies=[],
        actions=[skip_manifest],
        title="Meshtastic Manifest (skipped)",
        description="mtjson generation is skipped for native environments",
        always_build=True,
    )
else:
    env.AddCustomTarget(
        name="mtjson",
        dependencies=mtjson_deps,
        actions=[manifest_gather],
        title="Meshtastic Manifest",
        description="Generating Meshtastic manifest JSON + Checksums",
        always_build=True,
    )

    # Run manifest generation as part of the default build pipeline for non-native builds.
    env.Default("mtjson")

    env.AddPostAction("$BUILD_DIR/${PROGNAME}.elf", export_link_files)