"""
Trace models and utilities
"""

from dataclasses import dataclass
from typing import Optional, Dict, Any
from datetime import datetime
from enum import Enum


class TraceStatus(str, Enum):
    """Trace execution status."""
    SUCCESS = "success"
    ERROR = "error"
    TIMEOUT = "timeout"
    PENDING = "pending"


@dataclass
class Trace:
    """
    Represents an agent execution trace.
    """
    id: str
    agent_id: str
    input_data: Dict[str, Any]
    output_data: Dict[str, Any]
    status: str
    error_message: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    duration_ms: Optional[int] = None
    token_count: Optional[int] = None
    cost: Optional[float] = None
    created_at: Optional[datetime] = None
    
    @classmethod
    def from_dict(cls, data: Dict) -> "Trace":
        """Create Trace from API response dictionary."""
        return cls(
            id=data["id"],
            agent_id=data["agent_id"],
            input_data=data.get("input_data", {}),
            output_data=data.get("output_data", {}),
            status=data["status"],
            error_message=data.get("error_message"),
            metadata=data.get("metadata"),
            duration_ms=data.get("duration_ms"),
            token_count=data.get("token_count"),
            cost=data.get("cost"),
            created_at=datetime.fromisoformat(data["created_at"].replace("Z", "+00:00"))
                if data.get("created_at") else None
        )
    
    def is_success(self) -> bool:
        """Check if trace execution was successful."""
        return self.status == TraceStatus.SUCCESS
    
    def is_error(self) -> bool:
        """Check if trace execution had an error."""
        return self.status == TraceStatus.ERROR
