"""
Data models for trace objects using dataclasses for type safety.
"""
from dataclasses import dataclass
from typing import Any, Dict, 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,
        }


@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]
    trace_name: str
    tokens: TokenUsage
    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,
            "traceName": self.trace_name,
            "tokens": self.tokens.to_dict(),
        }

    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,
            "traceName": self.trace_name,
            "parentNode": self.parent_node,
            "tokens": self.tokens.to_dict(),
        }


# Backward compatibility aliases (to be removed after migration)
TraceDetail = Trace
TraceMetadata = Trace
