"""State snapshot capture from a live simulation System.

Provides a frozen dataclass representation of PE and SM state at a moment in time,
suitable for serialization and network transport.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

from cm_inst import Instruction, FrameSlotValue, Port
from emu.network import System
from sm_mod import Presence
from tokens import Token


@dataclass(frozen=True)
class PESnapshot:
    """Frozen snapshot of a Processing Element's state at a moment in time."""

    pe_id: int
    iram: dict[int, Instruction]
    frames: tuple[tuple[FrameSlotValue, ...], ...]
    tag_store: dict[int, tuple[int, int]]
    presence: tuple[tuple[tuple[bool, ...], ...], ...]
    port_store: tuple[tuple[tuple[Port | None, ...], ...], ...]
    match_data: tuple[tuple[tuple[int | None, ...], ...], ...]
    free_frames: tuple[int, ...]
    lane_count: int
    input_queue: tuple[Token, ...]
    output_log: tuple[Token, ...]


@dataclass(frozen=True)
class SMCellSnapshot:
    """Frozen snapshot of a Structure Memory cell's state."""

    pres: Presence
    data_l: Optional[int]
    data_r: Optional[int]


@dataclass(frozen=True)
class SMSnapshot:
    """Frozen snapshot of a Structure Memory's state at a moment in time."""

    sm_id: int
    cells: dict[int, SMCellSnapshot]
    deferred_read: Optional[dict]
    t0_store: tuple[int, ...]
    input_queue: tuple[Token, ...]


@dataclass(frozen=True)
class StateSnapshot:
    """Complete state snapshot of all PEs and SMs at a moment in time."""

    sim_time: float
    next_time: float
    pes: dict[int, PESnapshot]
    sms: dict[int, SMSnapshot]


def capture(system: System) -> StateSnapshot:
    """Capture a frozen snapshot of the entire system state at the current moment.

    Reads from live PE and SM attributes (frames, tag_store, presence, port_store,
    free_frames, iram, input_store.items, output_log, cells, deferred_read, t0_store).

    Args:
        system: A System instance from emu.network.build_topology()

    Returns:
        StateSnapshot containing frozen copies of all PE and SM state.
    """
    pes = {}
    for pe_id, pe in system.pes.items():
        frames = tuple(
            tuple(slot for slot in frame)
            for frame in pe.frames
        )
        tag_store = dict(pe.tag_store)
        presence = tuple(
            tuple(
                tuple(lane_val for lane_val in offset_lanes)
                for offset_lanes in frame_presence
            )
            for frame_presence in pe.presence
        )
        port_store = tuple(
            tuple(
                tuple(lane_val for lane_val in offset_lanes)
                for offset_lanes in frame_ports
            )
            for frame_ports in pe.port_store
        )
        match_data = tuple(
            tuple(
                tuple(lane_val for lane_val in offset_lanes)
                for offset_lanes in frame_match
            )
            for frame_match in pe.match_data
        )
        free_frames = tuple(pe.free_frames)

        pes[pe_id] = PESnapshot(
            pe_id=pe_id,
            iram=dict(pe.iram),
            frames=frames,
            tag_store=tag_store,
            presence=presence,
            port_store=port_store,
            match_data=match_data,
            free_frames=free_frames,
            lane_count=pe.lane_count,
            input_queue=tuple(pe.input_store.items),
            output_log=tuple(pe.output_log),
        )

    sms = {}
    for sm_id, sm in system.sms.items():
        cells = {}
        for i, cell in enumerate(sm.cells):
            if cell.pres != Presence.EMPTY or cell.data_l is not None:
                cells[i] = SMCellSnapshot(
                    pres=cell.pres,
                    data_l=cell.data_l,
                    data_r=cell.data_r,
                )

        dr = None
        if sm.deferred_read is not None:
            dr = {
                "cell_addr": sm.deferred_read.cell_addr,
                "return_route": sm.deferred_read.return_route,
            }

        sms[sm_id] = SMSnapshot(
            sm_id=sm_id,
            cells=cells,
            deferred_read=dr,
            t0_store=tuple(sm.t0_store),
            input_queue=tuple(sm.input_store.items),
        )

    return StateSnapshot(
        sim_time=system.env.now,
        next_time=system.env.peek(),
        pes=pes,
        sms=sms,
    )