# src/agentsh_secure_sandbox/toolset/langchain.py
"""LangChain/LangGraph integration — SecureSandboxToolkit.

Exposes agentsh-secured sandbox tools as LangChain BaseTool instances
for use with create_react_agent, ToolNode, or any LangChain agent.
"""
from __future__ import annotations

from typing import Any, Type

from typing_extensions import Self

from pydantic import BaseModel, Field, PrivateAttr
from langchain_core.tools import BaseTool

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


# --- Pydantic input schemas for BaseTool args_schema ---


class RunCommandInput(BaseModel):
    command: str = Field(description="Shell command to execute")
    cwd: str | None = Field(default=None, description="Working directory (default: /workspace)")


class WriteFileInput(BaseModel):
    path: str = Field(description="File path to write")
    content: str = Field(description="File content")


class ReadFileInput(BaseModel):
    path: str = Field(description="File path to read")


# --- BaseTool subclasses ---


class _RunCommandTool(BaseTool):
    name: str = "run_command"
    description: str = TOOL_DEFS["run_command"]["description"]
    args_schema: Type[BaseModel] = RunCommandInput
    _tools: SecureSandboxTools = PrivateAttr()

    def __init__(self, tools: SecureSandboxTools, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._tools = tools

    def _run(self, **kwargs: Any) -> str:
        raise NotImplementedError(
            "SecureSandboxToolkit tools are async-only. "
            "Use ainvoke() or an async LangGraph agent."
        )

    async def _arun(self, command: str, cwd: str | None = None, **kwargs: Any) -> str:
        return await self._tools.call_tool("run_command", {"command": command, "cwd": cwd})


class _WriteFileTool(BaseTool):
    name: str = "write_file"
    description: str = TOOL_DEFS["write_file"]["description"]
    args_schema: Type[BaseModel] = WriteFileInput
    _tools: SecureSandboxTools = PrivateAttr()

    def __init__(self, tools: SecureSandboxTools, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._tools = tools

    def _run(self, **kwargs: Any) -> str:
        raise NotImplementedError(
            "SecureSandboxToolkit tools are async-only. "
            "Use ainvoke() or an async LangGraph agent."
        )

    async def _arun(self, path: str, content: str, **kwargs: Any) -> str:
        return await self._tools.call_tool("write_file", {"path": path, "content": content})


class _ReadFileTool(BaseTool):
    name: str = "read_file"
    description: str = TOOL_DEFS["read_file"]["description"]
    args_schema: Type[BaseModel] = ReadFileInput
    _tools: SecureSandboxTools = PrivateAttr()

    def __init__(self, tools: SecureSandboxTools, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._tools = tools

    def _run(self, **kwargs: Any) -> str:
        raise NotImplementedError(
            "SecureSandboxToolkit tools are async-only. "
            "Use ainvoke() or an async LangGraph agent."
        )

    async def _arun(self, path: str, **kwargs: Any) -> str:
        return await self._tools.call_tool("read_file", {"path": path})


# --- Toolkit ---


class SecureSandboxToolkit:
    """LangChain toolkit providing agentsh-secured sandbox tools.

    Not concurrency-safe. Do not share across concurrent agent runs.

    Usage::

        toolkit = SecureSandboxToolkit(daytona(raw))
        async with toolkit:
            agent = create_react_agent(model, tools=toolkit.get_tools())
            result = agent.invoke({"messages": [...]})
    """

    def __init__(
        self,
        adapter: SandboxAdapter,
        config: SecureConfig | None = None,
    ) -> None:
        self._tools = SecureSandboxTools(adapter, config)

    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

    def get_tools(self) -> list[BaseTool]:
        return [
            _RunCommandTool(self._tools),
            _WriteFileTool(self._tools),
            _ReadFileTool(self._tools),
        ]
