# proxy to match cocotb and custom axi4_if

import logging

from cocotb.types import LogicArray

from cocotbext.axi.axi_channels import (
    AxiAWBus,
    AxiWBus,
    AxiBBus,
    AxiARBus,
    AxiRBus,
    AxiBus,
)


def install_cocotbext_axi_slice_compat():
    """
    cocotbext-axi StreamMonitor uses RisingEdge(self.valid/self.ready).
    That works only for real cocotb LogicObject handles.

    Our SliceSignal is a Python proxy over packed req/resp vector, so
    RisingEdge(SliceSignal) fails in cocotb 2.x.

    This patch makes StreamMonitor wake on signal.value_change instead.
    For SliceSignal this should return parent packed vector value_change.
    """

    # this sucks...
    from cocotbext.axi import stream as axi_stream

    async def _run_valid_monitor_value_change(self):
        while True:
            await self.valid.value_change
            self.wake_event.set()

    async def _run_ready_monitor_value_change(self):
        while True:
            await self.ready.value_change
            self.wake_event.set()

    axi_stream.StreamMonitor._run_valid_monitor = _run_valid_monitor_value_change
    axi_stream.StreamMonitor._run_ready_monitor = _run_ready_monitor_value_change


def _safe_int(value) -> int:
    try:
        return int(value)
    except Exception:

        s = str(value).upper()
        bits = []

        for ch in s:
            if ch in "01":
                bits.append(ch)
            elif ch in "XZUW-":
                bits.append("0")

        if not bits:
            return 0

        return int("".join(bits), 2)


def _logic_array(value: int, width: int):
    value &= (1 << width) - 1
    return LogicArray.from_unsigned(value, width)


def _param_int(obj, name: str, override):
    if override is not None:
        return int(override)

    p = getattr(obj, name)

    try:
        return int(p)
    except Exception:
        pass

    try:
        return int(p.value)
    except Exception:
        pass

    raise RuntimeError(f"Cannot read parameter {name} from {obj!r}")


def _packed_offsets(fields):
    """
    de-construct SV signal order.
    """
    total = sum(width for _, width in fields)
    pos = total
    offsets = {}

    for name, width in fields:
        pos -= width
        offsets[name] = (pos, width)

    return total, offsets


class PackedVector:
    def __init__(self, handle, name=None):
        self.handle = handle
        self.name = name or getattr(handle, "_path", repr(handle))
        self.width = len(handle)
        self.mask = (1 << self.width) - 1

        self.shadow = _safe_int(handle.value) & self.mask
        self.dirty = False

    def read(self, lo: int, width: int) -> int:
        src = self.shadow if self.dirty else (
            _safe_int(self.handle.value) & self.mask)
        return (src >> lo) & ((1 << width) - 1)

    def write(self, lo: int, width: int, value, immediate=False):
        value = _safe_int(value)
        field_mask = ((1 << width) - 1) << lo

        self.shadow &= ~field_mask
        self.shadow |= (value << lo) & field_mask
        self.shadow &= self.mask
        self.dirty = True

        v = _logic_array(self.shadow, self.width)

        if immediate:
            self.handle.setimmediatevalue(v)
        else:
            self.handle.value = v


class SliceSignal:
    """
    cocotbext-axi proxy

    looks like:
        sig.value
        sig.value = ...
        sig.setimmediatevalue(...)
        len(sig)

    """

    def __init__(self, name: str, parent: PackedVector, lo: int, width: int):
        self._name = name
        self._path = name
        self._log = logging.getLogger(f"cocotb.{name}")

        self.parent = parent
        self.lo = lo
        self.width = width

    def __len__(self):
        return self.width

    def __repr__(self):
        return f"<SliceSignal {self._path}[{self.lo + self.width - 1}:{self.lo}]>"

    @property
    def value(self):
        return _logic_array(self.parent.read(self.lo, self.width), self.width)

    @value.setter
    def value(self, value):
        self.parent.write(self.lo, self.width, value, immediate=False)

    def setimmediatevalue(self, value):
        self.parent.write(self.lo, self.width, value, immediate=True)

    def set(self, value):
        self.value = value

    def get(self):
        return self.value

    @property
    def value_change(self):
        # dirty
        return self.parent.handle.value_change


class ProxyEntity:
    def __init__(self, name: str, **signals):
        self._name = name
        self._path = name
        self._log = logging.getLogger(f"cocotb.{name}")
        self.__dict__.update(signals)

    def __dir__(self):
        return list(self.__dict__.keys())


def _sig(name, parent, base_lo, layout, field):
    lo, width = layout[field]
    return SliceSignal(name, parent, base_lo + lo, width)


def axi4_bus_from_packed_if(
    iface,
    name="axi",
    *,
    addr_width=None,
    data_width=None,
    id_width=None,
    user_width=None,
):
    """
    Build cocotbext-axi AxiBus from custom packed axi4_if.

    iface should be instance axi4_if:
        dut.axi_in
        dut.axi_out

    example:
        master = AxiMaster(
            axi4_bus_from_packed_if(dut.axi_in, "s_axi"),
            dut.clk,
            dut.rst,
        )

        ram = AxiRam(
            axi4_bus_from_packed_if(dut.axi_out, "m_axi"),
            dut.clk,
            dut.rst,
            size=2**20,
        )
    """

    ADDR_W = _param_int(iface, "ADDR_W", addr_width)
    DATA_W = _param_int(iface, "DATA_W", data_width)
    ID_W = _param_int(iface, "ID_W", id_width)
    USER_W = _param_int(iface, "USER_W", user_width)

    if DATA_W % 8 != 0:
        raise ValueError(f"AXI DATA_W must be divisible by 8, got {DATA_W}")

    STRB_W = DATA_W // 8

    aw_total, aw_layout = _packed_offsets([
        ("id", ID_W),
        ("addr", ADDR_W),
        ("len", 8),
        ("size", 3),
        ("burst", 2),
        ("lock", 1),
        ("cache", 4),
        ("prot", 3),
        ("qos", 4),
        ("region", 4),
        ("user", USER_W),
        ("valid", 1),
    ])

    w_total, w_layout = _packed_offsets([
        ("data", DATA_W),
        ("strb", STRB_W),
        ("last", 1),
        ("user", USER_W),
        ("valid", 1),
    ])

    b_total, b_layout = _packed_offsets([
        ("id", ID_W),
        ("resp", 2),
        ("user", USER_W),
        ("valid", 1),
    ])

    ar_total, ar_layout = _packed_offsets([
        ("id", ID_W),
        ("addr", ADDR_W),
        ("len", 8),
        ("size", 3),
        ("burst", 2),
        ("lock", 1),
        ("cache", 4),
        ("prot", 3),
        ("qos", 4),
        ("region", 4),
        ("user", USER_W),
        ("valid", 1),
    ])

    r_total, r_layout = _packed_offsets([
        ("id", ID_W),
        ("data", DATA_W),
        ("resp", 2),
        ("last", 1),
        ("user", USER_W),
        ("valid", 1),
    ])

    req_total, req_layout = _packed_offsets([
        ("aw", aw_total),
        ("w", w_total),
        ("b_ready", 1),
        ("ar", ar_total),
        ("r_ready", 1),
    ])

    resp_total, resp_layout = _packed_offsets([
        ("aw_ready", 1),
        ("w_ready", 1),
        ("b", b_total),
        ("ar_ready", 1),
        ("r", r_total),
    ])

    req = PackedVector(iface.req, f"{name}.req")
    resp = PackedVector(iface.resp, f"{name}.resp")

    if req_total != req.width:
        raise RuntimeError(
            f"{name}.req layout mismatch: calculated {req_total} bits, "
            f"simulator has {req.width} bits"
        )

    if resp_total != resp.width:
        raise RuntimeError(
            f"{name}.resp layout mismatch: calculated {resp_total} bits, "
            f"simulator has {resp.width} bits"
        )

    aw_lo, _ = req_layout["aw"]
    w_lo, _ = req_layout["w"]
    ar_lo, _ = req_layout["ar"]

    b_lo, _ = resp_layout["b"]
    r_lo, _ = resp_layout["r"]

    aw = AxiAWBus.from_entity(ProxyEntity(
        f"{name}_aw",

        awid=_sig(f"{name}_awid", req, aw_lo, aw_layout, "id"),
        awaddr=_sig(f"{name}_awaddr", req, aw_lo, aw_layout, "addr"),
        awlen=_sig(f"{name}_awlen", req, aw_lo, aw_layout, "len"),
        awsize=_sig(f"{name}_awsize", req, aw_lo, aw_layout, "size"),
        awburst=_sig(f"{name}_awburst", req, aw_lo, aw_layout, "burst"),
        awlock=_sig(f"{name}_awlock", req, aw_lo, aw_layout, "lock"),
        awcache=_sig(f"{name}_awcache", req, aw_lo, aw_layout, "cache"),
        awprot=_sig(f"{name}_awprot", req, aw_lo, aw_layout, "prot"),
        awqos=_sig(f"{name}_awqos", req, aw_lo, aw_layout, "qos"),
        awregion=_sig(f"{name}_awregion", req, aw_lo, aw_layout, "region"),
        awuser=_sig(f"{name}_awuser", req, aw_lo, aw_layout, "user"),
        awvalid=_sig(f"{name}_awvalid", req, aw_lo, aw_layout, "valid"),

        awready=SliceSignal(
            f"{name}_awready",
            resp,
            *resp_layout["aw_ready"],
        ),
    ))

    w = AxiWBus.from_entity(ProxyEntity(
        f"{name}_w",

        wdata=_sig(f"{name}_wdata", req, w_lo, w_layout, "data"),
        wstrb=_sig(f"{name}_wstrb", req, w_lo, w_layout, "strb"),
        wlast=_sig(f"{name}_wlast", req, w_lo, w_layout, "last"),
        wuser=_sig(f"{name}_wuser", req, w_lo, w_layout, "user"),
        wvalid=_sig(f"{name}_wvalid", req, w_lo, w_layout, "valid"),

        wready=SliceSignal(
            f"{name}_wready",
            resp,
            *resp_layout["w_ready"],
        ),
    ))

    b = AxiBBus.from_entity(ProxyEntity(
        f"{name}_b",

        bid=_sig(f"{name}_bid", resp, b_lo, b_layout, "id"),
        bresp=_sig(f"{name}_bresp", resp, b_lo, b_layout, "resp"),
        buser=_sig(f"{name}_buser", resp, b_lo, b_layout, "user"),
        bvalid=_sig(f"{name}_bvalid", resp, b_lo, b_layout, "valid"),

        bready=SliceSignal(
            f"{name}_bready",
            req,
            *req_layout["b_ready"],
        ),
    ))

    ar = AxiARBus.from_entity(ProxyEntity(
        f"{name}_ar",

        arid=_sig(f"{name}_arid", req, ar_lo, ar_layout, "id"),
        araddr=_sig(f"{name}_araddr", req, ar_lo, ar_layout, "addr"),
        arlen=_sig(f"{name}_arlen", req, ar_lo, ar_layout, "len"),
        arsize=_sig(f"{name}_arsize", req, ar_lo, ar_layout, "size"),
        arburst=_sig(f"{name}_arburst", req, ar_lo, ar_layout, "burst"),
        arlock=_sig(f"{name}_arlock", req, ar_lo, ar_layout, "lock"),
        arcache=_sig(f"{name}_arcache", req, ar_lo, ar_layout, "cache"),
        arprot=_sig(f"{name}_arprot", req, ar_lo, ar_layout, "prot"),
        arqos=_sig(f"{name}_arqos", req, ar_lo, ar_layout, "qos"),
        arregion=_sig(f"{name}_arregion", req, ar_lo, ar_layout, "region"),
        aruser=_sig(f"{name}_aruser", req, ar_lo, ar_layout, "user"),
        arvalid=_sig(f"{name}_arvalid", req, ar_lo, ar_layout, "valid"),

        arready=SliceSignal(
            f"{name}_arready",
            resp,
            *resp_layout["ar_ready"],
        ),
    ))

    r = AxiRBus.from_entity(ProxyEntity(
        f"{name}_r",

        rid=_sig(f"{name}_rid", resp, r_lo, r_layout, "id"),
        rdata=_sig(f"{name}_rdata", resp, r_lo, r_layout, "data"),
        rresp=_sig(f"{name}_rresp", resp, r_lo, r_layout, "resp"),
        rlast=_sig(f"{name}_rlast", resp, r_lo, r_layout, "last"),
        ruser=_sig(f"{name}_ruser", resp, r_lo, r_layout, "user"),
        rvalid=_sig(f"{name}_rvalid", resp, r_lo, r_layout, "valid"),

        rready=SliceSignal(
            f"{name}_rready",
            req,
            *req_layout["r_ready"],
        ),
    ))

    return AxiBus.from_channels(aw, w, b, ar, r)