"""
LLM response and result types for LLM interactions.

This module defines types for LLM API responses:
- ChatCompletionResult: Complete LLM API response
- ChatCompletionChunk: Streaming response chunk
- ModelClientError: Error exception

Message and usage types are defined in the root agentbyte.messages module.

All types use Pydantic for validation and serialization.
"""
from typing import Any, Dict, Optional

from pydantic import BaseModel, ConfigDict, Field

# Import message types and usage from root module (picoagents-aligned)
from agentbyte.messages import AssistantMessage, ToolCallRequest, Usage


class ChatCompletionResult(BaseModel):
    """
    Response from an LLM model API call.

    Represents the complete result of calling an LLM model, including:
    - The generated message
    - Token usage statistics
    - Model name
    - Provider-specific metadata
    - Optional structured output (if output_format was specified)

    Attributes:
        message: The AssistantMessage generated by the model
        usage: Token usage statistics for cost tracking
        model: Name of the model used
        metadata: Provider-specific data (finish_reason, etc.)
        finish_reason: Completion stop reason (if provided by the model)
        structured_output: Parsed structured output if output_format was provided (optional)

    Example:
        result = ChatCompletionResult(
            message=AssistantMessage(
                content="Paris is the capital of France."
            ),
            usage=Usage(tokens_input=50, tokens_output=15),
            model="gpt-4.1-mini",
            metadata={"finish_reason": "stop"}
        )
    """
    message: AssistantMessage = Field(
        ...,
        description="The generated message from the model"
    )
    usage: Usage = Field(
        ...,
        description="Token usage statistics"
    )
    model: str = Field(
        ...,
        description="Name of the model used (e.g., 'gpt-4.1-mini')"
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Provider-specific metadata (finish_reason, etc.)"
    )
    finish_reason: Optional[str] = Field(
        default=None,
        description="Completion stop reason (if provided by the model)"
    )
    structured_output: Optional[Any] = Field(
        default=None,
        description="Parsed structured output if output_format was provided to create()"
    )


class ChatCompletionChunk(BaseModel):
    """
    Single chunk in a streaming LLM response.

    When streaming is enabled, the response comes in chunks as the model
    generates tokens. Each chunk contains a small piece of the total response.

    Attributes:
        content: Text content of this chunk
        model: Name of the model
        is_complete: Whether this is the final chunk in the stream
        tool_call_chunk: Partial tool call data (if tool calls are streaming)
        usage: Token usage statistics (only present in final usage chunk)
        metadata: Provider-specific metadata

    Example:
        # Streaming example
        async for chunk in client.create_stream(messages=[...]):
            print(chunk.content, end="", flush=True)
            if chunk.is_complete:
                print("\nDone")
    """
    content: str = Field(
        ...,
        description="Text content of this chunk"
    )
    model: str = Field(
        ...,
        description="Name of the model"
    )
    is_complete: bool = Field(
        ...,
        description="Whether this is the final chunk in the stream"
    )
    tool_call_chunk: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Partial tool call data when tool calls are streamed"
    )
    usage: Optional[Usage] = Field(
        default=None,
        description="Token usage statistics (final usage chunk only)"
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Provider-specific metadata"
    )


class ModelPricing(BaseModel):
    """Typed pricing definition for model token costs (USD per 1M tokens)."""

    model_config = ConfigDict(frozen=True)

    input_per_1m: float = Field(
        ...,
        ge=0,
        description="Input token cost in USD per 1M tokens",
    )
    output_per_1m: float = Field(
        ...,
        ge=0,
        description="Output token cost in USD per 1M tokens",
    )

    def calculate_cost(self, input_tokens: int, output_tokens: int = 0) -> float:
        """
        Calculate total cost in USD from token counts.

        Args:
            input_tokens: Number of input tokens used
            output_tokens: Number of output tokens generated (default 0 for embeddings)

        Returns:
            Total cost in USD, rounded to 15 decimal places
        """
        input_cost = (input_tokens / 1_000_000) * self.input_per_1m
        output_cost = (output_tokens / 1_000_000) * self.output_per_1m
        return round(input_cost + output_cost, 15)


class EmbeddingResult(BaseModel):
    """
    Response payload for embedding generation.

    Attributes:
        embedding: First embedding vector in request order. For single-input
            calls, this is the embedding for the provided text.
        embeddings: Embedding vectors in request order.
        model: Model/deployment used by the provider.
        usage: Token usage and execution metadata.
        metadata: Provider-specific metadata.

    Invariant:
        `embedding == embeddings[0]`

        Single-input `create()` calls return exactly one vector in
        `embeddings`, while batch `create_batch()` calls return all vectors in
        request order and expose the first vector through `embedding`.
    """

    embedding: list[float] = Field(
        ...,
        description="First embedding vector in request order",
    )
    embeddings: list[list[float]] = Field(
        ...,
        description="Embedding vectors in request order",
    )
    model: str = Field(
        ...,
        description="Name of the model or deployment used",
    )
    usage: Usage = Field(
        ...,
        description="Token usage statistics",
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Provider-specific metadata",
    )


class ModelClientError(Exception):
    """
    Base exception for model client errors.

    Raised when there's an error communicating with an LLM provider,
    such as authentication failures, rate limiting, or API errors.

    Example:
        try:
            result = await client.create(messages=[...])
        except ModelClientError as e:
            print(f"LLM error: {e}")
    """
    pass


__all__ = [
    # Re-exported message types for convenience
    "AssistantMessage",
    "ToolCallRequest",
    # LLM response types
    "Usage",
    "ModelPricing",
    "ChatCompletionResult",
    "ChatCompletionChunk",
    "EmbeddingResult",
    # Exceptions
    "ModelClientError",
]

