"""Agent model and related types for the Swarm framework."""

import secrets
import string
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Type, Union

from pydantic import BaseModel, ConfigDict, Field, model_validator


def _generate_agent_id() -> str:
    """Generate a random agent ID in format agent-xxxxx (5 alphanumeric chars)."""
    chars = string.ascii_lowercase + string.digits
    suffix = "".join(secrets.choice(chars) for _ in range(5))
    return f"agent-{suffix}"


@dataclass
class ModelSettings:
    """Centralized settings for LLM calls across all providers."""

    # Generation parameters
    temperature: Optional[float] = None
    """The temperature to use when calling the model (0.0-2.0)."""

    max_tokens: Optional[int] = None
    """The maximum number of output tokens to generate."""

    top_p: Optional[float] = None
    """The top_p to use when calling the model."""

    frequency_penalty: Optional[float] = None
    """The frequency penalty to use when calling the model."""

    presence_penalty: Optional[float] = None
    """The presence penalty to use when calling the model."""

    stop: Optional[List[str]] = None
    """Stop sequences to use when generating."""

    seed: Optional[int] = None
    """Random seed for deterministic generation."""

    # Tool/Function calling
    tool_choice: Optional[str] = None
    """The tool choice to use when calling the model (auto/none/required)."""

    parallel_tool_calls: Optional[bool] = None
    """Controls whether the model can make multiple parallel tool calls."""

    # Response format
    response_format: Optional[Dict[str, Any]] = None
    """Response format (e.g., {"type": "json_object"})."""

    # Streaming and monitoring
    stream: bool = True
    """Whether to stream the response."""

    include_usage: Optional[bool] = None
    """Whether to include usage chunk for token tracking."""

    # Timeout and retries
    timeout: int = 60
    """Request timeout in seconds."""

    max_retries: int = 3
    """Maximum number of retry attempts."""

    def to_completion_params(self, has_tools: bool = False) -> Dict[str, Any]:
        """Extract completion request parameters, excluding None values.

        Args:
            has_tools: Whether the agent has tools defined (used for parallel_tool_calls)

        Returns only the fields that CompletionRequest expects,
        filtering out any None values.
        """
        # Core completion fields
        completion_fields = {
            "temperature",
            "max_tokens",
            "top_p",
            "frequency_penalty",
            "presence_penalty",
            "stop",
            "seed",
            "response_format",
            "tool_choice",
            "include_usage",
            "timeout",
            "max_retries",
        }

        # Only include parallel_tool_calls if there are tools
        if has_tools:
            completion_fields.add("parallel_tool_calls")

        return {
            field: getattr(self, field)
            for field in completion_fields
            if getattr(self, field) is not None
        }

    # Optional metadata
    metadata: Optional[Dict[str, str]] = None
    """Metadata to include with the model response call."""

    top_logprobs: Optional[int] = None
    """Number of top tokens to return logprobs for."""

    extra_body: Optional[Dict[str, Any]] = None
    """Additional body fields to provide with the request."""

    extra_headers: Optional[Dict[str, Any]] = None
    """Additional headers to provide with the request."""

    extra_args: Optional[Dict[str, Any]] = None
    """Additional arguments to provide with the request."""


class AgentSkill(BaseModel):
    """Skill that an agent can perform (inspired by A2A protocol)."""

    model_config = ConfigDict(extra="forbid")

    id: str = Field(..., description="Skill identifier, e.g., 'gather-requirements'")
    name: str = Field(..., description="Human-readable skill name")
    description: str = Field(..., description="What this skill does")
    tags: List[str] = Field(default_factory=list, description="Categorization tags")
    examples: List[str] = Field(
        default_factory=list, description="Example user inputs that trigger this skill"
    )


class AgentCard(BaseModel):
    """Agent identity and capabilities (inspired by A2A protocol)."""

    model_config = ConfigDict(extra="forbid")

    id: str = Field(..., description="Agent ID in format agent-xxxxx")
    name: str = Field(..., description="Human-readable agent name")
    description: str = Field(..., description="What this agent does")
    version: str = Field(default="1.0.0", description="Agent version")
    skills: List[AgentSkill] = Field(
        default_factory=list, description="Skills this agent provides"
    )


class AgentFunction(BaseModel):
    """Serialized representation of a function."""

    name: str
    module: str
    qualname: str


# Type aliases for functions
AgentFunctionType = Union[
    Callable[[], Union[str, "Agent", Dict[str, Any]]], AgentFunction
]
ToolFunctionType = Callable[..., Any]


class Agent(BaseModel):
    """Agent model with tools, hooks, and configuration."""

    model_config = ConfigDict(extra="forbid", protected_namespaces=())

    # Identity (AgentCard fields)
    id: Optional[str] = Field(
        default=None, description="Agent ID (auto-generated if not provided)"
    )
    name: str = Field(default="Agent", description="Human-readable agent name")
    description: str = Field(default="", description="What this agent does")
    version: str = Field(default="1.0.0", description="Agent version")
    skills: List[AgentSkill] = Field(
        default_factory=list, description="Skills this agent provides"
    )

    # LLM configuration
    model: Union[str, List[str]] = "openai/gpt-4o"
    instructions: Union[str, Callable[[Any], str]] = "You are a helpful agent."
    tools: List[ToolFunctionType] = []

    # Lifecycle hooks (observational — cannot alter execution flow)
    hooks: List[Any] = Field(default_factory=list)

    # Guardrails (validation — can halt execution via tripwire)
    input_guardrails: List[Any] = Field(default_factory=list)
    output_guardrails: List[Any] = Field(default_factory=list)

    # Optional structured output type
    output_type: Optional[Type[BaseModel]] = None

    # Model settings for fine-tuning LLM behavior
    model_settings: Optional[ModelSettings] = None

    @model_validator(mode="after")
    def _generate_id_if_missing(self) -> "Agent":
        """Auto-generate agent ID if not provided."""
        if self.id is None:
            object.__setattr__(self, "id", _generate_agent_id())
        return self

    @property
    def card(self) -> AgentCard:
        """Get agent's identity card."""
        return AgentCard(
            id=self.id or _generate_agent_id(),
            name=self.name,
            description=self.description,
            version=self.version,
            skills=self.skills,
        )

    def model_dump(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
        data = super().model_dump(*args, **kwargs)
        # Handle instructions
        if callable(data["instructions"]):
            data["instructions"] = data["instructions"]()

        # Handle functions
        if self.tools:
            data["functions"] = [
                {
                    "name": getattr(f, "__name__", getattr(f, "name", "unknown")),
                    "module": getattr(f, "__module__", getattr(f, "module", "unknown")),
                    "qualname": getattr(
                        f, "__qualname__", getattr(f, "qualname", "unknown")
                    ),
                }
                for f in self.tools
            ]

        # Exclude hooks and guardrails from serialization (not JSON-serializable)
        data.pop("hooks", None)
        data.pop("input_guardrails", None)
        data.pop("output_guardrails", None)

        # Exclude output_type from serialization (it's a type)
        data.pop("output_type", None)

        # Convert ModelSettings dataclass to dict if present
        if self.model_settings:
            from dataclasses import asdict

            data["model_settings"] = asdict(self.model_settings)

        return data

    @classmethod
    def model_validate(cls, obj: Any, *args: Any, **kwargs: Any) -> "Agent":
        if isinstance(obj, dict):
            # Handle functions
            if "functions" in obj and obj["functions"]:
                obj["functions"] = [
                    AgentFunction(**f) if isinstance(f, dict) else f
                    for f in obj["functions"]
                ]
        return super().model_validate(obj, *args, **kwargs)
