"""
LLM types for the Glueco Gateway.

Provides typed dataclasses for chat completion requests and responses,
compatible with OpenAI API format.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Literal, Optional, Union


# =============================================================================
# REQUEST TYPES
# =============================================================================

@dataclass
class ChatMessage:
    """A single message in a chat conversation.
    
    Attributes:
        role: Message role (system, user, assistant, tool).
        content: Message content.
        name: Optional sender name.
    """
    role: Literal["system", "user", "assistant", "tool"]
    content: str
    name: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to API-compatible dictionary."""
        result: Dict[str, Any] = {
            "role": self.role,
            "content": self.content,
        }
        if self.name:
            result["name"] = self.name
        return result


@dataclass
class ChatCompletionRequest:
    """Request for a chat completion.
    
    Attributes:
        model: Model name (e.g., "gpt-4o", "llama-3.1-8b-instant").
        messages: Conversation messages.
        temperature: Sampling temperature (0.0 to 2.0).
        max_tokens: Maximum tokens to generate.
        stream: Whether to stream responses.
        response_format: Response format specification.
    """
    model: str
    messages: List[Union[ChatMessage, Dict[str, Any]]]
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    stream: bool = False
    response_format: Optional[Dict[str, Any]] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to API-compatible dictionary."""
        # Normalize messages
        normalized_messages = []
        for msg in self.messages:
            if isinstance(msg, ChatMessage):
                normalized_messages.append(msg.to_dict())
            else:
                normalized_messages.append(msg)
        
        result: Dict[str, Any] = {
            "model": self.model,
            "messages": normalized_messages,
        }
        
        if self.temperature is not None:
            result["temperature"] = self.temperature
        if self.max_tokens is not None:
            result["max_tokens"] = self.max_tokens
        if self.stream:
            result["stream"] = True
        if self.response_format is not None:
            result["response_format"] = self.response_format
        
        return result


# =============================================================================
# RESPONSE TYPES
# =============================================================================

@dataclass
class ChatCompletionUsage:
    """Token usage information.
    
    Attributes:
        prompt_tokens: Tokens in the prompt.
        completion_tokens: Tokens in the completion.
        total_tokens: Total tokens used.
    """
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionUsage":
        """Create from API response dictionary."""
        return cls(
            prompt_tokens=data.get("prompt_tokens", 0),
            completion_tokens=data.get("completion_tokens", 0),
            total_tokens=data.get("total_tokens", 0),
        )


@dataclass
class ChatCompletionChoice:
    """A single choice in a chat completion response.
    
    Attributes:
        index: Choice index.
        message: Response message.
        finish_reason: Why generation stopped.
    """
    index: int
    message: ChatMessage
    finish_reason: Optional[str] = None
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionChoice":
        """Create from API response dictionary."""
        message_data = data.get("message", {})
        return cls(
            index=data.get("index", 0),
            message=ChatMessage(
                role=message_data.get("role", "assistant"),
                content=message_data.get("content", ""),
                name=message_data.get("name"),
            ),
            finish_reason=data.get("finish_reason"),
        )


@dataclass
class ChatCompletion:
    """Response from a chat completion request.
    
    Attributes:
        id: Completion ID.
        model: Model used.
        choices: Response choices.
        usage: Token usage (optional).
        created: Unix timestamp of creation.
    """
    id: str
    model: str
    choices: List[ChatCompletionChoice]
    usage: Optional[ChatCompletionUsage] = None
    created: Optional[int] = None
    
    @property
    def content(self) -> str:
        """Get the first choice's message content (convenience accessor)."""
        if self.choices:
            return self.choices[0].message.content
        return ""
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletion":
        """Create from API response dictionary."""
        choices = [
            ChatCompletionChoice.from_dict(c)
            for c in data.get("choices", [])
        ]
        
        usage = None
        if "usage" in data:
            usage = ChatCompletionUsage.from_dict(data["usage"])
        
        return cls(
            id=data.get("id", ""),
            model=data.get("model", ""),
            choices=choices,
            usage=usage,
            created=data.get("created"),
        )


# =============================================================================
# STREAMING TYPES
# =============================================================================

@dataclass
class ChatCompletionChunkDelta:
    """Delta content in a streaming chunk.
    
    Attributes:
        role: Message role (usually only in first chunk).
        content: Content fragment.
    """
    role: Optional[str] = None
    content: Optional[str] = None
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionChunkDelta":
        """Create from API response dictionary."""
        return cls(
            role=data.get("role"),
            content=data.get("content"),
        )


@dataclass
class ChatCompletionChunkChoice:
    """A choice in a streaming chunk.
    
    Attributes:
        index: Choice index.
        delta: Content delta.
        finish_reason: Why generation stopped (in final chunk).
    """
    index: int
    delta: ChatCompletionChunkDelta
    finish_reason: Optional[str] = None
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionChunkChoice":
        """Create from API response dictionary."""
        return cls(
            index=data.get("index", 0),
            delta=ChatCompletionChunkDelta.from_dict(data.get("delta", {})),
            finish_reason=data.get("finish_reason"),
        )


@dataclass
class ChatCompletionChunk:
    """A chunk in a streaming chat completion.
    
    Attributes:
        id: Completion ID.
        model: Model used.
        choices: Response choices.
        created: Unix timestamp.
    """
    id: str
    model: str
    choices: List[ChatCompletionChunkChoice]
    created: Optional[int] = None
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionChunk":
        """Create from API response dictionary."""
        choices = [
            ChatCompletionChunkChoice.from_dict(c)
            for c in data.get("choices", [])
        ]
        return cls(
            id=data.get("id", ""),
            model=data.get("model", ""),
            choices=choices,
            created=data.get("created"),
        )
