"""Event graph model for market execution tracing.

Nodes represent market events (ticks, signals, orders, fills, risk checks).
Edges encode causal relationships with precise timestamps and latency budgets.
"""

from __future__ import annotations

import enum
import time
import uuid
from dataclasses import dataclass, field
from typing import Optional


class EventType(enum.Enum):
    MD_TICK = "md_tick"
    SIGNAL = "signal"
    PLAN_DELTA = "plan_delta"
    ORDER = "order"
    FILL = "fill"
    RISK_CHECK = "risk_check"


@dataclass(frozen=True)
class Event:
    """A node in the event graph."""

    id: str
    event_type: EventType
    timestamp_ns: int
    payload: dict
    source_adapter: str = ""

    @staticmethod
    def create(event_type: EventType, payload: dict, source_adapter: str = "") -> Event:
        return Event(
            id=uuid.uuid4().hex,
            event_type=event_type,
            timestamp_ns=time.time_ns(),
            payload=payload,
            source_adapter=source_adapter,
        )


@dataclass(frozen=True)
class CausalEdge:
    """Directed edge encoding causality and latency between events."""

    from_id: str
    to_id: str
    latency_ns: int
    label: str = ""


class EventGraph:
    """Directed acyclic graph of market events with causal edges."""

    def __init__(self) -> None:
        self._events: dict[str, Event] = {}
        self._edges: list[CausalEdge] = []
        self._children: dict[str, list[str]] = {}
        self._parents: dict[str, list[str]] = {}

    def add_event(self, event: Event) -> None:
        self._events[event.id] = event
        self._children.setdefault(event.id, [])
        self._parents.setdefault(event.id, [])

    def add_edge(self, from_id: str, to_id: str, label: str = "") -> CausalEdge:
        if from_id not in self._events or to_id not in self._events:
            raise KeyError("Both events must exist in the graph")
        latency = self._events[to_id].timestamp_ns - self._events[from_id].timestamp_ns
        edge = CausalEdge(from_id=from_id, to_id=to_id, latency_ns=latency, label=label)
        self._edges.append(edge)
        self._children[from_id].append(to_id)
        self._parents[to_id].append(from_id)
        return edge

    def get_event(self, event_id: str) -> Event:
        return self._events[event_id]

    def get_roots(self) -> list[Event]:
        return [e for e in self._events.values() if not self._parents.get(e.id)]

    def get_children(self, event_id: str) -> list[Event]:
        return [self._events[cid] for cid in self._children.get(event_id, [])]

    def get_parents(self, event_id: str) -> list[Event]:
        return [self._events[pid] for pid in self._parents.get(event_id, [])]

    def topological_order(self) -> list[Event]:
        """Return events in topological (causal) order."""
        visited: set[str] = set()
        order: list[str] = []

        def dfs(eid: str) -> None:
            if eid in visited:
                return
            visited.add(eid)
            for child_id in self._children.get(eid, []):
                dfs(child_id)
            order.append(eid)

        for eid in self._events:
            dfs(eid)

        order.reverse()
        return [self._events[eid] for eid in order]

    @property
    def events(self) -> list[Event]:
        return list(self._events.values())

    @property
    def edges(self) -> list[CausalEdge]:
        return list(self._edges)

    def to_dict(self) -> dict:
        return {
            "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 self._events.values()
            ],
            "edges": [
                {
                    "from": edge.from_id,
                    "to": edge.to_id,
                    "latency_ns": edge.latency_ns,
                    "label": edge.label,
                }
                for edge in self._edges
            ],
        }

    @classmethod
    def from_dict(cls, data: dict) -> EventGraph:
        graph = cls()
        for e in data["events"]:
            event = Event(
                id=e["id"],
                event_type=EventType(e["type"]),
                timestamp_ns=e["timestamp_ns"],
                payload=e["payload"],
                source_adapter=e.get("source_adapter", ""),
            )
            graph.add_event(event)
        for edge in data["edges"]:
            graph.add_edge(edge["from"], edge["to"], edge.get("label", ""))
        return graph