"""
Tests for StructureMemory event firing (observability hooks).

Verifies acceptance criteria:
- or1-monitor.AC2.6: SM fires TokenReceived when a token is dequeued
- or1-monitor.AC2.7: SM fires CellWritten on any cell state change
- or1-monitor.AC2.8: SM fires DeferredRead when a read blocks
- or1-monitor.AC2.9: SM fires DeferredSatisfied when deferred read is satisfied
- or1-monitor.AC2.10: SM fires ResultSent when a result token is routed back
"""

import simpy

from cm_inst import MemOp
from emu.events import (
    TokenReceived, CellWritten, DeferredRead as DeferredReadEvent,
    DeferredSatisfied, ResultSent,
)
from emu.sm import StructureMemory
from sm_mod import Presence
from tokens import CMToken, SMToken


def inject_token(env: simpy.Environment, store: simpy.Store, token):
    """Helper to inject token into store via a process."""
    def _injector():
        yield store.put(token)

    env.process(_injector())


class TestAC2_6TokenReceived:
    """AC2.6: SM fires TokenReceived when a token is dequeued."""

    def test_token_received_write(self):
        """TokenReceived event fires for WRITE operation."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)

        # Inject WRITE token
        write_token = SMToken(target=0, addr=10, op=MemOp.WRITE, flags=None, data=0x1234, ret=None)
        inject_token(env, sm.input_store, write_token)

        env.run(until=100)

        # Verify TokenReceived event was fired
        token_received_events = [e for e in events if isinstance(e, TokenReceived)]
        assert len(token_received_events) >= 1
        assert token_received_events[0].token == write_token
        assert token_received_events[0].component == "sm:0"

    def test_token_received_read(self):
        """TokenReceived event fires for READ operation."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Pre-populate cell
        sm.cells[5].pres = Presence.FULL
        sm.cells[5].data_l = 0x5555

        # Inject READ token
        ret_route = CMToken(target=0, offset=0, act_id=0, data=0)
        read_token = SMToken(target=0, addr=5, op=MemOp.READ, flags=None, data=None, ret=ret_route)
        inject_token(env, sm.input_store, read_token)

        env.run(until=100)

        # Verify TokenReceived event was fired
        token_received_events = [e for e in events if isinstance(e, TokenReceived)]
        assert len(token_received_events) >= 1
        assert token_received_events[0].token == read_token


class TestAC2_7CellWritten:
    """AC2.7: SM fires CellWritten on any cell state change."""

    def test_cell_written_empty_to_full(self):
        """CellWritten fires when EMPTY cell becomes FULL."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)

        # Cell starts EMPTY
        assert sm.cells[20].pres == Presence.EMPTY

        # Inject WRITE to EMPTY cell
        write_token = SMToken(target=0, addr=20, op=MemOp.WRITE, flags=None, data=0xABCD, ret=None)
        inject_token(env, sm.input_store, write_token)

        env.run(until=100)

        # Verify CellWritten event was fired
        cell_written_events = [e for e in events if isinstance(e, CellWritten)]
        assert len(cell_written_events) >= 1
        cw = cell_written_events[0]
        assert cw.addr == 20
        assert cw.old_pres == Presence.EMPTY
        assert cw.new_pres == Presence.FULL

    def test_cell_written_clear(self):
        """CellWritten fires when cell is cleared to EMPTY."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)

        # Set cell to FULL
        sm.cells[30].pres = Presence.FULL
        sm.cells[30].data_l = 0x9999

        # Inject CLEAR token
        clear_token = SMToken(target=0, addr=30, op=MemOp.CLEAR, flags=None, data=None, ret=None)
        inject_token(env, sm.input_store, clear_token)

        env.run(until=100)

        # Verify CellWritten event was fired
        cell_written_events = [e for e in events if isinstance(e, CellWritten)]
        assert len(cell_written_events) >= 1
        cw = cell_written_events[0]
        assert cw.addr == 30
        assert cw.old_pres == Presence.FULL
        assert cw.new_pres == Presence.EMPTY

    def test_cell_written_alloc(self):
        """CellWritten fires when EMPTY cell is allocated to RESERVED."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)

        # Cell starts EMPTY
        assert sm.cells[40].pres == Presence.EMPTY

        # Inject ALLOC token
        alloc_token = SMToken(target=0, addr=40, op=MemOp.ALLOC, flags=None, data=None, ret=None)
        inject_token(env, sm.input_store, alloc_token)

        env.run(until=100)

        # Verify CellWritten event was fired
        cell_written_events = [e for e in events if isinstance(e, CellWritten)]
        assert len(cell_written_events) >= 1
        cw = cell_written_events[0]
        assert cw.addr == 40
        assert cw.old_pres == Presence.EMPTY
        assert cw.new_pres == Presence.RESERVED

    def test_cell_written_deferred_satisfaction(self):
        """CellWritten fires when deferred read is satisfied (WAITING -> FULL)."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Set up deferred read on cell 50
        ret_route = CMToken(target=0, offset=0, act_id=0, data=0)
        read_token = SMToken(target=0, addr=50, op=MemOp.READ, flags=None, data=None, ret=ret_route)
        inject_token(env, sm.input_store, read_token)

        env.run(until=10)

        # Clear events to focus on satisfaction
        events.clear()

        # Inject WRITE to satisfy deferred read
        write_token = SMToken(target=0, addr=50, op=MemOp.WRITE, flags=None, data=0xDEAD, ret=None)
        inject_token(env, sm.input_store, write_token)

        env.run(until=100)

        # Verify CellWritten event with WAITING -> FULL transition
        cell_written_events = [e for e in events if isinstance(e, CellWritten)]
        assert len(cell_written_events) >= 1
        cw = cell_written_events[0]
        assert cw.addr == 50
        assert cw.old_pres == Presence.WAITING
        assert cw.new_pres == Presence.FULL

    def test_cell_written_atomic_full_to_full(self):
        """CellWritten fires for atomic ops with FULL -> FULL transition."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Set cell to FULL
        sm.cells[100].pres = Presence.FULL
        sm.cells[100].data_l = 42

        # Inject RD_INC token
        ret_route = CMToken(target=0, offset=0, act_id=0, data=0)
        inc_token = SMToken(target=0, addr=100, op=MemOp.RD_INC, flags=None, data=None, ret=ret_route)
        inject_token(env, sm.input_store, inc_token)

        env.run(until=100)

        # Verify CellWritten event with FULL -> FULL transition
        cell_written_events = [e for e in events if isinstance(e, CellWritten)]
        assert len(cell_written_events) >= 1
        cw = cell_written_events[0]
        assert cw.addr == 100
        assert cw.old_pres == Presence.FULL
        assert cw.new_pres == Presence.FULL


class TestAC2_8DeferredRead:
    """AC2.8: SM fires DeferredRead when a read blocks on a non-FULL cell."""

    def test_deferred_read_on_empty(self):
        """DeferredRead event fires when READ blocks on EMPTY cell."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Cell starts EMPTY
        assert sm.cells[60].pres == Presence.EMPTY

        # Inject READ token
        ret_route = CMToken(target=0, offset=0, act_id=0, data=0)
        read_token = SMToken(target=0, addr=60, op=MemOp.READ, flags=None, data=None, ret=ret_route)
        inject_token(env, sm.input_store, read_token)

        env.run(until=100)

        # Verify DeferredRead event was fired
        deferred_read_events = [e for e in events if isinstance(e, DeferredReadEvent)]
        assert len(deferred_read_events) >= 1
        dr = deferred_read_events[0]
        assert dr.addr == 60
        assert dr.component == "sm:0"


class TestAC2_9DeferredSatisfied:
    """AC2.9: SM fires DeferredSatisfied when a subsequent write satisfies a deferred read."""

    def test_deferred_satisfied(self):
        """DeferredSatisfied event fires when deferred read is satisfied."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Set up deferred read on cell 70
        ret_route = CMToken(target=0, offset=0, act_id=0, data=0)
        read_token = SMToken(target=0, addr=70, op=MemOp.READ, flags=None, data=None, ret=ret_route)
        inject_token(env, sm.input_store, read_token)

        env.run(until=10)

        # Clear events to focus on satisfaction
        events.clear()

        # Inject WRITE to satisfy deferred read
        write_token = SMToken(target=0, addr=70, op=MemOp.WRITE, flags=None, data=0x7777, ret=None)
        inject_token(env, sm.input_store, write_token)

        env.run(until=100)

        # Verify DeferredSatisfied event was fired
        deferred_satisfied_events = [e for e in events if isinstance(e, DeferredSatisfied)]
        assert len(deferred_satisfied_events) >= 1
        ds = deferred_satisfied_events[0]
        assert ds.addr == 70
        assert ds.data == 0x7777
        assert ds.component == "sm:0"


class TestAC2_10ResultSent:
    """AC2.10: SM fires ResultSent when a result token is routed back to a PE."""

    def test_result_sent_on_read_full(self):
        """ResultSent event fires when READ on FULL cell returns result."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Pre-populate cell with data
        sm.cells[80].pres = Presence.FULL
        sm.cells[80].data_l = 0xCAFE

        # Inject READ token
        ret_route = CMToken(target=0, offset=5, act_id=2, data=0)
        read_token = SMToken(target=0, addr=80, op=MemOp.READ, flags=None, data=None, ret=ret_route)
        inject_token(env, sm.input_store, read_token)

        env.run(until=100)

        # Verify ResultSent event was fired
        result_sent_events = [e for e in events if isinstance(e, ResultSent)]
        assert len(result_sent_events) >= 1
        rs = result_sent_events[0]
        assert rs.token.data == 0xCAFE
        assert rs.token.offset == 5
        assert rs.token.act_id == 2
        assert rs.component == "sm:0"

    def test_result_sent_on_atomic(self):
        """ResultSent event fires when atomic op returns old value."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Pre-populate cell
        sm.cells[110].pres = Presence.FULL
        sm.cells[110].data_l = 100

        # Inject RD_INC token
        ret_route = CMToken(target=0, offset=12, act_id=1, data=0)
        inc_token = SMToken(target=0, addr=110, op=MemOp.RD_INC, flags=None, data=None, ret=ret_route)
        inject_token(env, sm.input_store, inc_token)

        env.run(until=100)

        # Verify ResultSent event was fired with old value
        result_sent_events = [e for e in events if isinstance(e, ResultSent)]
        assert len(result_sent_events) >= 1
        rs = result_sent_events[0]
        assert rs.token.data == 100  # Old value
        assert rs.component == "sm:0"

    def test_result_sent_on_deferred_satisfaction(self):
        """ResultSent event fires when deferred read is satisfied."""
        env = simpy.Environment()
        events = []

        def on_event(event):
            events.append(event)

        sm = StructureMemory(env, 0, on_event=on_event)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Set up deferred read on cell 120
        ret_route = CMToken(target=0, offset=15, act_id=3, data=0)
        read_token = SMToken(target=0, addr=120, op=MemOp.READ, flags=None, data=None, ret=ret_route)
        inject_token(env, sm.input_store, read_token)

        env.run(until=10)

        # Clear events
        events.clear()

        # Inject WRITE to satisfy deferred read
        write_token = SMToken(target=0, addr=120, op=MemOp.WRITE, flags=None, data=0x8888, ret=None)
        inject_token(env, sm.input_store, write_token)

        env.run(until=100)

        # Verify ResultSent event was fired with written data
        result_sent_events = [e for e in events if isinstance(e, ResultSent)]
        assert len(result_sent_events) >= 1
        rs = result_sent_events[0]
        assert rs.token.data == 0x8888
        assert rs.token.offset == 15
        assert rs.token.act_id == 3