"""Pydantic types for the Influxion SDK.

Parameter types are provided as a convenience for callers who want to construct
parameter objects before passing them to the client. The client methods also
accept individual keyword arguments directly.

Agent config types represent configuration returned by the API (e.g. from
:meth:`~influxion.InfluxionClient.get_config`).
"""

from __future__ import annotations

from typing import Any
from uuid import UUID

from pydantic import BaseModel, PrivateAttr, model_validator


class CreateSessionParams(BaseModel):
    """Parameters accepted by :meth:`~influxion.InfluxionClient.create_session`."""

    name: str
    project_id: str
    agent_id: str
    tags: list[str] | None = None
    metadata: dict[str, Any] | None = None


class UpdateSessionParams(BaseModel):
    """Parameters accepted by :meth:`~influxion.InfluxionClient.update_session`."""

    success: bool | None = None
    failure_message: str | None = None
    tags: list[str] | None = None
    metadata: dict[str, Any] | None = None


class CreateArtifactParams(BaseModel):
    """Parameters accepted by :meth:`~influxion.InfluxionClient.create_session_artifact`."""

    session_id: str
    name: str
    artifact: str
    type: str = "text"
    tags: list[str] | None = None
    metadata: dict[str, Any] | None = None


class CreateToolCallParams(BaseModel):
    """Parameters accepted by :meth:`~influxion.InfluxionClient.create_session_tool_call`."""

    session_id: str
    name: str
    input: dict[str, Any] | None = None
    output: dict[str, Any] | None = None
    tags: list[str] | None = None
    metadata: dict[str, Any] | None = None


# ---------------------------------------------------------------------------
# Agent config response types
# ---------------------------------------------------------------------------


class AgentConfigPromptMessage(BaseModel):
    """A prompt message configuration for an agent."""

    role: str
    content: str


class AgentConfigPrompt(BaseModel):
    """A prompt configuration for an agent."""

    label: str
    messages: list[AgentConfigPromptMessage]


class AgentConfigModel(BaseModel):
    """A model configuration for an agent."""

    label: str
    base_url: str | None
    name: str


class AgentConfig(BaseModel):
    """Configuration parameters stored for an agent in Influxion.

    Returned by :meth:`~influxion.InfluxionClient.get_config`. Use
    :meth:`get_model` and :meth:`get_prompt` to look up entries by label.
    """

    id: UUID
    config_id: str
    name: str
    description: str | None
    is_default: bool
    models: list[AgentConfigModel]
    prompts: list[AgentConfigPrompt]

    _models_by_label: dict[str, AgentConfigModel] = PrivateAttr(default_factory=dict)
    _prompts_by_label: dict[str, AgentConfigPrompt] = PrivateAttr(default_factory=dict)

    @model_validator(mode="after")
    def _build_label_indexes(self) -> AgentConfig:
        self._models_by_label = {m.label: m for m in self.models}
        self._prompts_by_label = {p.label: p for p in self.prompts}
        return self

    def get_model(self, label: str) -> AgentConfigModel | None:
        """Return the model configuration with the given label, or ``None`` if not found."""
        return self._models_by_label.get(label)

    def get_prompt(self, label: str) -> AgentConfigPrompt | None:
        """Return the prompt configuration with the given label, or ``None`` if not found."""
        return self._prompts_by_label.get(label)
