# 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 .._api import secure_sandbox

# Type alias -- toolset is deps-agnostic, works with any agent deps type.
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.

    Usage:
        from agentsh_secure_sandbox import SecureSandboxToolset
        from agentsh_secure_sandbox.adapters import e2b
        from e2b import AsyncSandbox

        raw = await AsyncSandbox.create()
        toolset = SecureSandboxToolset(e2b(raw))

        agent = Agent("claude-sonnet-4-20250514", toolsets=[toolset])

        async with agent:
            result = await agent.run("Install express and create a hello world server")
    """

    _TOOL_DEFS: dict[str, ToolDefinition] = {
        "run_command": ToolDefinition(
            name="run_command",
            description=(
                "Execute a shell command in the secure sandbox. "
                "All commands are mediated by security policy. "
                "Dangerous operations (reading secrets, network exfiltration, "
                "privilege escalation) will be blocked."
            ),
            parameters_json_schema={
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Shell command to execute",
                    },
                    "cwd": {
                        "type": "string",
                        "description": "Working directory (default: /workspace)",
                    },
                },
                "required": ["command"],
                "additionalProperties": False,
            },
        ),
        "write_file": ToolDefinition(
            name="write_file",
            description=(
                "Write content to a file in the secure sandbox. "
                "Writes to sensitive paths (secrets, config files) "
                "will be blocked by policy."
            ),
            parameters_json_schema={
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "File path to write",
                    },
                    "content": {
                        "type": "string",
                        "description": "File content",
                    },
                },
                "required": ["path", "content"],
                "additionalProperties": False,
            },
        ),
        "read_file": ToolDefinition(
            name="read_file",
            description=(
                "Read a file from the secure sandbox. "
                "Reads from sensitive paths (secrets, credentials) "
                "will be blocked by policy."
            ),
            parameters_json_schema={
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "File path to read",
                    },
                },
                "required": ["path"],
                "additionalProperties": False,
            },
        ),
    }

    def __init__(
        self,
        adapter: SandboxAdapter,
        config: SecureConfig | None = None,
        *,
        toolset_id: str = "agentsh-sandbox",
    ) -> None:
        self._adapter = adapter
        self._config = config or SecureConfig()
        self._sandbox: SecuredSandbox | None = None
        self._toolset_id = toolset_id

    # --- AbstractToolset required properties ---

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

    # --- Lifecycle ---

    async def __aenter__(self) -> Self:
        """Provision agentsh inside the sandbox (eager, not lazy)."""
        self._sandbox = await secure_sandbox(self._adapter, self._config)
        return self

    async def __aexit__(self, *args: Any) -> None:
        """Stop the agentsh session and the underlying provider sandbox."""
        if self._sandbox:
            await self._sandbox.stop()
            self._sandbox = None

    # --- Tool listing and execution ---

    @property
    def sandbox(self) -> SecuredSandbox:
        """Access the underlying SecuredSandbox for direct use."""
        if self._sandbox is None:
            raise RuntimeError(
                "Toolset not initialized. Use 'async with agent:' to "
                "enter the toolset context before running."
            )
        return self._sandbox

    # Shared pass-through validator for dict args (no strict schema enforcement
    # beyond what pydantic-ai already does from the JSON schema).
    _ARGS_VALIDATOR = SchemaValidator(core_schema.dict_schema())

    async def get_tools(
        self,
        ctx: RunContext[AgentDepsT],
    ) -> dict[str, ToolsetTool[AgentDepsT]]:
        """Return tool definitions. ctx is available but not used."""
        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:
        """Execute a tool call against the secured sandbox."""
        sb = self.sandbox

        if name == "run_command":
            result = await sb.exec(
                tool_args["command"],
                cwd=tool_args.get("cwd"),
            )
            if result.exit_code != 0:
                parts = [f"Command failed (exit {result.exit_code}):"]
                if result.stdout:
                    parts.append(f"stdout:\n{result.stdout}")
                if result.stderr:
                    parts.append(f"stderr:\n{result.stderr}")
                return "\n".join(parts)
            return result.stdout

        if name == "write_file":
            write_result = await sb.write_file(tool_args["path"], tool_args["content"])
            if write_result.success:
                return f"Written to {write_result.path}"
            return f"Write denied: {write_result.error}"

        if name == "read_file":
            read_result = await sb.read_file(tool_args["path"])
            if read_result.success:
                return read_result.content
            return f"Read denied: {read_result.error}"

        raise ValueError(f"Unknown tool: {name}")
