# Copyright 2026 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

import os
import sys
import argparse
import subprocess
import json
import re
from pathlib import Path

import kconfig_utils

SUPPORTED_ARCHITECTURES = ["arc", "arm", "arm64", "mips", "posix", "riscv", "sparc", "xtensa", "x86", "rx"]

# Add Zephyr's kconfiglib to sys.path
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--zephyr-base", required=True)
    parser.add_argument("--app-dir", required=True)
    parser.add_argument("--board", required=True)
    parser.add_argument("--board-dir", required=True)
    parser.add_argument("--parent-platform", required=True)
    parser.add_argument("--modules", nargs="+", default=[])
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--oot-dts-roots", nargs="+", default=[])
    parser.add_argument("--app-name", default="zephyr_app")
    parser.add_argument("--toolchain-variant", default=None)
    parser.add_argument("--conf-fragments", nargs="+", default=[])
    return parser.parse_args()

def setup_environment(args, zephyr_base, output_dir, board_dir):
    os.environ["srctree"] = zephyr_base
    os.environ["ZEPHYR_BASE"] = zephyr_base
    os.environ["KCONFIG_BINARY_DIR"] = output_dir

    # Replicate Zephyr SDK compiler capabilities for Kconfig
    os.environ["TOOLCHAIN_HAS_PICOLIBC"] = "y"
    os.environ["TOOLCHAIN_HAS_NEWLIB"] = "y"

    if "/" in args.board:
        board_name, board_qualifiers = args.board.split("/", 1)
        os.environ["BOARD"] = board_name
        os.environ["BOARD_QUALIFIERS"] = board_qualifiers
    else:
        os.environ["BOARD"] = args.board
        os.environ["BOARD_QUALIFIERS"] = ""

    os.environ["ARCH"] = "*"
    os.environ["HWM_SCHEME"] = "v2"
    os.environ["KCONFIG_FUNCTIONS"] = "kconfigfunctions"
    os.environ["TOOLCHAIN_HAS_PICOLIBC"] = "y"
    os.environ["TOOLCHAIN_HAS_NEWLIB"] = "y"
    os.environ["ZEPHYR_TOOLCHAIN_VARIANT"] = "zephyr"

    # Determine toolchain variant similar to Zephyr's FindHostTools.cmake
    toolchain_variant = args.toolchain_variant
    if not toolchain_variant:
        board_dir_lower = board_dir.lower()
        if "native" in args.board or "boards/native" in board_dir_lower or "posix" in board_dir_lower:
            toolchain_variant = "host"
        else:
            toolchain_variant = "zephyr"

    os.environ["TOOLCHAIN_KCONFIG_DIR"] = os.path.join(zephyr_base, "cmake", "toolchain", toolchain_variant)
    os.environ["KCONFIG_BOARD_DIR"] = board_dir
    os.environ["SOC_NAME"] = "none" # Will be overridden by defconfig

    kconfig_env_file = os.path.join(output_dir, "kconfig_module_dirs.env")
    os.makedirs(output_dir, exist_ok=True)
    with open(kconfig_env_file, "w") as f:
        if args.modules:
            for m in args.modules:
                if os.path.isdir(m):
                    module_name = os.path.basename(m.rstrip('/')).split('+')[-1].replace('-', '_').upper()
                    m_abs = os.path.abspath(m)
                    f.write(f"ZEPHYR_{module_name}_MODULE_DIR={m_abs}\n")
                    os.environ[f"ZEPHYR_{module_name}_MODULE_DIR"] = m_abs
    os.environ["KCONFIG_ENV_FILE"] = kconfig_env_file

def generate_kconfig_aggregations(zephyr_base, args, output_dir, board_dir):
    # Generate aggregation Kconfig files
    soc_dirs = []

    # Try to find SoC name from board.yml
    soc_names = []
    for board_file in ["board.yml", "board.yaml"]:
        board_yml_path = os.path.join(board_dir, board_file)
        if os.path.exists(board_yml_path):
            with open(board_yml_path, 'r') as f:
                content = f.read()
                # Simple heuristic to find soc name in board.yml
                # It looks like:
                # socs:
                #   - name: same70q21b
                matches = re.findall(r"-\s*name:\s*(\S+)", content)
                if matches:
                    soc_names.extend(matches)
            break

    soc_roots = [Path(zephyr_base)] + ([Path(m) for m in args.modules] if args.modules else [])
    for sroot in soc_roots:
        if not sroot.exists():
            continue
        for root, _, files in os.walk(sroot, followlinks=True):
            if "soc.yml" in files:
                soc_dirs.append(root)

    kconfig_utils.generate_aggregation_file(os.path.join(output_dir, "soc", "Kconfig.defconfig"), soc_dirs, "SoC defconfigs")
    kconfig_utils.generate_aggregation_file(os.path.join(output_dir, "soc", "Kconfig"), soc_dirs, "SoC Kconfigs")
    kconfig_utils.generate_aggregation_file(os.path.join(output_dir, "soc", "Kconfig.soc"), soc_dirs, "SoC Kconfigs")
    kconfig_utils.generate_aggregation_file(os.path.join(output_dir, "soc", "Kconfig.sysbuild"), soc_dirs, "SoC sysbuild Kconfigs")

    shield_defconfig = os.path.join(output_dir, "Kconfig.shield.defconfig")
    os.makedirs(os.path.dirname(shield_defconfig), exist_ok=True)
    with open(shield_defconfig, "w") as f:
        f.write("# Empty shield defconfig\n")

    arch_dirs = []
    arch_root = os.path.join(zephyr_base, "arch")
    for d in os.listdir(arch_root):
        if os.path.isdir(os.path.join(arch_root, d)):
            arch_dirs.append(os.path.join(arch_root, d))
    kconfig_utils.generate_aggregation_file(os.path.join(output_dir, "arch", "Kconfig"), arch_dirs, "Arch Kconfigs")

    board_dirs = [board_dir]
    kconfig_utils.generate_aggregation_file(os.path.join(output_dir, "boards", "Kconfig.defconfig"), board_dirs, "Board defconfigs")
    kconfig_utils.generate_aggregation_file(os.path.join(output_dir, "boards", "Kconfig"), board_dirs, "Board Kconfigs")
    kconfig_utils.generate_aggregation_file(os.path.join(output_dir, "boards", "Kconfig.sysbuild"), board_dirs, "Board sysbuild Kconfigs")

    board_name = os.environ["BOARD"]
    board_kconfig_agg = os.path.join(output_dir, "boards", f"Kconfig.{board_name}")
    os.makedirs(os.path.dirname(board_kconfig_agg), exist_ok=True)
    with open(board_kconfig_agg, "w") as f:
        board_kconfig = os.path.join(board_dir, "Kconfig")
        if not os.path.exists(board_kconfig):
             board_kconfig = os.path.join(board_dir, f"Kconfig.{board_name}")
        if os.path.exists(board_kconfig):
            f.write(f'osource "{os.path.abspath(board_kconfig)}"\n')

    kconfig_modules = os.path.join(output_dir, "Kconfig.modules")
    kconfig_utils.generate_kconfig_modules(zephyr_base, args.modules, kconfig_modules)

def discover_overlays(app_dir, board_name, board_qualifiers):
    board_qualifiers = board_qualifiers.lstrip("/")
    overlays = []

    # 1. Board-specific overlays in app
    if board_qualifiers:
        safe_q = board_qualifiers.replace("/", "_")
        soc_board_overlay = os.path.join(
            app_dir, "boards", f"{board_name}_{safe_q}.overlay"
        )
        if os.path.exists(soc_board_overlay):
            overlays.append(soc_board_overlay)

    board_overlay = os.path.join(app_dir, "boards", f"{board_name}.overlay")
    if os.path.exists(board_overlay):
        overlays.append(board_overlay)

    # 2. Default app overlay
    app_overlay = os.path.join(app_dir, "app.overlay")
    if os.path.exists(app_overlay):
        overlays.append(app_overlay)

    return overlays

def preprocess_dts(args, zephyr_base, output_dir, board_dir):
    # Determine input DTS file
    board_name = os.environ["BOARD"]
    board_qualifiers = os.environ.get("BOARD_QUALIFIERS", "")

    candidates = [
        os.path.join(board_dir, f"{args.board.replace('/', '_')}.dts"),
        os.path.join(board_dir, f"{board_name}.dts"),
    ]

    # In HWM v2, try to find a DTS file matching any SoC in board.yml
    for board_file in ["board.yml", "board.yaml"]:
        board_yml_path = os.path.join(board_dir, board_file)
        if os.path.exists(board_yml_path):
            with open(board_yml_path, 'r') as f:
                content = f.read()
                matches = re.findall(r"-\s*name:\s*(\S+)", content)
                for soc in matches:
                    candidates.append(os.path.join(board_dir, f"{board_name}_{soc}.dts"))
            break

    dts_file = None
    for candidate in candidates:
        if os.path.exists(candidate):
            dts_file = candidate
            break

    if not dts_file:
        raise RuntimeError(f"Could not find DTS file for board {args.board} in {board_dir}. Tried: {candidates}")

    overlays = discover_overlays(args.app_dir, board_name, board_qualifiers)

    dts_input_path = os.path.join(output_dir, "dts_input.c")
    os.makedirs(output_dir, exist_ok=True)
    with open(dts_input_path, "w") as f:
        f.write("/* Generated by zephyr-bazel - do not edit! */\n")
        f.write(f'#include "{os.path.abspath(dts_file)}"\n')
        for overlay in overlays:
            f.write(f'#include "{os.path.abspath(overlay)}"\n')

    merged_dts_path = os.path.join(output_dir, "zephyr.dts")

    # Preprocess DTS using pcpp
    import pcpp
    pcmd_path = os.path.join(os.path.dirname(pcpp.__file__), "pcmd.py")
    pcpp_args = [
        sys.executable,
        pcmd_path,
        dts_input_path,
    ]

    raw_roots = [Path(r) for r in args.oot_dts_roots] + [Path(m) for m in args.modules] + [Path(zephyr_base)]
    dts_roots = []
    for r in raw_roots:
        if r not in dts_roots:
            dts_roots.append(r)
    for dts_root in dts_roots:
        directories = [
            "include",
            "include/zephyr",
            "dts",
            "dts/common",
        ] + ["dts/" + a for a in SUPPORTED_ARCHITECTURES] + ["dts/vendor"]
        for directory in directories:
            path = dts_root / directory
            if path.exists():
                pcpp_args.append("-I" + str(path.resolve()))
    pcpp_args.extend(["-o", merged_dts_path])

    subprocess.run(pcpp_args, check=True)
    return merged_dts_path, dts_roots

def generate_edt_pickle(zephyr_base, output_dir, dts_roots, merged_dts_path):
    output_edt_pickle_path = os.path.join(output_dir, "edt.pickle")

    bindings_dirs = []
    for dts_root in dts_roots:
        bdir = dts_root / "zephyr" / "dts" / "bindings"
        if bdir.exists():
            if str(bdir.resolve()) not in bindings_dirs:
                bindings_dirs.append(str(bdir.resolve()))
        else:
            bdir = dts_root / "dts" / "bindings"
            if bdir.exists():
                if str(bdir.resolve()) not in bindings_dirs:
                    bindings_dirs.append(str(bdir.resolve()))

    gen_edt_args = [
        sys.executable,
        os.path.join(zephyr_base, "scripts", "dts", "gen_edt.py"),
        "--edt-pickle-out",
        output_edt_pickle_path,
        "--dts",
        merged_dts_path,
        "--bindings-dirs",
    ]
    gen_edt_args.extend(bindings_dirs)
    gen_edt_args.extend([
        "--dts-out",
        os.path.join(output_dir, "zephyr.dts.debug"),
        "--dtc-flags",
        "Wno-simple_bus_reg",
    ])

    # Add PYTHONPATH for edtlib
    env = os.environ.copy()
    env["PYTHONPATH"] = os.path.join(zephyr_base, "scripts", "dts") + ":" + os.path.join(zephyr_base, "scripts", "dts", "python-devicetree", "src") + ":" + env.get("PYTHONPATH", "")
    print(f"DEBUG PYTHONPATH: {env['PYTHONPATH']}")
 
    subprocess.run(gen_edt_args, env=env, check=True)
    os.environ["EDT_PICKLE"] = output_edt_pickle_path

def generate_dts_headers(zephyr_base, output_dir, edt_pickle_path):
    gen_defines_py = os.path.join(zephyr_base, "scripts", "dts", "gen_defines.py")
    output_header = os.path.join(output_dir, "zephyr", "devicetree_generated.h")
    os.makedirs(os.path.dirname(output_header), exist_ok=True)

    cmd = [
        sys.executable,
        gen_defines_py,
        "--edt-pickle", edt_pickle_path,
        "--header-out", output_header,
    ]

    # Ensure edtlib is in PYTHONPATH (portable using os.pathsep)
    env = os.environ.copy()
    env["PYTHONPATH"] = os.pathsep.join([
        os.path.join(zephyr_base, "scripts", "dts"),
        os.path.join(zephyr_base, "scripts", "dts", "python-devicetree", "src"),
        env.get("PYTHONPATH", ""),
    ])

    subprocess.run(cmd, env=env, check=True)

def get_conf_files(args, board_dir):
    board_name = os.environ["BOARD"]
    board_qualifiers = os.environ.get("BOARD_QUALIFIERS", "").lstrip("/")

    conf_files = []

    # Try qualified defconfig (HWM v2)
    if board_qualifiers:
        safe_qualifiers = board_qualifiers.replace("/", "_")
        qualified_defconfig = os.path.join(board_dir, f"{board_name}_{safe_qualifiers}_defconfig")
        if os.path.exists(qualified_defconfig):
            conf_files.append(qualified_defconfig)

    # Fallback to base defconfig
    if not conf_files:
        board_defconfig = os.path.join(board_dir, f"{board_name}_defconfig")
        if os.path.exists(board_defconfig):
            conf_files.append(board_defconfig)

    prj_conf = os.path.join(args.app_dir, "prj.conf")
    if os.path.exists(prj_conf):
        conf_files.append(prj_conf)

    # Also try board-specific conf in app dir
    if board_qualifiers:
        soc_board_conf = os.path.join(args.app_dir, "boards", f"{board_name}_{board_qualifiers}.conf")
        if os.path.exists(soc_board_conf):
            conf_files.append(soc_board_conf)

    board_conf = os.path.join(args.app_dir, "boards", f"{board_name}.conf")
    if os.path.exists(board_conf):
        conf_files.append(board_conf)

    return conf_files

def generate_config_c(kconf, output_dir, args):
    os.makedirs(os.path.join(output_dir, "zephyr"), exist_ok=True)
    with open(os.path.join(output_dir, "zephyr", "config.c"), "w") as f:
        f.write("/* AUTO-GENERATED by kconfig_gen_values.py, do not edit! */\n\n")
        f.write("#include <zephyr/toolchain.h>\n\n")
        f.write("GEN_ABS_SYM_BEGIN (_ConfigAbsSyms)\n")

        # Call scripts/build/generate_configs_dot_c.py methodology
        for sym in kconf.unique_defined_syms:
            if not sym.nodes:
                continue

            val = None
            # Import here because it's not available at the top level due to path setup.
            import kconfiglib
            if sym.type in [kconfiglib.BOOL, kconfiglib.TRISTATE]:
                if sym.str_value == "y":
                    val = "1"
            elif sym.type == kconfiglib.INT:
                if sym.str_value:
                    val = sym.str_value
            elif sym.type == kconfiglib.HEX:
                if sym.str_value:
                    val = "0x" + sym.str_value.replace("0x", "")
            elif sym.type == kconfiglib.STRING:
                if sym.str_value:
                    val = "1"

            if val is not None:
                f.write(f"GEN_ABSOLUTE_SYM_KCONFIG(CONFIG_{sym.name}, {val});\n")

        f.write("GEN_ABS_SYM_END\n")

def generate_bazel_build(args, kconf, output_dir, kconfiglib):
    resolved_arch = "arm"
    for a in SUPPORTED_ARCHITECTURES:
        sym_name = "ARCH_POSIX" if a == "posix" else a.upper()
        sym = kconf.syms.get(sym_name)
        if sym and sym.str_value == "y":
            resolved_arch = a
            break

    if resolved_arch == "arm":
        # Fallback check for ARC if standard symbol is missing
        sym = kconf.syms.get("DT_HAS_SNPS_ARCEM_ENABLED")
        if sym and sym.str_value == "y":
            resolved_arch = "arc"

    arch_macro = "__arm__"
    if resolved_arch == "arm64":
        arch_macro = "__aarch64__"
    elif resolved_arch == "x86":
        arch_macro = "__i386__"
    elif resolved_arch == "riscv":
        arch_macro = "__riscv"

    soc_name = kconf.syms.get("SOC").str_value
    if not soc_name:
        # Try SOC_SERIES or other fallbacks if SOC is empty
        soc_name = kconf.syms.get("SOC_SERIES").str_value
    if not soc_name:
        soc_name = os.environ.get("BOARD", "unknown")

    target_name = kconfig_utils.sanitize_name_for_target(soc_name)
    print(f"DEBUG: soc_name={soc_name}, target_name={target_name}")

    with open(os.path.join(output_dir, "BUILD.bazel"), "w") as f:
        f.write('load("@rules_cc//cc:defs.bzl", "cc_library")\n')
        f.write('load("@zephyr//:cc.bzl", "zephyr_cc_library")\n')
        f.write('package(default_visibility = ["//visibility:public"])\n\n')
        f.write('exports_files(["zephyr/autoconf.h", "zephyr/devicetree_generated.h", ".config"])\n\n')

        f.write('cc_library(\n')
        f.write('    name = "devicetree_generated",\n')
        f.write('    hdrs = ["zephyr/devicetree_generated.h"],\n')
        f.write('    includes = ["."],\n')
        f.write(')\n\n')

        f.write('cc_library(\n')
        f.write('    name = "autoconf_library",\n')
        f.write('    hdrs = ["zephyr/autoconf.h"],\n')
        f.write('    includes = ["."],\n')
        f.write('    deps = [":devicetree_generated"],\n')
        f.write(')\n\n')

        f.write('zephyr_cc_library(\n')
        f.write('    name = "autoconf_symbols",\n')
        f.write('    srcs = ["zephyr/config.c"],\n')
        f.write('    copts = ["-D%s"],\n' % arch_macro)
        f.write('    deps = [\n')
        f.write('        ":autoconf_library",\n')
        f.write('        "@zephyr//include:zephyr",\n')
        f.write('    ],\n')
        f.write(')\n\n')

        f.write('platform(\n')
        f.write('    name = "%s",\n' % target_name)
        f.write('    parents = ["%s"],\n' % args.parent_platform)
        f.write('    flags = [\n')
        for sym in kconf.unique_defined_syms:
             if not sym.nodes:
                 continue

             sname = kconfig_utils.sanitize_name_for_target(sym.name)
             full_name = "CONFIG_" + sname

             if sym.type in [kconfiglib.BOOL, kconfiglib.TRISTATE]:
                 val = "true" if sym.str_value == "y" else "false"
                 f.write('        "--@zephyr_kconfig//:%s=%s",\n' % (full_name, val))
             elif sym.type == kconfiglib.INT:
                 if sym.str_value:
                     f.write('        "--@zephyr_kconfig//:%s=%s",\n' % (full_name, sym.str_value))
             elif sym.type == kconfiglib.HEX:
                 if sym.str_value:
                     try:
                         val = int(sym.str_value, 16)
                         f.write('        "--@zephyr_kconfig//:%s=%d",\n' % (full_name, val))
                     except:
                         pass
             elif sym.type == kconfiglib.STRING:
                 if sym.str_value:
                     f.write('        "--@zephyr_kconfig//:%s=%s",\n' % (full_name, sym.str_value))

        # Resolve board DTS target
        board_label = args.parent_platform
        board_pkg = board_label.split(":")[0]
        f.write('        "--@zephyr//:autoconf_file=@zc_target//:zephyr/autoconf.h",\n')
        f.write('        "--@zephyr//:autoconf_library=@zc_target//:autoconf_library",\n')
        f.write('        "--@zephyr//:autoconf_symbols_to_link=@zc_target//:autoconf_symbols",\n')
        f.write('        "--@zephyr//:dts_cc_library=@zc_target//:devicetree_generated",\n')
        f.write('    ],\n')
        f.write(')\n\n')
        f.write('alias(name = "platform", actual = ":%s")\n' % target_name)

def main():
    args = parse_args()
    kconfig_utils.setup_kconfiglib(args.zephyr_base)
    # kconfiglib is in Zephyr's scripts/kconfig directory, which is added to path by setup_kconfiglib above.
    import kconfiglib

    zephyr_base = os.path.abspath(args.zephyr_base)
    output_dir = os.path.abspath(args.output_dir)
    board_dir = os.path.abspath(args.board_dir)

    board_name = args.board.split("/")[0]
    if not os.path.exists(os.path.join(board_dir, "board.yml")) and not os.path.exists(os.path.join(board_dir, f"{board_name}.dts")):
        found = False
        for root in [board_dir] + args.oot_dts_roots:
            for dirpath, _, filenames in os.walk(root, followlinks=True):
                if f"{board_name}.dts" in filenames:
                    board_dir = os.path.abspath(dirpath)
                    found = True
                    break
            if found: break

    setup_environment(args, zephyr_base, output_dir, board_dir)

    generate_kconfig_aggregations(zephyr_base, args, output_dir, board_dir)

    merged_dts_path, dts_roots = preprocess_dts(args, zephyr_base, output_dir, board_dir)
    generate_edt_pickle(zephyr_base, output_dir, dts_roots, merged_dts_path)
    generate_dts_headers(zephyr_base, output_dir, os.path.join(output_dir, "edt.pickle"))
    kconfig_utils.generate_kconfig_dts(zephyr_base, os.path.join(output_dir, "Kconfig.dts"), dts_roots)

    # Re-verify environment for kconfiglib
    os.environ["ARCH"] = "*"
    os.environ["ZEPHYR_BASE"] = zephyr_base

    conf_files = get_conf_files(args, board_dir)
    conf_files.extend(args.conf_fragments)

    os.makedirs(os.path.join(output_dir, "zephyr"), exist_ok=True)

    # Call Zephyr's kconfig.py
    kconfig_py = os.path.join(zephyr_base, "scripts", "kconfig", "kconfig.py")
    cmd = [
        sys.executable,
        kconfig_py,
        "--handwritten-input-configs",
        os.path.join(zephyr_base, "Kconfig"),
        os.path.join(output_dir, ".config"),
        os.path.join(output_dir, "zephyr", "autoconf.h"),
        os.path.join(output_dir, "kconfig_list.txt"),
    ] + conf_files

    env = os.environ.copy()
    
    res = subprocess.run(cmd, env=env, capture_output=True, text=True)
    
    if res.stderr:
        print(res.stderr, file=sys.stderr)

    if res.returncode != 0:
        sys.exit(f"Kconfig generation failed with exit code {res.returncode}")

    # Add the app name to the autoconf.h since bazel does not propagate build flags to dependencies.
    with open(os.path.join(output_dir, "zephyr", "autoconf.h"), "a") as f:
        f.write(f'#define APP_NAME "{args.app_name}"\n')

    # Load generated .config with kconfiglib
    kconf = kconfiglib.Kconfig(os.path.join(zephyr_base, "Kconfig"))
    kconf.load_config(os.path.join(output_dir, ".config"))

    generate_config_c(kconf, output_dir, args)
    generate_bazel_build(args, kconf, output_dir, kconfiglib)

if __name__ == "__main__":
    main()