from datetime import datetime
from typing import Any, List, Optional

from pydantic import BaseModel, Field


class Usage(BaseModel):
    prompt_tokens: Optional[int] = Field(
        default=None,
        description="Number of tokens used in the prompt",
    )
    completion_tokens: Optional[int] = Field(
        default=None,
        description="Number of tokens used in the completion",
    )
    total_tokens: Optional[int] = Field(
        default=None, description="Total number of tokens used"
    )


class ModelConfig(BaseModel):
    model_name: Optional[str] = Field(
        default=None,
        description="Name of the model used for generation",
    )
    temperature: Optional[float] = Field(
        default=None,
        description="Temperature setting used for generation",
    )
    top_p: Optional[float] = Field(
        default=None, description="Top-p setting used for generation"
    )
    max_tokens: Optional[int] = Field(
        default=None,
        description="Maximum number of tokens to generate",
    )
    frequency_penalty: Optional[float] = Field(
        default=None,
        description="Frequency penalty used for generation",
    )
    presence_penalty: Optional[float] = Field(
        default=None,
        description="Presence penalty used for generation",
    )


class AgentCompletionResponse(BaseModel):
    id: Optional[str] = Field(
        default=None, description="Unique identifier for the response"
    )
    agent_name: Optional[str] = Field(
        default=None,
        description="Name of the agent that generated the response",
    )
    agent_description: Optional[str] = Field(
        default=None, description="Description of the agent"
    )
    outputs: Optional[List[Any]] = Field(
        default=None,
        description="List of outputs generated by the agent",
    )
    usage: Optional[Usage] = Field(
        default=None, description="Token usage statistics"
    )
    model_config: Optional[ModelConfig] = Field(
        default=None, description="Model configuration"
    )
    timestamp: Optional[str] = Field(
        default_factory=lambda: datetime.now().isoformat(),
        description="Timestamp of when the response was generated",
    )
