# src/agentsh_secure_sandbox/toolset/pydantic_ai.py
from __future__ import annotations

from typing import Any

from typing_extensions import Self

from pydantic_ai import RunContext
from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets import AbstractToolset, ToolsetTool
from pydantic_core import SchemaValidator, core_schema

from ..core.types import SecuredSandbox, SandboxAdapter, SecureConfig
from ._base import SecureSandboxTools, TOOL_DEFS

AgentDepsT = Any


class SecureSandboxToolset(AbstractToolset[AgentDepsT]):
    """Pydantic AI Toolset that wraps a sandbox provider with agentsh
    security enforcement.

    Lifecycle:
    - One sandbox per toolset instance (1:1 ownership).
    - Not concurrency-safe. Do not share across concurrent agent runs.
    - After stop(), the instance is dead. Create a new one for a new run.
    - Provisioning happens eagerly in __aenter__, not lazily per tool call.
    """

    _TOOL_DEFS: dict[str, ToolDefinition] = {
        name: ToolDefinition(
            name=defn["name"],
            description=defn["description"],
            parameters_json_schema=defn["parameters"],
        )
        for name, defn in TOOL_DEFS.items()
    }

    def __init__(
        self,
        adapter: SandboxAdapter,
        config: SecureConfig | None = None,
        *,
        toolset_id: str = "agentsh-sandbox",
    ) -> None:
        self._tools = SecureSandboxTools(adapter, config)
        self._toolset_id = toolset_id

    @property
    def id(self) -> str:
        return self._toolset_id

    async def __aenter__(self) -> Self:
        await self._tools.__aenter__()
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self._tools.__aexit__(*args)

    @property
    def sandbox(self) -> SecuredSandbox:
        return self._tools.sandbox

    _ARGS_VALIDATOR = SchemaValidator(core_schema.dict_schema())

    async def get_tools(
        self,
        ctx: RunContext[AgentDepsT],
    ) -> dict[str, ToolsetTool[AgentDepsT]]:
        return {
            name: ToolsetTool(
                tool_def=defn,
                toolset=self,
                max_retries=0,
                args_validator=self._ARGS_VALIDATOR,
            )
            for name, defn in self._TOOL_DEFS.items()
        }

    async def call_tool(
        self,
        name: str,
        tool_args: dict[str, Any],
        ctx: RunContext[AgentDepsT],
        tool: ToolsetTool[AgentDepsT],
    ) -> Any:
        return await self._tools.call_tool(name, tool_args)
