"""
Data models for trace objects using dataclasses for type safety.
"""
from dataclasses import dataclass
from typing import Any, Dict, Literal, Optional


@dataclass
class TokenUsage:
    """Token usage statistics for a trace."""

    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    estimated_cost: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {
            "promptTokens": self.prompt_tokens,
            "completionTokens": self.completion_tokens,
            "totalTokens": self.total_tokens,
            "estimatedCost": self.estimated_cost,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TokenUsage":
        """Build TokenUsage from dictionary payload."""
        return cls(
            prompt_tokens=int(data.get("promptTokens", 0) or 0),
            completion_tokens=int(data.get("completionTokens", 0) or 0),
            total_tokens=int(data.get("totalTokens", 0) or 0),
            estimated_cost=float(data.get("estimatedCost", 0.0) or 0.0),
        )


@dataclass
class Trace:
    """
    Trace object containing all trace data.
    Supports flexible serialization for different API endpoints.
    """

    id: str
    name: str
    result: str
    begin: Optional[float]
    duration: Optional[float]
    tokens: TokenUsage
    source: Literal["dataset", "pasted"] = "dataset"
    parent_node: Optional[Dict[str, Any]] = None  # Full trace tree with assigned IDs

    def to_list_dict(self) -> Dict[str, Any]:
        """
        Convert to dictionary for list endpoint (lightweight, no trace tree).

        Returns:
            Dictionary with metadata fields only (id, name, result, etc.)
        """
        return {
            "id": self.id,
            "name": self.name,
            "result": self.result,
            "begin": self.begin,
            "duration": self.duration,
            "tokens": self.tokens.to_dict(),
            "source": self.source,
        }

    def to_detail_dict(self) -> Dict[str, Any]:
        """
        Convert to dictionary for detail endpoints (includes trace tree).

        Returns:
            Dictionary with all fields including parentNode
        """
        return {
            "id": self.id,
            "name": self.name,
            "result": self.result,
            "begin": self.begin,
            "duration": self.duration,
            "parentNode": self.parent_node,
            "tokens": self.tokens.to_dict(),
            "source": self.source,
        }

    @classmethod
    def from_dict(
        cls,
        data: Dict[str, Any],
        source: Literal["dataset", "pasted"] = "dataset",
    ) -> "Trace":
        """Build Trace from detail-like dictionary payload."""
        tokens_data = data.get("tokens")
        tokens = (
            TokenUsage.from_dict(tokens_data)
            if isinstance(tokens_data, dict)
            else TokenUsage()
        )

        return cls(
            id=str(data.get("id") or ""),
            name=str(data.get("name") or "?"),
            result=str(data.get("result") or ""),
            begin=data.get("begin"),
            duration=data.get("duration"),
            tokens=tokens,
            source=source,
            parent_node=data.get("parentNode"),
        )
