# SPDX-FileCopyrightText: 2024 Justin Simon <justin@simonctl.com>
#
# SPDX-License-Identifier: MIT

import base64
import dataclasses
import json
import uuid
from dataclasses import field
from enum import IntEnum
from pathlib import Path
from typing import Any, Protocol, TypeVar, runtime_checkable

from mashumaro import DataClassDictMixin, field_options
from mashumaro.config import BaseConfig

from ..interfaces import AnyPacketType


class MsgTypes(IntEnum):
    """
    Messages used to support initialization and configuration of MCTP communication within an MCTP network, as
    specified in DSP0236
    """

    CTRL = 0x00
    PLDM = 0x01
    NCSI = 0x02
    Ethernet = 0x03
    NVMeMgmtMsg = 0x04
    SPDM = 0x05
    SECUREDMSG = 0x06
    CXL_FM_API = 0x07
    CXL_CCI = 0x08
    VDPCI = 0x7E
    VDIANA = 0x7F


class VendorIdFormat(IntEnum):
    PCI_VENDOR_ID = 0x0
    """16-bit Unsigned Integer that identifies the manufacturer
    of the device. Valid vendor identifiers are allocated by the
    PCI SIG to ensure uniqueness."""

    IANA_ENT_NUMBER = 0x1
    """The IANA enterprise number for the organization or vendor
    expressed as a 32-bit unsigned binary number."""


@dataclasses.dataclass(frozen=True)
class VendorCapabilitySet(DataClassDictMixin):
    vendor_id: int
    id_format: VendorIdFormat = VendorIdFormat.PCI_VENDOR_ID
    command_set_type: int = 0


class PhysicalTransportBindingId(IntEnum):
    Reserved = 0x0
    MCTPoverSMBus = 0x01
    MCTPoverPcieVdm = 0x02
    MCTPoverKCS = 0x04
    MCTPoverSerial = 0x05
    MCTPoverI3C = 0x06
    VendorDefined = 0xFF


class PhysicalMediumIdentifiers(IntEnum):
    Unspecified = 0x0
    SMBUS_2_0_100khz = 0x1
    SMBUS_2_0_I2C_100khz = 0x2
    I2C_100khz = 0x3
    SMBUS_3_0_I2C_400khz = 0x4
    SMBUS_3_0_I2C_1mhz = 0x5
    I2C_3_4mhz = 0x6
    PCIeRev_1_1 = 0x8
    PCIeRev_2_0 = 0x9
    PCIeRev_2_1 = 0xA
    PCIeRev_3 = 0xB
    PCIeRev_4 = 0xC
    PCIeRev_5 = 0xD
    PCICompatible = 0xF
    USB_1_1 = 0x10
    USB_2 = 0x11
    USB_3 = 0x12
    NCSIOverRBT = 0x18
    KCSLegacy = 0x20
    KCSPCI = 0x21
    SerialHostLegacy = 0x22
    SerialHostPCI = 0x23
    AsyncSerial = 0x24
    I3CBasic = 0x30


class EntryType(IntEnum):
    SINGLE_ENDPOINT = 0
    BRIDGE_AND_DOWNSTREAM_ENDPOINT = 1
    SINGLE_BRIDGE_ENDPOINT = 2
    ADDITIONAL_BRIDGE_EID_RANGE = 3


@dataclasses.dataclass(frozen=True)
class RoutingTableEntry(DataClassDictMixin):
    starting_eid: int
    port_number: int
    phy_address: list[int]
    phys_transport_binding_id: PhysicalTransportBindingId = PhysicalTransportBindingId.MCTPoverSMBus
    phy_media_type_id: PhysicalMediumIdentifiers = PhysicalMediumIdentifiers.I2C_100khz
    entry_type: EntryType = EntryType.SINGLE_ENDPOINT
    eid_range: int = 1
    static_eid: bool = False


@dataclasses.dataclass(frozen=True)
class Smbus7bitAddress(DataClassDictMixin):
    address: int

    def __post_init__(self):
        if self.address > 0x7F:
            msg = f"SMBUS address {hex(self.address)} is more than 7 bits."
            raise ValueError(msg)

    def read(self) -> int:
        return self.address << 1 | 0x01

    def write(self) -> int:
        return self.address << 1 | 0x00


@dataclasses.dataclass(frozen=True)
class Smbus10bitAddress(DataClassDictMixin):
    address: int

    def __post_init__(self):
        if self.address > 0x3FF:
            msg = f"SMBUS address {hex(self.address)} is more than 10 bits."
            raise ValueError(msg)


AnyPhysicalAddress = TypeVar("AnyPhysicalAddress", Smbus7bitAddress, Smbus10bitAddress)


@dataclasses.dataclass(frozen=True)
class MctpResponse(DataClassDictMixin):
    request: list[int]
    response: list[int]
    processing_delay: int = field(metadata=field_options(alias="processing-delay"))
    key: bytes = field(init=False)
    data: bytes = field(init=False)
    response_index: int = 0
    description: str = "<missing>"

    def __post_init__(self):
        req = bytes(self.request)
        key = base64.b64encode(req)
        object.__setattr__(self, "key", key)

        packet = bytes(self.response)
        object.__setattr__(self, "data", packet)


# @dataclasses.dataclass(frozen=True)
# class VdPciResponeList(DataClassDictMixin):
#     vendors: Dict[str, Dict[str, List[MctpResponse]]]


@dataclasses.dataclass(frozen=True)
class MctpResponseList(DataClassDictMixin):
    responses: dict[MsgTypes, list[MctpResponse] | dict[str, dict[str, list[MctpResponse]]]]


def deserialize_msg_types(values: list[MsgTypes] | list[str]) -> list[MsgTypes]:
    if not isinstance(values, list):
        msg = f"Unknown msg type: {type(values)}"
        raise TypeError(msg)
    if all(isinstance(y, MsgTypes) for y in values):
        return values
    msg_types = []
    for mt in values:
        if isinstance(mt, str) and hasattr(MsgTypes, mt):
            msg_types += [MsgTypes[mt]]
            continue
        msg = f"Unknown msg type: {mt}"
        raise TypeError(msg)
    return msg_types


@dataclasses.dataclass
class EndpointContext(DataClassDictMixin):
    physical_address: Smbus7bitAddress
    static_eid: int | None = None
    assigned_eid: int = 0
    discovered: bool = False
    is_bus_owner: bool = False
    pool_size: int = 0
    mtu_size: int = 240
    allocated_pool: list[int] | None = None
    endpoint_uuid: uuid.UUID = dataclasses.field(default_factory=lambda: uuid.uuid4())
    supported_msg_types: list[MsgTypes] = dataclasses.field(default_factory=lambda: [MsgTypes.CTRL])
    supported_vdm_msg_types: list[VendorCapabilitySet] = dataclasses.field(default_factory=list)
    mctp_responses: MctpResponseList | None = None
    routing_table_ready: bool = False
    routing_table: list[RoutingTableEntry] = dataclasses.field(default_factory=list)

    class Config(BaseConfig):
        serialization_strategy = {list[MsgTypes]: {"deserialize": deserialize_msg_types}}

    @property
    def eid(self) -> int:
        if self.assigned_eid:
            return self.assigned_eid
        if self.static_eid:
            return self.static_eid
        # no EID assigned, just return the null EID
        return 0

    def import_json_responses(self, resp_file: Path):
        with resp_file.open("r") as f:
            data: dict[str, Any] = json.load(f)
            data = {MsgTypes[k]: v for k, v in data.items() if MsgTypes[k] in self.supported_msg_types}
            self.mctp_responses = MctpResponseList.from_dict({"responses": data})

    def get_response(
        self, msg_type: MsgTypes, req_bytes: bytes, vendor_id: str = "", vdm_cmd_code: str = ""
    ) -> MctpResponse | None:
        if msg_type not in MsgTypes:
            return None
        if msg_type not in self.mctp_responses.responses:
            return None
        key = base64.b64encode(req_bytes)
        if msg_type != MsgTypes.VDPCI:
            for resp in self.mctp_responses.responses[msg_type]:
                if resp.key == key:
                    return resp
        else:
            responses = self.mctp_responses.responses[msg_type]
            if vendor_id not in responses:
                return None
            responses = responses[vendor_id]
            if vdm_cmd_code not in responses:
                return None
            for resp in responses[vdm_cmd_code]:
                if resp.key == key:
                    return resp
        return None


@runtime_checkable
class ICanReply(Protocol):
    def is_request(self, check_payload: bool = True) -> bool:
        pass

    def make_reply(self, ctx: EndpointContext) -> AnyPacketType:
        pass
