"""Base data models for agent execution trace events.

This module defines the core event types and base TraceEvent class used to capture
and analyze agent execution traces. Events form a tree structure where each
event can have a parent, enabling hierarchical trace analysis.
"""

from __future__ import annotations

import sys
import uuid
from dataclasses import asdict, dataclass, field, fields
from datetime import datetime, timezone
from enum import Enum
from typing import Any

# Python 3.10 compatibility: StrEnum was added in Python 3.11
if sys.version_info >= (3, 11):
    from enum import StrEnum
else:

    class StrEnum(str, Enum):
        """Compatibility shim for StrEnum in Python 3.10."""

        def __str__(self) -> str:
            return str(self.value)


class EventType(StrEnum):
    """Enumeration of all trace event types.

    These represent the different phases and actions during agent execution:
    - AGENT_START/END: Session lifecycle events
    - TOOL_CALL/RESULT: Tool execution events
    - LLM_REQUEST/RESPONSE: LLM API interaction events
    - DECISION: Agent decision points with reasoning
    - ERROR: Error occurrences during execution
    - CHECKPOINT: State snapshots for time-travel debugging
    """

    AGENT_START = "agent_start"
    AGENT_END = "agent_end"
    TOOL_CALL = "tool_call"
    TOOL_RESULT = "tool_result"
    LLM_REQUEST = "llm_request"
    LLM_RESPONSE = "llm_response"
    DECISION = "decision"
    ERROR = "error"
    CHECKPOINT = "checkpoint"
    SAFETY_CHECK = "safety_check"
    REFUSAL = "refusal"
    POLICY_VIOLATION = "policy_violation"
    PROMPT_POLICY = "prompt_policy"
    AGENT_TURN = "agent_turn"
    BEHAVIOR_ALERT = "behavior_alert"
    REPAIR_ATTEMPT = "repair_attempt"


class SessionStatus(StrEnum):
    """Session lifecycle status values."""

    RUNNING = "running"
    COMPLETED = "completed"
    ERROR = "error"


class RiskLevel(StrEnum):
    """Shared risk/severity labels across domain events."""

    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class SafetyOutcome(StrEnum):
    """Explicit outcome labels for safety checks."""

    PASS = "pass"
    FAIL = "fail"
    WARN = "warn"
    BLOCK = "block"


BASE_EVENT_FIELDS = {
    "id",
    "session_id",
    "parent_id",
    "event_type",
    "timestamp",
    "name",
    "data",
    "metadata",
    "importance",
    "upstream_event_ids",
}


def _serialize_field_value(value: Any) -> Any:
    """Convert dataclass field values into JSON-serializable payloads."""
    if isinstance(value, datetime):
        return value.isoformat()
    if isinstance(value, Enum):
        return value.value
    if isinstance(value, list):
        return [_serialize_field_value(item) for item in value]
    if isinstance(value, tuple):
        return [_serialize_field_value(item) for item in value]
    if isinstance(value, dict):
        return {key: _serialize_field_value(item) for key, item in value.items()}
    return value


@dataclass(kw_only=True)
class TraceEvent:
    """Base dataclass for all trace events.

    Events form a tree structure via parent_id.
    """

    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    session_id: str = ""
    parent_id: str | None = None
    event_type: EventType = EventType.AGENT_START
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    name: str = ""
    data: dict[str, Any] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)
    importance: float = 0.5
    upstream_event_ids: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a JSON-serializable dictionary."""
        # Use asdict (C-implemented) for better performance
        raw_dict = asdict(self)
        # Post-process for datetime/Enum serialization
        return {k: _serialize_field_value(v) for k, v in raw_dict.items()}

    @classmethod
    def _typed_field_names(cls) -> set[str]:
        """Return event-specific dataclass fields beyond the shared base payload."""
        return {field_info.name for field_info in fields(cls) if field_info.name not in BASE_EVENT_FIELDS}

    def to_storage_data(self) -> dict[str, Any]:
        """Merge event-specific fields into the storage payload."""
        payload = dict(self.data)
        for field_name in self._typed_field_names():
            payload[field_name] = getattr(self, field_name)
        return payload

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> TraceEvent:
        """Deserialize a dictionary to a TraceEvent."""
        # Handle timestamp deserialization
        if isinstance(data.get("timestamp"), str):
            data["timestamp"] = datetime.fromisoformat(data["timestamp"])

        # Handle EventType deserialization
        if isinstance(data.get("event_type"), str):
            data["event_type"] = EventType(data["event_type"])

        return cls(**data)

    @classmethod
    def from_data(
        cls,
        event_type: EventType,
        base_kwargs: dict[str, Any],
        data: dict[str, Any],
    ) -> TraceEvent:
        """Build the typed event instance for the given event_type."""
        from agent_debugger_sdk.core.events import EVENT_TYPE_REGISTRY

        # Use try/except to trigger lazy loading via __missing__
        try:
            event_cls = EVENT_TYPE_REGISTRY[event_type]
        except KeyError:
            # Event type not in registry, use base TraceEvent class
            event_cls = cls

        typed_field_names = event_cls._typed_field_names()
        typed_kwargs = {field_name: data[field_name] for field_name in typed_field_names if field_name in data}
        payload = {key: value for key, value in data.items() if key not in typed_field_names}
        return event_cls(
            **base_kwargs,
            event_type=event_type,
            data=payload,
            **typed_kwargs,
        )
