# Copyright (c) 2025 Nordic Semiconductor ASA
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import importlib.resources
import pickle
from collections import defaultdict
from enum import auto, Flag
from functools import lru_cache
from itertools import chain
from pathlib import Path
from typing import Iterable, Optional, Sequence, TYPE_CHECKING

import svd
from intelhex import IntelHex

if TYPE_CHECKING:
    from devicetree import edtlib

from . import log, uicr, __version__
from .platform import (
    Address,
    AddressRegion,
    CTRLSEL_DEFAULT,
    DomainID,
    OwnerID,
    peripheral_id_get,
    ProcessorID,
)
from ._common import get_field
from ._svd_util import SvdEnum


class MigrateFlag(Flag):
    DEFAULTS = auto()

    @classmethod
    def all(cls) -> MigrateFlag:
        return ~MigrateFlag(0)

    @classmethod
    def default(cls) -> MigrateFlag:
        return cls.all()


def uicr_migrate_hex_files_to_periphconf(
    uicr_hex_files: list[Path],
    edt_pickle_file: Optional[Path] = None,
    flags: MigrateFlag = MigrateFlag.default(),
) -> str:
    if edt_pickle_file is not None:
        with edt_pickle_file.open("rb") as file:
            dt = pickle.load(file)

        periph_annotation_lookup = _build_periph_annotation_dt(dt)
    else:
        periph_annotation_lookup = {}

    to_migrate = []

    for uicr_hex_file in uicr_hex_files:
        ihex = IntelHex()
        with uicr_hex_file.open("r", encoding="utf-8") as fp:
            ihex.loadhex(fp)

        uicr_bytes = bytearray(ihex.tobinstr())
        to_migrate.append(uicr.Uicr.from_bytes(uicr_bytes))

    return uicr_migrate(
        to_migrate,
        flags=flags,
        periph_annotation_lookup=periph_annotation_lookup,
    )


def _build_periph_annotation_dt(dt: edtlib.EDT) -> dict[int, str]:
    lut = {}

    for node in dt.nodes:
        if not node.regs:
            continue

        addr = node.regs[0].addr
        if node.labels:
            lut[addr] = node.labels[0]
        else:
            lut[addr] = node.path

    return lut


def uicr_migrate(
    to_migrate: Sequence[uicr.Uicr],
    flags: MigrateFlag = MigrateFlag.default(),
    periph_annotation_lookup: dict[int, str] = {},
) -> str:
    header_lines = [
        f"/* Generated by nrf-regtool version {__version__} */",
        "#include <uicr/uicr.h>",
        "",
    ]

    ipcmap_idx = 0

    # For grouping lines by type so that e.g. GPIO settings from two different UICRs are grouped
    # together (the group name is arbitrary).
    pconf_lines_by_type = defaultdict(list)

    for builder in to_migrate:
        uicr_svd = builder._uicr

        default_domain = DomainID(uicr_svd["HEADER"]["DOMAIN"]["DOMAIN"].content)
        default_processor = ProcessorID.from_domain(default_domain)
        default_owner = OwnerID.from_processor(default_processor)

        for reg in uicr_svd["IPCMAP"]:
            if not reg.modified:
                break

            source_domain = DomainID(reg["DOMAINIDSOURCE"].content)
            source_ch = reg["IPCTCHSOURCE"].content
            sink_domain = DomainID(reg["DOMAINIDSINK"].content)
            sink_ch = reg["IPCTCHSINK"].content

            pconf_lines_by_type["IPCMAP"].append(
                f"/* {source_domain.name} IPCT ch. {source_ch} => "
                f"{sink_domain.name} IPCT ch. {sink_ch} */"
            )
            pconf_lines_by_type["IPCMAP"].append(
                (
                    "UICR_IPCMAP_CHANNEL_CFG("
                    f"{ipcmap_idx}, "
                    f"{source_domain.c_enum}, {source_ch}, "
                    f"{sink_domain.c_enum}, {sink_ch});"
                )
            )
            ipcmap_idx += 1

        for reg in uicr_svd["GPIOTE"]:
            if not reg["INSTANCE"].modified:
                break

            address = reg["INSTANCE"].content
            instance_name = periph_annotation_lookup.get(address, _c_hex_addr(address))
            spu_address = _get_spu_addr_for_periph(address)
            spu_name = _SPU_ADDR_TO_NAME[spu_address]

            for i, secure in _iter_owned_secure(reg["CH"]["OWN"], reg["CH"]["SECURE"]):
                if MigrateFlag.DEFAULTS not in flags and (secure, default_owner) == (
                    True,
                    OwnerID.APPLICATION,
                ):
                    continue

                pconf_lines_by_type["SPU_GPIOTE_CH"].append(
                    f"/* {spu_name} feature configuration for {instance_name} ch. {i} */"
                )
                pconf_lines_by_type["SPU_GPIOTE_CH"].append(
                    (
                        "UICR_SPU_FEATURE_GPIOTE_CH_SET("
                        f"{_c_hex_addr(spu_address)}, "
                        f"0, {i}, {_c_bool(secure)}, {default_owner.c_enum});"
                    )
                )

        for reg in uicr_svd["IPCT"]:
            if not reg["INSTANCE"].modified:
                break

            address = reg["INSTANCE"].content
            instance_name = periph_annotation_lookup.get(address, _c_hex_addr(address))
            spu_address = _get_spu_addr_for_periph(address)
            spu_name = _SPU_ADDR_TO_NAME[spu_address]

            for i, secure in _iter_owned_secure(reg["CH"]["OWN"], reg["CH"]["SECURE"]):
                if MigrateFlag.DEFAULTS not in flags and (secure, default_owner) == (
                    True,
                    OwnerID.APPLICATION,
                ):
                    continue

                pconf_lines_by_type["SPU_IPCT_CH"].append(
                    f"/* {spu_name} feature configuration for {instance_name} ch. {i} */"
                )
                pconf_lines_by_type["SPU_IPCT_CH"].append(
                    (
                        "UICR_SPU_FEATURE_IPCT_CH_SET("
                        f"{_c_hex_addr(spu_address)}, "
                        f"{i}, {_c_bool(secure)}, {default_owner.c_enum});"
                    )
                )

            for i, secure in _iter_owned_secure(
                reg["INTERRUPT"]["OWN"], reg["INTERRUPT"]["SECURE"]
            ):
                if MigrateFlag.DEFAULTS not in flags and (secure, default_owner) == (
                    True,
                    OwnerID.APPLICATION,
                ):
                    continue

                pconf_lines_by_type["SPU_IPCT_INT"].append(
                    f"/* {spu_name} feature configuration for {instance_name} interrupt {i} */"
                )
                pconf_lines_by_type["SPU_IPCT_INT"].append(
                    (
                        "UICR_SPU_FEATURE_IPCT_INTERRUPT_SET("
                        f"{_c_hex_addr(spu_address)}, "
                        f"{i}, {_c_bool(secure)}, {default_owner.c_enum});"
                    )
                )

        for reg in uicr_svd["DPPI"]:
            if not reg["INSTANCE"].modified:
                break

            address = reg["INSTANCE"].content
            instance_name = _DPPIC_ADDR_TO_NAME[address]
            if instance_name == "DPPIC130":
                continue

            spu_address = _get_spu_addr_for_periph(address)
            spu_name = _SPU_ADDR_TO_NAME[spu_address]

            for i, secure in _iter_owned_secure(reg["CH"]["OWN"], reg["CH"]["SECURE"]):
                if MigrateFlag.DEFAULTS not in flags and (secure, default_owner) == (
                    True,
                    OwnerID.APPLICATION,
                ):
                    continue

                pconf_lines_by_type["SPU_DPPI_CH"].append(
                    f"/* {spu_name} feature configuration for {instance_name} ch. {i} */"
                )
                pconf_lines_by_type["SPU_DPPI_CH"].append(
                    (
                        "UICR_SPU_FEATURE_DPPIC_CH_SET("
                        f"{_c_hex_addr(spu_address)}, "
                        f"{i}, {_c_bool(secure)}, {default_owner.c_enum});"
                    )
                )

            for i, secure in _iter_owned_secure(
                reg["CHG"]["OWN"], reg["CHG"]["SECURE"]
            ):
                if MigrateFlag.DEFAULTS not in flags and (secure, default_owner) == (
                    True,
                    OwnerID.APPLICATION,
                ):
                    continue

                pconf_lines_by_type["SPU_DDPI_CHG"].append(
                    f"/* {spu_name} feature configuration for {instance_name} ch. group {i} */"
                )
                pconf_lines_by_type["SPU_DPPI_CHG"].append(
                    (
                        "UICR_SPU_FEATURE_DPPIC_CHG_SET("
                        f"{_c_hex_addr(spu_address)}, "
                        f"{i}, {_c_bool(secure)}, {default_owner.c_enum});"
                    )
                )

            reg_link = reg["CH"]["LINK"]
            for ch, (enabled, direction) in enumerate(
                zip(reg_link["EN"].values(), reg_link["DIR"].values())
            ):
                if enabled.content_enum != "Enabled":
                    continue

                local_ppib_name, local_ppib_ch_map = _DPPIC_TO_PPIB[instance_name]
                local_ppib_ch = local_ppib_ch_map[ch]

                remote_ppib_name, remote_ppib_ch_map = _PPIB_TO_PPIB[local_ppib_name]
                remote_ppib_ch = remote_ppib_ch_map[local_ppib_ch]

                if direction.content_enum == "Source":
                    sub_name, sub_ch = local_ppib_name, local_ppib_ch
                    pub_name, pub_ch = remote_ppib_name, remote_ppib_ch
                else:
                    sub_name, sub_ch = remote_ppib_name, remote_ppib_ch
                    pub_name, pub_ch = local_ppib_name, local_ppib_ch

                sub_addr = _PPIB_NAME_TO_ADDR[sub_name]
                sub_final_name = periph_annotation_lookup.get(sub_addr, sub_name)
                pub_addr = _PPIB_NAME_TO_ADDR[pub_name]
                pub_final_name = periph_annotation_lookup.get(pub_addr, pub_name)

                pconf_lines_by_type["DPPI"].append(
                    f"/* {sub_final_name} ch. {sub_ch} => {pub_final_name} ch. {pub_ch} */"
                )
                pconf_lines_by_type["DPPI"].append(
                    f"UICR_PPIB_SUBSCRIBE_SEND_ENABLE({_c_hex_addr(sub_addr)}, {sub_ch});"
                )
                pconf_lines_by_type["DPPI"].append(
                    f"UICR_PPIB_PUBLISH_RECEIVE_ENABLE({_c_hex_addr(pub_addr)}, {pub_ch});"
                )

        grtc_spu_address = _get_spu_addr_for_periph(_GRTC_ADDR)
        grtc_spu_name = _SPU_ADDR_TO_NAME[grtc_spu_address]

        for i, secure in _iter_owned_secure(
            uicr_svd["GRTC"]["CC"]["OWN"], uicr_svd["GRTC"]["CC"]["SECURE"]
        ):
            if MigrateFlag.DEFAULTS in flags or (
                not secure or default_owner != OwnerID.APPLICATION
            ):
                pconf_lines_by_type["SPU_GRTC_CC"].append(
                    f"/* {grtc_spu_name} feature configuration for GRTC CC{i} */"
                )
                pconf_lines_by_type["SPU_GRTC_CC"].append(
                    (
                        "UICR_SPU_FEATURE_GRTC_CC_SET("
                        f"{_c_hex_addr(grtc_spu_address)}, "
                        f"{i}, {_c_bool(secure)}, {default_owner.c_enum});"
                    )
                )

        for i, reg in enumerate(uicr_svd["GPIO"]):
            if not reg["INSTANCE"].modified:
                break

            address = reg["INSTANCE"].content
            instance_name = periph_annotation_lookup.get(address, _c_hex_addr(address))
            spu_address = _get_spu_addr_for_periph(address)
            spu_name = _SPU_ADDR_TO_NAME[spu_address]
            gpio_port = _get_gpio_port_num(address)

            for pin, own in enumerate(reg["OWN"].values()):
                if own.content_enum != "Own":
                    continue

                secure = reg["SECURE"][f"PIN_{i}"].content_enum == "Secure"

                if MigrateFlag.DEFAULTS in flags or (
                    not secure or default_owner != OwnerID.APPLICATION
                ):
                    pconf_lines_by_type["SPU_GPIO"].append(
                        f"/* {spu_name} feature configuration for {instance_name}, "
                        f"P{gpio_port}.{pin} */"
                    )
                    pconf_lines_by_type["SPU_GPIO"].append(
                        (
                            "UICR_SPU_FEATURE_GPIO_PIN_SET("
                            f"{_c_hex_addr(spu_address)}, "
                            f"{gpio_port}, {pin}, {_c_bool(secure)}, {default_owner.c_enum});"
                        )
                    )

                ctrlsel = uicr_svd["GPIO_PIN"][i]["CTRLSEL"][pin]
                if not ctrlsel.modified:
                    continue

                ctrlsel_val = ctrlsel.content

                if MigrateFlag.DEFAULTS in flags or ctrlsel_val != CTRLSEL_DEFAULT:
                    pconf_lines_by_type["GPIO"].append(
                        f"/* {instance_name} - P{gpio_port}.{pin} CTRLSEL = {ctrlsel_val} */"
                    )
                    pconf_lines_by_type["GPIO"].append(
                        f"UICR_GPIO_PIN_CNF_CTRLSEL_SET({_c_hex_addr(address)}, "
                        f"{pin}, {ctrlsel.content});"
                    )

        for reg in uicr_svd["PERIPH"]:
            reg_cfg = reg["CONFIG"]
            if not reg_cfg.modified:
                break

            address = reg_cfg["ADDRESS"].content << reg_cfg["ADDRESS"].bit_offset
            irq_processor = ProcessorID(reg_cfg["PROCESSOR"].content)
            secure = reg_cfg["SECURE"].content_enum == "Secure"
            dma_secure = reg_cfg["DMASEC"].content_enum == "Secure"

            pconf_lines_by_type["SPU_PERIPH"].extend(
                _make_spu_periph_perm_set(
                    periph_addr=address,
                    secure=secure,
                    dma_secure=dma_secure,
                    owner=default_owner,
                    flags=flags,
                    periph_annotation_lookup=periph_annotation_lookup,
                )
            )
            pconf_lines_by_type["IRQMAP"].extend(
                _make_irqmap_irq_sink_set(
                    periph_addr=address,
                    irq_processor=irq_processor,
                    flags=flags,
                    periph_annotation_lookup=periph_annotation_lookup,
                )
            )

    full_lines = header_lines + list(chain.from_iterable(pconf_lines_by_type.values()))

    return "\n".join(full_lines)


def _iter_owned_secure(
    reg_own: svd.Register, reg_sec: svd.Register
) -> Iterable[tuple[int, bool]]:
    for i, (own, sec) in enumerate(zip(reg_own.values(), reg_sec.values())):
        owned = not own.content
        if not owned:
            continue

        yield i, sec.content_enum == SvdEnum.SECURE


_SPU_PERIPH_PERM_OFFSET = 0x500
_SPU_PERIPH_PERM_DMA_NODMA = 0


def _make_spu_periph_perm_set(
    periph_addr: int,
    secure: bool,
    dma_secure: bool,
    owner: OwnerID,
    flags: MigrateFlag,
    periph_annotation_lookup: dict[int, str],
) -> list[str]:
    periph_name = periph_annotation_lookup.get(periph_addr, _c_hex_addr(periph_addr))
    pconf_defaults = _load_pconf_defaults()

    if periph_addr in (_VPR121_ADDR, _VPR130_ADDR):
        log.debug(f"skip SPU.PERIPH[n].PERM for VPR peripheral {periph_name}")
        return []

    parsed_addr = Address(periph_addr)
    if parsed_addr.region == AddressRegion.STM:
        log.debug(f"skip SPU.PERIPH[n].PERM for TDD peripheral {periph_name}")
        return []

    spu_address = _get_spu_addr_for_periph(periph_addr)
    spu_name = _SPU_ADDR_TO_NAME[spu_address]
    periph_slave_index = parsed_addr.slave_index

    default_value = pconf_defaults[
        spu_address + _SPU_PERIPH_PERM_OFFSET + 4 * periph_slave_index
    ]
    default_fields = {
        "SECUREMAPPING": get_field(default_value, 0, 0b11),
        "DMA": get_field(default_value, 2, 0b11 << 2),
        "SECATTR": bool(get_field(default_value, 4, 0b1 << 4)),
        "DMASEC": bool(get_field(default_value, 5, 0b1 << 5)),
        "OWNERID": get_field(default_value, 16, 0b1111 << 16),
        "OWNERPROG": get_field(default_value, 30, 0b1 << 30),
    }

    if default_fields["OWNERID"] in (OwnerID.SECURE, OwnerID.SYSCTRL):
        log.debug(
            f"skip unassignable SPU.PERIPH[{periph_slave_index}].PERM "
            f"(default owner is {default_fields['OWNERID']})"
        )
        return []

    # If there is no DMA, use the default value to avoid readback errors
    if default_fields["DMA"] == _SPU_PERIPH_PERM_DMA_NODMA:
        dma_secure = default_fields["DMASEC"]
        log.debug(
            f"use default DMASEC={dma_secure} for SPU.PERIPH[{periph_slave_index}].PERM (noDMA)"
        )

    if (secure, dma_secure, owner) == (
        default_fields["SECATTR"],
        default_fields["DMASEC"],
        default_fields["OWNERID"],
    ) and MigrateFlag.DEFAULTS not in flags:
        log.debug(
            f"skip default configuration of SPU.PERIPH[{periph_slave_index}].PERM "
            f"(default values are {default_fields['SECATTR']=}, {default_fields['DMASEC']=}, "
            f"{default_fields['OWNERID']=})"
        )
        return []

    return [
        f"/* {spu_name} configuration for {periph_name} */",
        (
            "UICR_SPU_PERIPH_PERM_SET("
            f"{_c_hex_addr(spu_address)}, {periph_slave_index}, "
            f"{_c_bool(secure)}, {_c_bool(dma_secure)}, {owner.c_enum});"
        ),
    ]


_IRQMAP_ADDR = 0x5F92_4000
_IRQMAP_IRQ_SINK_OFFSET = 0x400


def _make_irqmap_irq_sink_set(
    periph_addr: int,
    irq_processor: ProcessorID,
    flags: MigrateFlag,
    periph_annotation_lookup: dict[int, str],
) -> list[str]:
    periph_name = periph_annotation_lookup.get(periph_addr, _c_hex_addr(periph_addr))
    pconf_defaults = _load_pconf_defaults()

    if periph_addr == _TBM_ADDR:
        periph_irqn = _TBM_IRQN
    else:
        periph_irqn = peripheral_id_get(periph_addr)

    default_value = pconf_defaults[
        _IRQMAP_ADDR + _IRQMAP_IRQ_SINK_OFFSET + 4 * periph_irqn
    ]
    default_processor_int = get_field(default_value, 8, 0b1111 << 8)

    if default_processor_int in (0, ProcessorID.SECURE, ProcessorID.SYSCTRL):
        log.debug(
            f"skip unassignable IRQMAP.IRQ[{periph_irqn}].SINK config "
            f"(default processor is {default_processor_int})"
        )
        return []

    if irq_processor == default_processor_int and MigrateFlag.DEFAULTS not in flags:
        log.debug(
            f"skip default IRQMAP.IRQ[{periph_irqn}].SINK config "
            f"(default processor is {default_processor_int})"
        )
        return []

    return [
        f"/* {periph_name} IRQ => {irq_processor.name} */",
        f"UICR_IRQMAP_IRQ_SINK_SET({periph_irqn}, " f"{irq_processor.c_enum});",
    ]


@lru_cache(1)
def _load_pconf_defaults() -> dict[int, int]:
    pconf_hex_path = importlib.resources.files("nrfregtool.resources").joinpath(
        "nrf54h20_periphconf_defaults.hex"
    )

    with importlib.resources.as_file(pconf_hex_path) as hex_path:
        with open(hex_path, "r", encoding="utf-8") as fp:
            pconf_hex = IntelHex()
            pconf_hex.loadhex(fp)

    defaults = {}
    pconf_addresses = pconf_hex.addresses()
    for i in range(0, len(pconf_addresses), 4):
        addresses = [pconf_addresses[i + j] for j in range(4)]
        value_parts = [pconf_hex[a] for a in addresses]
        value = int.from_bytes(value_parts, "little")
        defaults[addresses[0]] = value

    return defaults


def _get_spu_addr_for_periph(periph_addr: int) -> int:
    address = Address(periph_addr)
    address.slave_index = 0
    address.address_space = 0
    # Usually the above is sufficient to get the SPU address, but some buses
    # are special due to having > 16 slaves and need to be handled here.
    address.bus = _SPU_ADDRESS_BUS_REMAPPING.get(address.bus, address.bus)

    return int(address)


_SPU_ADDRESS_BUS_REMAPPING = {
    # Both of these bus IDs represent APB32 and should use the same SPU instance with bus ID 146.
    147: 146,
}


def _get_gpio_port_num(gpio_addr: int) -> int:
    return (gpio_addr - _P0_ADDR) // _GPIO_SIZE


_P0_ADDR = 0x5F93_8000
_GPIO_SIZE = 0x200

_GRTC_ADDR = 0x5F99_C000

_TDDCONF_ADDR = 0xBF00_1000

_TBM_ADDR = 0xBF00_3000
_TBM_IRQN = 127

_VPR121_ADDR = 0x5F8D_4000
_VPR130_ADDR = 0x5F90_8000


def _c_hex_addr(address: int) -> str:
    return f"0x{address:08x}UL"


def _c_bool(value: bool) -> str:
    return "true" if value else "false"


_SPU_ADDR_TO_NAME = {
    0x5F08_0000: "SPU110",
    0x5F09_0000: "SPU111",
    0x5F8C_0000: "SPU120",
    0x5F8D_0000: "SPU121",
    0x5F8E_0000: "SPU122",
    0x5F90_0000: "SPU130",
    0x5F92_0000: "SPU131",
    0x5F98_0000: "SPU132",
    0x5F99_0000: "SPU133",
    0x5F9A_0000: "SPU134",
    0x5F9B_0000: "SPU135",
    0x5F9C_0000: "SPU136",
    0x5F9D_0000: "SPU137",
}

_DPPIC_ADDR_TO_NAME = {
    0x5F8E_1000: "DPPIC120",
    0x5F92_2000: "DPPIC130",
    0x5F98_1000: "DPPIC131",
    0x5F99_1000: "DPPIC132",
    0x5F9A_1000: "DPPIC133",
    0x5F9B_1000: "DPPIC134",
    0x5F9C_1000: "DPPIC135",
    0x5F9D_1000: "DPPIC136",
}

_PPIB_NAME_TO_ADDR = {
    "PPIB121": 0x5F8E_F000,
    "PPIB130": 0x5F92_5000,
    "PPIB131": 0x5F92_6000,
    "PPIB132": 0x5F98_D000,
    "PPIB133": 0x5F99_D000,
    "PPIB134": 0x5F9A_D000,
    "PPIB135": 0x5F9B_D000,
    "PPIB136": 0x5F9C_D000,
    "PPIB137": 0x5F9D_D000,
}

_DPPIC_TO_PPIB = {
    "DPPIC120": ("PPIB121", range(0, 8)),
    "DPPIC131": ("PPIB132", range(0, 8)),
    "DPPIC132": ("PPIB133", range(0, 8)),
    "DPPIC133": ("PPIB134", range(0, 8)),
    "DPPIC134": ("PPIB135", range(0, 8)),
    "DPPIC135": ("PPIB136", range(0, 8)),
    "DPPIC136": ("PPIB137", range(0, 8)),
}

_PPIB_TO_PPIB = {
    "PPIB132": ("PPIB130", range(0, 8)),
    "PPIB133": ("PPIB130", range(8, 16)),
    "PPIB134": ("PPIB130", range(16, 24)),
    "PPIB135": ("PPIB130", range(24, 32)),
    "PPIB136": ("PPIB131", range(0, 8)),
    "PPIB137": ("PPIB131", range(8, 16)),
    "PPIB121": ("PPIB131", range(16, 24)),
}
