from enum import Enum
from typing import Any, Optional

from galtea.utils.from_camel_case_base_model import FromCamelCaseBaseModel


class NodeType(str, Enum):
    """Types of nodes in a trace hierarchy.

    .. deprecated::
        This enum is deprecated and will be removed in a future version.
        Use :class:`TraceType` instead.

    Used to categorize traces based on their functional role.
    """

    import warnings

    warnings.warn(
        "NodeType is deprecated and will be removed in a future version. Use TraceType instead.",
        DeprecationWarning,
        stacklevel=2,
    )

    TOOL = "TOOL"
    """Tool/function calls (e.g., external APIs)"""

    CHAIN = "CHAIN"
    """Links between different application steps"""

    RETRIEVER = "RETRIEVER"
    """Data retrieval steps (vector store, database)"""

    CUSTOM = "SPAN"
    """Generic durations of work in a trace"""

    LLM = "GENERATION"
    """AI model generations including prompts, token usage, and costs"""


class TraceType(str, Enum):
    """Types of traces.

    Used to categorize traces based on their functional role.

    See: https://docs.galtea.ai/concepts/product/version/session/trace#trace-types
    """

    SPAN = "SPAN"
    """Generic durations of work in a trace"""

    GENERATION = "GENERATION"
    """AI model generations including prompts, token usage, and costs"""

    EVENT = "EVENT"
    """Discrete point-in-time events"""

    AGENT = "AGENT"
    """Agent that orchestrates flow and uses tools with LLM guidance"""

    TOOL = "TOOL"
    """Tool/function calls (e.g., external APIs)"""

    CHAIN = "CHAIN"
    """Links between different application steps"""

    RETRIEVER = "RETRIEVER"
    """Data retrieval steps (vector store, database)"""

    EVALUATOR = "EVALUATOR"
    """Functions that assess LLM outputs"""

    EMBEDDING = "EMBEDDING"
    """Embedding model calls"""

    GUARDRAIL = "GUARDRAIL"
    """Components that protect against malicious content"""


class TraceBase(FromCamelCaseBaseModel):
    """Base model for creating traces.

    A trace represents a single operation during an inference,
    capturing input, output, timing, and error information.
    """

    inference_result_id: str
    name: str
    id: Optional[str] = None
    description: Optional[str] = None
    type: Optional[TraceType] = None
    input_data: Optional[Any] = None
    output_data: Optional[Any] = None
    error: Optional[str] = None
    latency_ms: Optional[float] = None
    metadata: Optional[Any] = None
    parent_trace_id: Optional[str] = None
    start_time: Optional[str] = None
    end_time: Optional[str] = None


class Trace(TraceBase):
    """Complete trace model returned from API.

    Includes all fields from TraceBase plus server-generated fields
    like id, created_at, and deleted_at.
    """

    created_at: str
    deleted_at: Optional[str] = None
