from saleae.analyzers import AnalyzerFrame
from SecondLayerHLAs import SecondLayerAnalyzer
from collections import namedtuple
import logging

MCTP_MSG_TYPES = {
    0x00: "MCTP Control",
    0x01: "PLDM",
    0x02: "NCSI",
    0x03: "Ethernet",
    0x04: "NVMe-MI",
    0x05: "CXL-FM-API",
    0x06: "CXL-CCI",
    0x7E: "Vendor (IANA)",
    0x7F: "Vendor (PCI)",
}

# MIPI Mandatory Data Byte value for MCTP pending read notification (DSP0233 §5.2.2.1)
MCTP_PENDING_READ_MDB = 0xAE

_Parsed = namedtuple("_Parsed", [
    "data_array", "base", "hdr_ver", "dest_eid", "src_eid",
    "som", "eom", "seq", "to", "tag", "ic", "msg_type_name",
    "payload", "pec_status", "direction", "key",
])


def _crc8(data: list) -> int:
    """CRC-8 per SMBus PEC: polynomial x^8+x^2+x+1 (0x07), init=0, no final XOR.
    DSP0233 §5.3.1: covers the I3C address byte + all MCTP packet bytes,
    excluding 7'h7E broadcast address, Start/Stop, T-bits, ACK/NACK.
    """
    crc = 0
    for byte in data:
        crc ^= byte
        for _ in range(8):
            crc = ((crc << 1) ^ 0x07) if (crc & 0x80) else (crc << 1)
            crc &= 0xFF
    return crc


class MCTPBase(SecondLayerAnalyzer.SecondLayerHLA):
    """Shared MCTP header parsing, state tracking, and orchestration.
    Subclasses override _on_packet() and optionally _make_orphan_frame() to
    control what frame is emitted for each packet.
    """

    MCTP_logger = logging.getLogger("MCTP")

    def __init__(self):
        # Keyed by (src_eid, dest_eid, tag, to) — uniquely identifies an in-progress multi-packet message
        self._partial = {}

    def update_ccc(self, ccc_dict: dict) -> None:
        pass

    def _make_incomplete_frame(self, entry: dict, start_time, end_time) -> AnalyzerFrame:
        """Build an 'MCTP Incomplete Message' frame from a partial-message entry.
        Length is an integer so SDRDecoder callers can compare it with > 0.
        """
        return AnalyzerFrame("MCTP Incomplete Message", start_time, end_time, {
            "DestEID":   f"0x{entry['dest_eid']:02X}",
            "SrcEID":    f"0x{entry['src_eid']:02X}",
            "Tag":       str(entry["tag"]),
            "MsgType":   entry["msg_type"],
            "Direction": entry["direction"],
            "Data":      ", ".join(f"0x{b:02X}" for b in entry["data"]),
            "Length":    entry["length"],   # integer; SDRDecoder converts to str after > 0 check
            "Packets":   str(entry["packets"]),
        })

    def _make_orphan_frame(self, orphan_entry: dict, current_frame: AnalyzerFrame) -> AnalyzerFrame:
        """Build incomplete frame for an orphaned in-progress message.
        Uses current_frame.start_time by default (safe since prior packet frames were already emitted).
        MCTPMessageHla overrides this to use orphan_entry["start_time"] for proper time spanning.
        """
        return self._make_incomplete_frame(orphan_entry, current_frame.start_time, current_frame.end_time)

    def _parse_frame(self, frame: AnalyzerFrame):
        """Parse MCTP header fields from a completed I3C transaction frame.
        Returns a _Parsed namedtuple, or None if the frame is not MCTP-bearing or is malformed.
        Reads self._partial for msg_type lookup on non-SOM packets but does not modify it.
        """
        if frame.type not in ("Message", "I3C Transaction", "I3C IBI Transaction"):
            return None

        try:
            data_array = [int(x, 16) for x in frame.data["Data"].split(",")]
        except (ValueError, KeyError):
            return None

        # Byte layout (DSP0233 Table 1 / Table 3):
        #   "Message"         → data[0]=addr,  data[1]=MCTP hdr ver, ..., data[-1]=PEC
        #   "I3C Transaction" → data[0]=0x7E,  data[1]=addr, data[2]=MCTP hdr ver, ..., data[-1]=PEC
        base = 2 if frame.type == "I3C Transaction" else 1

        # Minimum length: MCTP header (4 bytes at base) + PEC (1 byte at end)
        if len(data_array) < base + 5:
            return None

        # PEC validation (DSP0233 §5.3.1): reconstruct the 8-bit wire address byte
        # (7-bit addr << 1 | R/W bit) and prepend it to the CRC input.
        # The 0x7E broadcast marker is excluded per spec; for "I3C Transaction" frames
        # the target address is at data_array[1]; for other frames at data_array[0].
        wr_field = frame.data.get("W/R")
        if not wr_field:
            self.MCTP_logger.warning(
                f"MCTP PEC: W/R field missing in {frame.type!r} frame; "
                f"defaulting to Write for CRC computation"
            )
        direction = (wr_field or "W")[0]
        addr_byte = data_array[1] if frame.type == "I3C Transaction" else data_array[0]
        wire_byte = (addr_byte << 1) | (1 if direction == "R" else 0)
        pec_ok = _crc8([wire_byte] + data_array[base:-1]) == data_array[-1]
        pec_status = "OK" if pec_ok else "Error"

        hdr_ver = data_array[base] & 0x0F
        if hdr_ver != 0x01:
            self.MCTP_logger.warning(f"MCTP unexpected header version: 0x{hdr_ver:X} (expected 0x01)")
        # Decode regardless — be permissive with non-compliant devices

        dest_eid = data_array[base + 1]
        src_eid  = data_array[base + 2]
        flags    = data_array[base + 3]
        som      = (flags >> 7) & 0x1
        eom      = (flags >> 6) & 0x1
        seq      = (flags >> 4) & 0x3
        to       = (flags >> 3) & 0x1
        tag      = flags & 0x7

        # SOM=1 requires one additional byte for IC|MsgType before PEC
        if som and len(data_array) < base + 6:
            return None

        key = (src_eid, dest_eid, tag, to)

        # IC bit and message type (IC|MsgType byte only present on SOM=1, per spec)
        if som:
            type_byte     = data_array[base + 4]
            ic            = (type_byte >> 7) & 0x1
            msg_type_val  = type_byte & 0x7F
            msg_type_name = MCTP_MSG_TYPES.get(msg_type_val, f"0x{msg_type_val:02X}")
            payload       = data_array[base + 5:-1]
        else:
            ic            = 0
            msg_type_name = self._partial.get(key, {}).get("msg_type", "Unknown")
            payload       = data_array[base + 4:-1]

        return _Parsed(
            data_array=data_array, base=base, hdr_ver=hdr_ver,
            dest_eid=dest_eid, src_eid=src_eid,
            som=som, eom=eom, seq=seq, to=to, tag=tag,
            ic=ic, msg_type_name=msg_type_name, payload=payload,
            pec_status=pec_status, direction=direction, key=key,
        )

    def _on_packet(self, frame, parsed, pkt_frame, orphan_frame, eom_entry):
        """Decide what frame to emit for a successfully parsed MCTP packet.
        Called by instruction_decoder after all _partial state transitions are done.

        Args:
            frame:        The raw AnalyzerFrame from the I3C HLA.
            parsed:       _Parsed namedtuple with all decoded MCTP fields.
            pkt_frame:    Pre-built AnalyzerFrame("MCTP Packet", ...) for this packet.
            orphan_frame: Pre-built AnalyzerFrame("MCTP Incomplete Message", ...) for an
                          orphaned in-progress message, or None.
            eom_entry:    The popped _partial entry dict if this packet is EOM, or None.
        """
        raise NotImplementedError

    def instruction_decoder(self, frame: AnalyzerFrame):
        # IBI: detect MCTP pending read notification (MDB = 0xAE, DSP0233 §5.2.2.1).
        # For other IBI payloads, fall through to the MCTP packet decoder — some devices
        # embed MCTP data directly in the IBI payload with the MCTP header version byte
        # serving as the MDB, rather than using the separate Private Read flow.
        if frame.type == "I3C IBI Transaction":
            try:
                ibi_bytes = [int(x, 16) for x in frame.data["Data"].split(",")]
                if len(ibi_bytes) >= 2 and ibi_bytes[1] == MCTP_PENDING_READ_MDB:
                    return AnalyzerFrame("MCTP Pending Read Notification",
                        frame.start_time, frame.end_time, {
                            "IBI_Addr": f"0x{ibi_bytes[0]:02X}",
                            "MDB":      f"0x{MCTP_PENDING_READ_MDB:02X}",
                        })
            except (ValueError, KeyError, IndexError):
                pass
            # Fall through to MCTP packet decode

        parsed = self._parse_frame(frame)
        if parsed is None:
            return frame

        p = parsed
        pkt_frame = AnalyzerFrame("MCTP Packet", frame.start_time, frame.end_time, {
            "DestEID":   f"0x{p.dest_eid:02X}",
            "SrcEID":    f"0x{p.src_eid:02X}",
            "HdrVer":    str(p.hdr_ver),
            "SOM":       str(p.som),
            "EOM":       str(p.eom),
            "Seq":       str(p.seq),
            "TO":        str(p.to),
            "Tag":       str(p.tag),
            "IC":        str(p.ic),
            "MsgType":   p.msg_type_name,
            "PEC":       p.pec_status,
            "Direction": p.direction,
            "Data":      ", ".join(f"0x{b:02X}" for b in p.payload),
            "Length":    str(len(p.payload)),
        })

        self.MCTP_logger.debug(
            f"MCTP Pkt SOM={p.som} EOM={p.eom} Seq={p.seq} Tag={p.tag} "
            f"Dest=0x{p.dest_eid:02X} Src=0x{p.src_eid:02X} Type={p.msg_type_name} PEC={p.pec_status}"
        )

        # Orphaned partial: new SOM for a key that already has an in-progress message.
        orphan_frame = None
        if p.som and p.key in self._partial:
            orphan = self._partial.pop(p.key)
            self.MCTP_logger.warning(
                f"MCTP orphaned partial: "
                f"Src=0x{orphan['src_eid']:02X} Dest=0x{orphan['dest_eid']:02X} "
                f"Tag={orphan['tag']} TO={orphan.get('to', '?')}"
            )
            orphan_frame = self._make_orphan_frame(orphan, frame)

        # _partial state transitions
        if p.som:
            self._partial[p.key] = {
                "dest_eid":   p.dest_eid,
                "src_eid":    p.src_eid,
                "tag":        p.tag,
                "to":         p.to,
                "msg_type":   p.msg_type_name,
                "direction":  p.direction,
                "data":       list(p.payload),
                "length":     len(p.payload),
                "packets":    1,
                "ic":         p.ic,
                "seq":        p.seq,
                "start_time": frame.start_time,
            }
        elif p.key in self._partial:
            expected_seq = (self._partial[p.key]["seq"] + 1) % 4
            if p.seq != expected_seq:
                orphan = self._partial.pop(p.key)
                self.MCTP_logger.warning(
                    f"MCTP seq error discarded partial: "
                    f"Src=0x{orphan['src_eid']:02X} Tag={orphan['tag']} "
                    f"expected={expected_seq} got={p.seq}"
                )
                return AnalyzerFrame("MCTP Seq Error", frame.start_time, frame.end_time, {
                    "DestEID":     f"0x{p.dest_eid:02X}",
                    "SrcEID":      f"0x{p.src_eid:02X}",
                    "Tag":         str(p.tag),
                    "Direction":   p.direction,
                    "ExpectedSeq": str(expected_seq),
                    "Seq":         str(p.seq),
                })
            self._partial[p.key]["seq"] = p.seq
            self._partial[p.key]["data"].extend(p.payload)
            self._partial[p.key]["length"] += len(p.payload)
            self._partial[p.key]["packets"] += 1
        else:
            # Non-SOM packet with no matching SOM in flight — pass through as original frame
            # to suppress false-positive MCTP Packet displays for unrelated bus traffic.
            return frame

        # Return the orphan's incomplete-message frame if we detected one.
        # The new SOM has already been stored in _partial above.
        # Edge case: if the new SOM is also EOM (single-packet), pop it now so it does
        # not linger in _partial — we cannot return two frames, so the new complete
        # message is sacrificed in favour of making the orphan visible.
        if orphan_frame is not None:
            if p.eom:
                dropped = self._partial.pop(p.key, None)
                if dropped:
                    self.MCTP_logger.warning(
                        f"MCTP complete message dropped (orphan takes priority): "
                        f"Src=0x{dropped['src_eid']:02X} Dest=0x{dropped['dest_eid']:02X} "
                        f"Tag={dropped['tag']} TO={dropped.get('to','?')} "
                        f"Type={dropped['msg_type']}"
                    )
            return self._on_packet(frame, parsed, pkt_frame, orphan_frame, None)

        # Pop partial on EOM and pass to _on_packet as eom_entry
        eom_entry = None
        if p.eom and p.key in self._partial:
            eom_entry = self._partial.pop(p.key)

        return self._on_packet(frame, parsed, pkt_frame, None, eom_entry)

    def handle_abort_message(self, temp_frame: AnalyzerFrame, end_time) -> AnalyzerFrame:
        if not self._partial:
            return temp_frame

        # Emit the first in-flight partial as 'MCTP Incomplete Message'.
        # The interface allows only one return value; if there are multiple partials,
        # log a warning for the extras.
        if len(self._partial) > 1:
            self.MCTP_logger.warning(
                f"MCTP abort: {len(self._partial)} in-progress partials; "
                f"emitting first as MCTP Incomplete Message, discarding the rest"
            )
        entry = next(iter(self._partial.values()))
        self._partial.clear()
        # Return integer Length so SDRDecoder's `> 0` check works; SDRDecoder converts to str.
        return self._make_incomplete_frame(entry, temp_frame.start_time, end_time)


class MCTPPacketHla(MCTPBase):
    """MCTP HLA that emits only individual 'MCTP Packet' frames.
    Use alongside MCTPMessageHla (add the analyzer twice in Saleae) to see both
    packet-level and message-level bubbles simultaneously.
    """

    def update_result_types(self, result_types: dict) -> None:
        result_types["MCTP Packet"] = {
            'format':  'MCTP Pkt [{{data.Direction}}] Dest [{{data.DestEID}}] Src [{{data.SrcEID}}] HdrVer [{{data.HdrVer}}] SOM [{{data.SOM}}] EOM [{{data.EOM}}] Seq [{{data.Seq}}] Tag [{{data.Tag}}] TO [{{data.TO}}] IC [{{data.IC}}] Type [{{data.MsgType}}] PEC [{{data.PEC}}] [{{data.Length}}]: [ {{data.Data}} ]',
            'console': 'MCTP Pkt [{Direction}] Dest [{DestEID}] Src [{SrcEID}] HdrVer [{HdrVer}] SOM [{SOM}] EOM [{EOM}] Seq [{Seq}] Tag [{Tag}] TO [{TO}] IC [{IC}] Type [{MsgType}] PEC [{PEC}] [{Length}]: [ {Data} ]',
        }
        result_types["MCTP Incomplete Message"] = {
            'format':  'MCTP Incomplete Msg [{{data.Direction}}] Dest [{{data.DestEID}}] Src [{{data.SrcEID}}] Tag [{{data.Tag}}] Type [{{data.MsgType}}] [{{data.Packets}} pkts] [{{data.Length}}]: [ {{data.Data}} ]',
            'console': 'MCTP Incomplete Msg [{Direction}] Dest [{DestEID}] Src [{SrcEID}] Tag [{Tag}] Type [{MsgType}] [{Packets} pkts] [{Length}]: [ {Data} ]',
        }
        result_types["MCTP Seq Error"] = {
            'format':  'MCTP Seq Error [{{data.Direction}}] Dest [{{data.DestEID}}] Src [{{data.SrcEID}}] Tag [{{data.Tag}}] Expected Seq [{{data.ExpectedSeq}}] Got [{{data.Seq}}]',
            'console': 'MCTP Seq Error [{Direction}] Dest [{DestEID}] Src [{SrcEID}] Tag [{Tag}] Expected Seq [{ExpectedSeq}] Got [{Seq}]',
        }
        result_types["MCTP Pending Read Notification"] = {
            'format':  'MCTP Pending Read [{{data.IBI_Addr}}] MDB [{{data.MDB}}]',
            'console': 'MCTP Pending Read [{IBI_Addr}] MDB [{MDB}]',
        }

    def _on_packet(self, frame, parsed, pkt_frame, orphan_frame, eom_entry):
        return orphan_frame if orphan_frame is not None else pkt_frame


class MCTPMessageHla(MCTPBase):
    """MCTP HLA that emits only 'MCTP Message' frames, spanning from the SOM packet's
    start_time to the EOM packet's end_time.  Returns None for SOM and middle packets so
    that Saleae permits the spanning timestamp without backward-time violations.
    Use alongside MCTPPacketHla (add the analyzer twice in Saleae) to see both layers.
    """

    def update_result_types(self, result_types: dict) -> None:
        result_types["MCTP Message"] = {
            'format':  'MCTP Msg [{{data.Direction}}] Dest [{{data.DestEID}}] Src [{{data.SrcEID}}] Tag [{{data.Tag}}] Type [{{data.MsgType}}] IC [{{data.IC}}] PEC [{{data.PEC}}] HdrVer [{{data.HdrVer}}] [{{data.Packets}} pkts] [{{data.Length}}]: [ {{data.Data}} ]',
            'console': 'MCTP Msg [{Direction}] Dest [{DestEID}] Src [{SrcEID}] Tag [{Tag}] Type [{MsgType}] IC [{IC}] PEC [{PEC}] HdrVer [{HdrVer}] [{Packets} pkts] [{Length}]: [ {Data} ]',
        }
        result_types["MCTP Incomplete Message"] = {
            'format':  'MCTP Incomplete Msg [{{data.Direction}}] Dest [{{data.DestEID}}] Src [{{data.SrcEID}}] Tag [{{data.Tag}}] Type [{{data.MsgType}}] [{{data.Packets}} pkts] [{{data.Length}}]: [ {{data.Data}} ]',
            'console': 'MCTP Incomplete Msg [{Direction}] Dest [{DestEID}] Src [{SrcEID}] Tag [{Tag}] Type [{MsgType}] [{Packets} pkts] [{Length}]: [ {Data} ]',
        }
        result_types["MCTP Seq Error"] = {
            'format':  'MCTP Seq Error [{{data.Direction}}] Dest [{{data.DestEID}}] Src [{{data.SrcEID}}] Tag [{{data.Tag}}] Expected Seq [{{data.ExpectedSeq}}] Got [{{data.Seq}}]',
            'console': 'MCTP Seq Error [{Direction}] Dest [{DestEID}] Src [{SrcEID}] Tag [{Tag}] Expected Seq [{ExpectedSeq}] Got [{Seq}]',
        }
        result_types["MCTP Pending Read Notification"] = {
            'format':  'MCTP Pending Read [{{data.IBI_Addr}}] MDB [{{data.MDB}}]',
            'console': 'MCTP Pending Read [{IBI_Addr}] MDB [{MDB}]',
        }

    def _make_orphan_frame(self, orphan_entry: dict, current_frame: AnalyzerFrame) -> AnalyzerFrame:
        # Span from orphan's stored SOM start_time — no frames were emitted for that region.
        return self._make_incomplete_frame(
            orphan_entry, orphan_entry["start_time"], current_frame.end_time
        )

    def _on_packet(self, frame, parsed, pkt_frame, orphan_frame, eom_entry):
        if orphan_frame is not None:
            return orphan_frame
        if eom_entry is not None:
            # Emit message spanning from SOM start_time to EOM end_time.
            # Valid because no frames were emitted for SOM/middle packets (they returned None).
            return AnalyzerFrame("MCTP Message", eom_entry["start_time"], frame.end_time, {
                "DestEID":   f"0x{eom_entry['dest_eid']:02X}",
                "SrcEID":    f"0x{eom_entry['src_eid']:02X}",
                "Tag":       str(eom_entry["tag"]),
                "MsgType":   eom_entry["msg_type"],
                "Direction": eom_entry["direction"],
                "Data":      ", ".join(f"0x{b:02X}" for b in eom_entry["data"]),
                "Length":    str(eom_entry["length"]),
                "Packets":   str(eom_entry["packets"]),
                "PEC":       parsed.pec_status,
                "IC":        str(eom_entry.get("ic", parsed.ic)),
                "HdrVer":    str(parsed.hdr_ver),
            })
        return None  # SOM-only or middle packet — suppress, wait for EOM

    def handle_abort_message(self, temp_frame: AnalyzerFrame, end_time) -> AnalyzerFrame:
        if not self._partial:
            return temp_frame
        if len(self._partial) > 1:
            self.MCTP_logger.warning(
                f"MCTP abort: {len(self._partial)} in-progress partials; "
                f"emitting first as MCTP Incomplete Message, discarding the rest"
            )
        entry = next(iter(self._partial.values()))
        self._partial.clear()
        # Span from stored SOM start_time — no frames were emitted for that region.
        return self._make_incomplete_frame(entry, entry["start_time"], end_time)
