"""Tests for the deterministic replay engine."""

import unittest

from deltatrace.event_graph import Event, EventGraph, EventType
from deltatrace.replay_engine import ReplayEngine


class TestReplayEngine(unittest.TestCase):
    def _build_graph(self):
        graph = EventGraph()
        tick = Event.create(EventType.MD_TICK, {"symbol": "AAPL"})
        signal = Event.create(EventType.SIGNAL, {"direction": "buy"})
        order = Event.create(EventType.ORDER, {"qty": 100})
        fill = Event.create(EventType.FILL, {"price": 185.0})

        for ev in [tick, signal, order, fill]:
            graph.add_event(ev)
        graph.add_edge(tick.id, signal.id)
        graph.add_edge(signal.id, order.id)
        graph.add_edge(order.id, fill.id)
        return graph

    def test_replay_all_events(self):
        graph = self._build_graph()
        engine = ReplayEngine()
        result = engine.replay(graph)
        self.assertEqual(result.events_replayed, 4)
        self.assertEqual(result.fidelity_score, 1.0)
        self.assertEqual(len(result.event_log), 4)

    def test_replay_deterministic(self):
        graph = self._build_graph()
        engine = ReplayEngine()
        r1 = engine.replay(graph)
        r2 = engine.replay(graph)
        ids1 = [e["id"] for e in r1.event_log]
        ids2 = [e["id"] for e in r2.event_log]
        self.assertEqual(ids1, ids2)

    def test_replay_with_filter(self):
        graph = self._build_graph()
        engine = ReplayEngine()
        result = engine.replay(graph, filter_types={EventType.MD_TICK, EventType.FILL})
        self.assertEqual(result.events_replayed, 2)
        types = {e["type"] for e in result.event_log}
        self.assertEqual(types, {"md_tick", "fill"})

    def test_replay_handlers(self):
        graph = self._build_graph()
        engine = ReplayEngine()
        seen = []
        engine.register_handler(EventType.ORDER, lambda e: seen.append(e.id))
        engine.replay(graph)
        self.assertEqual(len(seen), 1)

    def test_replay_from_snapshot(self):
        graph = self._build_graph()
        snapshot = graph.to_dict()
        engine = ReplayEngine()
        result = engine.replay_from_snapshot(snapshot)
        self.assertEqual(result.events_replayed, 4)

    def test_replay_checkpoint_resume(self):
        graph = self._build_graph()
        engine = ReplayEngine()

        # Create a checkpoint up to the SIGNAL event and replay just that prefix
        signal = next(e for e in graph.topological_order() if e.event_type == EventType.SIGNAL)
        checkpoint = engine.create_checkpoint(graph, upto_event_id=signal.id)
        res = engine.replay_from_snapshot(checkpoint)
        # tick + signal expected
        self.assertEqual(res.events_replayed, 2)

    def test_determinism_score(self):
        graph = self._build_graph()
        engine = ReplayEngine()

        # Reference run
        r1 = engine.replay(graph)
        ref = r1.event_log

        # Mutate timestamps in a snapshot to force a different replay ordering
        snapshot = graph.to_dict()
        # Find the signal event and make its timestamp far later
        for e in snapshot["events"]:
            if e["type"] == "signal":
                e["timestamp_ns"] = e["timestamp_ns"] + 10_000_000_000
                break

        mutated = EventGraph.from_dict(snapshot)
        r2 = engine.replay(mutated, reference_event_log=ref)
        # Expect fidelity to be lower than perfect (not equal to 1.0)
        self.assertLess(r2.fidelity_score, 1.0)


if __name__ == "__main__":
    unittest.main()