"""AXI4 protocol decoder."""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Optional

from ewal.protocol.base import Transaction, Violation

# AXI response codes
OKAY = 0
EXOKAY = 1
SLVERR = 2
DECERR = 3

_RESP_NAMES = {OKAY: "OKAY", EXOKAY: "EXOKAY", SLVERR: "SLVERR", DECERR: "DECERR"}


@dataclass
class AXI4WriteTxn:
    addr: int
    data: list[int]
    strb: list[int]
    length: int
    resp: int
    start: int
    end: int

    @property
    def resp_name(self) -> str:
        return _RESP_NAMES.get(self.resp, f"UNKNOWN({self.resp})")


@dataclass
class AXI4ReadTxn:
    addr: int
    data: list[int]
    length: int
    resp: int
    start: int
    end: int

    @property
    def resp_name(self) -> str:
        return _RESP_NAMES.get(self.resp, f"UNKNOWN({self.resp})")


@dataclass
class AXI4Violation:
    type: str
    cycle: int
    message: str = ""


class AXI4:
    """AXI4 protocol decoder."""

    OKAY = OKAY
    EXOKAY = EXOKAY
    SLVERR = SLVERR
    DECERR = DECERR

    @staticmethod
    def bind(**signals) -> "BoundAXI4":
        return BoundAXI4(signals)


class BoundAXI4:
    """AXI4 protocol bound to trace signals."""

    def __init__(self, signals: dict):
        self.signals = signals
        self._trace = None
        for sig in signals.values():
            if hasattr(sig, "_trace"):
                self._trace = sig._trace
                break

    def _get(self, name):
        return self.signals.get(name)

    def _val(self, name, index):
        sig = self._get(name)
        if sig is None:
            return None
        old = self._trace._index
        self._trace._index = index
        self._trace._virtual_cache.clear()
        v = sig.value
        self._trace._index = old
        return v

    def write_transactions(self) -> list[AXI4WriteTxn]:
        """Extract all write transactions."""
        t = self._trace
        if t is None:
            return []

        awvalid = self._get("awvalid")
        awready = self._get("awready")
        wvalid = self._get("wvalid")
        wready = self._get("wready")
        bvalid = self._get("bvalid")
        bready = self._get("bready")

        if not all([awvalid, awready, wvalid, wready, bvalid, bready]):
            return []

        # Find AW handshakes
        aw_indices = t.find(awvalid & awready)
        # Find W handshakes
        w_indices = t.find(wvalid & wready)
        # Find B handshakes
        b_indices = t.find(bvalid & bready)

        results = []
        w_ptr = 0
        b_ptr = 0

        for aw_idx in aw_indices:
            addr = self._val("awaddr", aw_idx) or 0
            awlen = self._val("awlen", aw_idx)
            burst_len = (awlen + 1) if awlen is not None else 1

            # Collect data beats
            data_beats = []
            strb_beats = []
            for _ in range(burst_len):
                if w_ptr < len(w_indices) and w_indices[w_ptr] >= aw_idx:
                    w_idx = w_indices[w_ptr]
                    data_beats.append(self._val("wdata", w_idx) or 0)
                    strb_beats.append(self._val("wstrb", w_idx) or 0xF)
                    w_ptr += 1
                elif w_ptr < len(w_indices):
                    w_ptr += 1

            # Find response
            resp = OKAY
            end_idx = aw_idx
            if b_ptr < len(b_indices):
                while b_ptr < len(b_indices) and b_indices[b_ptr] < aw_idx:
                    b_ptr += 1
                if b_ptr < len(b_indices):
                    b_idx = b_indices[b_ptr]
                    resp = self._val("bresp", b_idx) or OKAY
                    end_idx = b_idx
                    b_ptr += 1

            results.append(AXI4WriteTxn(
                addr=addr, data=data_beats, strb=strb_beats,
                length=burst_len, resp=resp,
                start=aw_idx, end=end_idx,
            ))

        return results

    def read_transactions(self) -> list[AXI4ReadTxn]:
        """Extract all read transactions."""
        t = self._trace
        if t is None:
            return []

        arvalid = self._get("arvalid")
        arready = self._get("arready")
        rvalid = self._get("rvalid")
        rready = self._get("rready")

        if not all([arvalid, arready, rvalid, rready]):
            return []

        ar_indices = t.find(arvalid & arready)
        r_indices = t.find(rvalid & rready)

        results = []
        r_ptr = 0

        for ar_idx in ar_indices:
            addr = self._val("araddr", ar_idx) or 0
            arlen = self._val("arlen", ar_idx)
            burst_len = (arlen + 1) if arlen is not None else 1

            data_beats = []
            resp = OKAY
            end_idx = ar_idx

            for _ in range(burst_len):
                if r_ptr < len(r_indices) and r_indices[r_ptr] >= ar_idx:
                    r_idx = r_indices[r_ptr]
                    data_beats.append(self._val("rdata", r_idx) or 0)
                    resp = self._val("rresp", r_idx) or OKAY
                    end_idx = r_idx
                    r_ptr += 1

            results.append(AXI4ReadTxn(
                addr=addr, data=data_beats, length=burst_len,
                resp=resp, start=ar_idx, end=end_idx,
            ))

        return results

    def check(self) -> list[AXI4Violation]:
        """Check for AXI4 protocol violations."""
        t = self._trace
        if t is None:
            return []

        violations = []

        # Check: VALID must not depend on READY (VALID should assert first)
        # Simplified check: if READY without VALID, that's fine.
        # If VALID deasserts without handshake, that's a violation.
        for channel, vname, rname in [
            ("AW", "awvalid", "awready"),
            ("W", "wvalid", "wready"),
            ("B", "bvalid", "bready"),
        ]:
            vsig = self._get(vname)
            rsig = self._get(rname)
            if vsig is None or rsig is None:
                continue

            # Check: once VALID asserted, must stay high until handshake
            prev_valid = False
            for idx in range(t.max_index + 1):
                old = t._index
                t._index = idx
                t._virtual_cache.clear()
                v = bool(vsig.value)
                r = bool(rsig.value)
                t._index = old

                if prev_valid and not v and not r:
                    violations.append(AXI4Violation(
                        type=f"{channel}_VALID_DEASSERT_WITHOUT_HANDSHAKE",
                        cycle=idx,
                        message=f"{vname} deasserted at cycle {idx} without {rname}",
                    ))
                prev_valid = v

        return violations
