"""Deterministic replay engine for event graphs.

Replays a captured event stream in exact causal order, reproducing
decision paths for incident analysis and strategy validation.
"""

from __future__ import annotations

import heapq
import json
from dataclasses import dataclass
from typing import Callable, Optional

from .event_graph import Event, EventGraph, EventType


@dataclass
class ReplayResult:
    """Result of a deterministic replay run."""

    events_replayed: int
    total_latency_ns: int
    fidelity_score: float  # 1.0 = perfect determinism
    event_log: list[dict]


class ReplayEngine:
    """Replays event graphs deterministically using a priority queue."""

    def __init__(self) -> None:
        self._handlers: dict[EventType, list[Callable[[Event], None]]] = {}

    def register_handler(self, event_type: EventType, handler: Callable[[Event], None]) -> None:
        self._handlers.setdefault(event_type, []).append(handler)

    def replay(self, graph: EventGraph, filter_types: Optional[set[EventType]] = None) -> ReplayResult:
        """Replay events in deterministic topological + timestamp order.

        Optional behavior:
        - If `filter_types` is provided, only events of those types are replayed.
        - If `reference_event_log` is provided (list of event dicts), compute a
          fidelity score in [0.0, 1.0] comparing this run to the reference.
        """
        ordered = graph.topological_order()

        if filter_types:
            ordered = [e for e in ordered if e.event_type in filter_types]

        # Secondary sort by timestamp for determinism within same topo level
        ordered.sort(key=lambda e: e.timestamp_ns)

        event_log = []
        replayed = 0
        min_ts = ordered[0].timestamp_ns if ordered else 0
        max_ts = ordered[-1].timestamp_ns if ordered else 0

        for event in ordered:
            # Fire registered handlers
            for handler in self._handlers.get(event.event_type, []):
                handler(event)

            event_log.append({
                "id": event.id,
                "type": event.event_type.value,
                "timestamp_ns": event.timestamp_ns,
                "payload": event.payload,
            })
            replayed += 1

        return ReplayResult(
            events_replayed=replayed,
            total_latency_ns=max_ts - min_ts,
            fidelity_score=1.0,  # Deterministic replay = perfect fidelity
            event_log=event_log,
        )

    def replay(self, graph: EventGraph, filter_types: Optional[set[EventType]] = None,
               reference_event_log: Optional[list[dict]] = None) -> ReplayResult:
        """Replay events and optionally score fidelity against a reference log.

        Backwards-compatible: if called without `reference_event_log`, behavior
        is unchanged and fidelity_score is 1.0.
        """
        ordered = graph.topological_order()

        if filter_types:
            ordered = [e for e in ordered if e.event_type in filter_types]

        # Secondary sort by timestamp for determinism within same topo level
        ordered.sort(key=lambda e: e.timestamp_ns)

        event_log = []
        replayed = 0
        min_ts = ordered[0].timestamp_ns if ordered else 0
        max_ts = ordered[-1].timestamp_ns if ordered else 0

        for event in ordered:
            # Fire registered handlers
            for handler in self._handlers.get(event.event_type, []):
                handler(event)

            event_log.append({
                "id": event.id,
                "type": event.event_type.value,
                "timestamp_ns": event.timestamp_ns,
                "payload": event.payload,
            })
            replayed += 1

        # Compute a simple fidelity score vs a provided reference log: fraction of
        # matching ids at the same indices (normalized by the longer of the two).
        fidelity = 1.0
        if reference_event_log is not None:
            ref_ids = [e.get("id") for e in reference_event_log]
            cur_ids = [e.get("id") for e in event_log]
            matches = sum(1 for i in range(min(len(ref_ids), len(cur_ids))) if ref_ids[i] == cur_ids[i])
            denom = max(len(ref_ids), len(cur_ids)) if max(len(ref_ids), len(cur_ids)) > 0 else 1
            fidelity = matches / denom

        return ReplayResult(
            events_replayed=replayed,
            total_latency_ns=max_ts - min_ts,
            fidelity_score=fidelity,
            event_log=event_log,
        )

    def replay_from_snapshot(self, snapshot: dict, reference_event_log: Optional[list[dict]] = None) -> ReplayResult:
        """Replay from a serialized event graph snapshot."""
        graph = EventGraph.from_dict(snapshot)
        return self.replay(graph, reference_event_log=reference_event_log)

    def create_checkpoint(self, graph: EventGraph, upto_event_id: Optional[str] = None) -> dict:
        """Create a serialized checkpoint of the event graph.

        If `upto_event_id` is provided, include events in topological order up to
        and including that event. Otherwise returns full snapshot (graph.to_dict()).
        """
        if upto_event_id is None:
            return graph.to_dict()

        ordered = graph.topological_order()
        included_ids = []
        for e in ordered:
            included_ids.append(e.id)
            if e.id == upto_event_id:
                break

        # Build snapshot consisting of included events and edges between them
        events = [
            {
                "id": e.id,
                "type": e.event_type.value,
                "timestamp_ns": e.timestamp_ns,
                "payload": e.payload,
                "source_adapter": e.source_adapter,
            }
            for e in ordered if e.id in included_ids
        ]

        edges = [
            {"from": edge.from_id, "to": edge.to_id, "latency_ns": edge.latency_ns, "label": edge.label}
            for edge in graph.edges
            if edge.from_id in included_ids and edge.to_id in included_ids
        ]

        return {"events": events, "edges": edges}