"""Data models for the Lumo SDK."""

import json
from enum import Enum
from typing import List, Optional, Any, Dict

from pydantic import BaseModel, Field, field_validator


class MessageRole(str, Enum):
    """Message role types."""

    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"
    TOOL = "tool"

class Function(BaseModel):
    
    name: str = Field(..., description="Name of the function")
    arguments: Dict[str, Any] = Field(..., description="Arguments of the function")
    
    @field_validator('arguments', mode='before')
    @classmethod
    def parse_arguments(cls, v):
        """Parse JSON string arguments into a dictionary if needed."""
        if isinstance(v, str):
            try:
                return json.loads(v)
            except json.JSONDecodeError:
                # If it's not valid JSON, return as is (will fail validation)
                return v
        return v

class ToolCall(BaseModel):
    """Represents a tool call in a message."""

    id: str = Field(..., description="Tool call ID")
    type: str = Field(..., description="Type of tool call")
    function: Function = Field(..., description="Function details")


class Message(BaseModel):
    """Represents a message in the conversation history."""

    role: MessageRole = Field(..., description="Role of the message sender")
    content: str = Field(..., description="Content of the message")
    tool_call_id: Optional[str] = Field(None, description="ID of the tool call if applicable")
    tool_calls: Optional[List[ToolCall]] = Field(None, description="Tool calls in this message")


class MCPServer(BaseModel):
    """Represents an MCP server configuration."""

    command: str = Field(..., description="Command to run the MCP server (e.g., 'npx')")
    args: List[str] = Field(..., description="Arguments for the command")
    env: Optional[Dict[str, str]] = Field(None, description="Environment variables for the server")


class RunTaskRequest(BaseModel):
    """Request model for running a task."""

    task: str = Field(..., description="The task to execute")
    model: str = Field(..., description="Model ID (e.g., gpt-4, qwen2.5, gemini-2.0-flash)")
    base_url: str = Field(..., description="Base URL for the upstream API")
    tools: Optional[List[str]] = Field(None, description="Array of tool names to make available")
    max_steps: Optional[int] = Field(None, description="Maximum number of steps to take")
    agent_type: Optional[str] = Field(
        None, description="Type of agent to use (function-calling or mcp)"
    )
    history: Optional[List[Message]] = Field(None, description="Prior conversation context")
    mcp_servers: Optional[List[MCPServer]] = Field(
        None, description="List of MCP servers to use when agent_type is 'mcp'"
    )

    class Config:
        json_schema_extra = {
            "example": {
                "task": "What is the weather in Berlin?",
                "model": "gpt-4.1-mini",
                "base_url": "https://api.openai.com/v1/chat/completions",
                "tools": ["ExaSearchTool", "VisitWebsite"],
            }
        }

class TokenUsage(BaseModel):
    """Token usage for a step."""
    
    prompt_tokens: int = Field(..., description="Number of tokens in the prompt")
    completion_tokens: int = Field(..., description="Number of tokens in the completion")
    total_tokens: int = Field(..., description="Total number of tokens")

class Step(BaseModel):
    """Represents a step in the task execution."""

    tool_calls: List[ToolCall] = Field(default_factory=list, description="Tool calls in this step")
    llm_output: Optional[str] = Field(None, description="LLM output for this step")
    token_usage: Optional[TokenUsage] = Field(None, description="Token usage for this step")


class RunTaskResponse(BaseModel):
    """Response model for running a task."""

    final_answer: str = Field(..., description="Final answer from the task execution")
    steps: List[Step] = Field(..., description="Steps taken during task execution")
    token_usage: Optional[TokenUsage] = Field(None, description="Token usage for the entire task")


class StreamTokenEvent(BaseModel):
    """Token event from streaming response."""

    type: str = Field("token", description="Event type")
    content: str = Field(..., description="Token content")


class StreamStepEvent(BaseModel):
    """Step event from streaming response."""

    type: str = Field("step", description="Event type")
    step: Step = Field(..., description="Step data containing tool_calls and output")


class StreamDoneEvent(BaseModel):
    """Done event from streaming response."""

    token_usage: Optional[TokenUsage] = Field(None, description="Token usage for the entire task")
    type: str = Field("done", description="Event type")

