"""
Streaming event types for AugmentedLLM streaming support.

This module defines the event types and models used for streaming LLM responses,
including text deltas, tool execution events, and iteration boundaries.
"""

from enum import Enum
from typing import Any, Dict, Optional, Union
from pydantic import BaseModel, Field
import time


class StreamEventType(str, Enum):
    """Types of streaming events emitted during LLM generation.

    Streaming events provide real-time updates about the generation process,
    including incremental text content, tool usage, and iteration boundaries.
    """

    # Content events
    TEXT_DELTA = "text_delta"
    """Incremental text content as it's generated by the LLM."""

    THINKING = "thinking"
    """Extended thinking content (for models that support extended thinking)."""

    # Tool events
    TOOL_USE_START = "tool_use_start"
    """Indicates the LLM has initiated a tool call."""

    TOOL_USE_END = "tool_use_end"
    """Indicates a tool call has completed execution."""

    TOOL_RESULT = "tool_result"
    """Contains the result from tool execution."""

    # Iteration events
    ITERATION_START = "iteration_start"
    """Start of an agentic iteration in a multi-turn loop."""

    ITERATION_END = "iteration_end"
    """End of an agentic iteration."""

    # Completion events
    COMPLETE = "complete"
    """Generation has fully completed."""

    ERROR = "error"
    """An error occurred during generation."""


class StreamEvent(BaseModel):
    """A streaming event with full context.

    StreamEvent provides structured information about each stage of LLM generation,
    enabling real-time monitoring and progressive UI updates.

    Attributes:
        type: The type of streaming event
        content: Event-specific content (text delta, tool info, error message, etc.)
        iteration: The current iteration number in the agentic loop
        metadata: Additional event-specific metadata
        timestamp: Unix timestamp when the event was created
        model: The model identifier (optional)
        stop_reason: The reason generation stopped (optional)
        usage: Token usage information (optional)

    Examples:
        Text delta event:
        >>> event = StreamEvent(
        ...     type=StreamEventType.TEXT_DELTA,
        ...     content="Hello, ",
        ...     iteration=0
        ... )

        Tool use event:
        >>> event = StreamEvent(
        ...     type=StreamEventType.TOOL_USE_START,
        ...     content={"name": "search", "input": {"query": "weather"}},
        ...     iteration=1,
        ...     metadata={"tool_id": "tool_123"}
        ... )
    """

    type: StreamEventType = Field(..., description="The type of streaming event")

    content: Optional[Union[str, Dict[str, Any]]] = Field(
        default=None,
        description="Event-specific content (text, tool data, error info, etc.)",
    )

    iteration: int = Field(
        default=0, description="Current iteration number in the agentic loop"
    )

    metadata: Dict[str, Any] = Field(
        default_factory=dict, description="Additional event-specific metadata"
    )

    timestamp: float = Field(
        default_factory=lambda: time.time(),
        description="Unix timestamp when the event was created",
    )

    # Optional context fields
    model: Optional[str] = Field(
        default=None, description="Model identifier (e.g., 'claude-3-7-sonnet-latest')"
    )

    stop_reason: Optional[str] = Field(
        default=None,
        description="Reason generation stopped (e.g., 'end_turn', 'tool_use', 'max_tokens')",
    )

    usage: Optional[Dict[str, int]] = Field(
        default=None,
        description="Token usage information (input_tokens, output_tokens, etc.)",
    )

    class Config:
        """Pydantic model configuration."""

        json_schema_extra = {
            "examples": [
                {
                    "type": "text_delta",
                    "content": "Hello, world!",
                    "iteration": 0,
                    "metadata": {},
                    "timestamp": 1704724800.0,
                    "model": "claude-3-7-sonnet-latest",
                },
                {
                    "type": "tool_use_start",
                    "content": {"name": "search_tool", "input": {"query": "test"}},
                    "iteration": 1,
                    "metadata": {"tool_id": "tool_abc123"},
                    "timestamp": 1704724801.0,
                },
                {
                    "type": "complete",
                    "content": None,
                    "iteration": 2,
                    "metadata": {},
                    "timestamp": 1704724802.0,
                    "stop_reason": "end_turn",
                    "usage": {"input_tokens": 100, "output_tokens": 50},
                },
            ]
        }
