"""LangChain integration for Baponi sandbox execution.

Usage::

    from baponi.langchain import code_sandbox

    agent = create_react_agent(llm, tools=[code_sandbox])

For custom configuration::

    from baponi.langchain import create_code_sandbox

    sandbox = create_code_sandbox(api_key="sk-...", thread_id="my-session")
"""

from __future__ import annotations

from typing import Any

try:
    from langchain_core.tools import BaseTool
except ImportError as e:
    raise ImportError(
        "langchain-core is required for the LangChain integration. "
        "Install it with: pip install baponi[langchain]"
    ) from e

from pydantic import BaseModel, Field

from baponi._async_client import AsyncBaponi
from baponi._base import result_to_llm_text
from baponi._client import Baponi

_TOOL_DESCRIPTION = (
    "Execute code in a secure, isolated sandbox. Supports bash, python, and node.\n\n"
    "Stateless by default. Nothing persists between calls. Pass a thread_id "
    "for persistent state (files, pip packages) across calls.\n\n"
    "Storage volumes mounted at /data/<volume-name>/ if configured.\n\n"
    "Returns stdout, stderr, and exit code."
)


class CodeSandboxInput(BaseModel):
    """Input schema for the code sandbox tool."""

    code: str = Field(description="The code to execute.")
    language: str = Field(
        default="python",
        description="Programming language: python, bash, node.",
    )
    thread_id: str | None = Field(
        default=None,
        description=(
            "Pass for persistent state across calls (files, pip packages). "
            "Must be unique. Use a descriptive prefix + random suffix "
            "(e.g. data-analysis-x8k2m9p4z1) or a UUID."
        ),
    )
    timeout: int | None = Field(
        default=None,
        description="Execution timeout in seconds (default 30). Increase for long-running tasks, max 3600.",
    )
    metadata: dict[str, str] | None = Field(
        default=None,
        description="Key-value metadata for audit trail. Max 10 keys.",
    )
    env_vars: dict[str, str] | None = Field(
        default=None,
        description="Environment variables to set. Keys must be uppercase (e.g. MY_VAR).",
    )


class BaponiCodeSandbox(BaseTool):
    """LangChain tool for Baponi sandbox execution."""

    name: str = "code_sandbox"
    description: str = _TOOL_DESCRIPTION
    args_schema: type[BaseModel] = CodeSandboxInput

    # Internal config, not part of tool schema
    _api_key: str | None
    _base_url: str | None
    _default_thread_id: str | None
    _default_timeout: int
    _default_metadata: dict[str, str] | None
    _default_env_vars: dict[str, str] | None
    _sync_client: Baponi | None
    _async_client: AsyncBaponi | None

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        thread_id: str | None = None,
        timeout: int = 30,
        metadata: dict[str, str] | None = None,
        env_vars: dict[str, str] | None = None,
    ) -> None:
        super().__init__()  # type: ignore[call-arg]
        self._api_key = api_key
        self._base_url = base_url
        self._default_thread_id = thread_id
        self._default_timeout = timeout
        self._default_metadata = metadata
        self._default_env_vars = env_vars
        self._sync_client = None
        self._async_client = None

    def _get_sync_client(self) -> Baponi:
        if self._sync_client is None:
            kwargs: dict[str, Any] = {}
            if self._api_key is not None:
                kwargs["api_key"] = self._api_key
            if self._base_url is not None:
                kwargs["base_url"] = self._base_url
            self._sync_client = Baponi(**kwargs, _client_id="baponi-python/langchain")
        return self._sync_client

    def _get_async_client(self) -> AsyncBaponi:
        if self._async_client is None:
            kwargs: dict[str, Any] = {}
            if self._api_key is not None:
                kwargs["api_key"] = self._api_key
            if self._base_url is not None:
                kwargs["base_url"] = self._base_url
            self._async_client = AsyncBaponi(**kwargs, _client_id="baponi-python/langchain")
        return self._async_client

    def _merge_params(
        self,
        thread_id: str | None,
        timeout: int | None,
        metadata: dict[str, str] | None,
        env_vars: dict[str, str] | None = None,
    ) -> tuple[str | None, int, dict[str, str] | None, dict[str, str] | None]:
        """Merge per-call params with defaults. Per-call values take priority."""
        merged_thread = thread_id if thread_id is not None else self._default_thread_id
        merged_timeout = timeout if timeout is not None else self._default_timeout
        if metadata is not None:
            merged_meta = {**(self._default_metadata or {}), **metadata}
        else:
            merged_meta = self._default_metadata
        if env_vars is not None:
            merged_env = {**(self._default_env_vars or {}), **env_vars}
        else:
            merged_env = self._default_env_vars
        return merged_thread, merged_timeout, merged_meta, merged_env

    def _run(
        self,
        code: str,
        language: str = "python",
        thread_id: str | None = None,
        timeout: int | None = None,
        metadata: dict[str, str] | None = None,
        env_vars: dict[str, str] | None = None,
    ) -> str:
        client = self._get_sync_client()
        tid, to, meta, evars = self._merge_params(thread_id, timeout, metadata, env_vars)
        result = client.execute(
            code, language=language, timeout=to, thread_id=tid, metadata=meta, env_vars=evars
        )
        return result_to_llm_text(result)

    async def _arun(
        self,
        code: str,
        language: str = "python",
        thread_id: str | None = None,
        timeout: int | None = None,
        metadata: dict[str, str] | None = None,
        env_vars: dict[str, str] | None = None,
    ) -> str:
        client = self._get_async_client()
        tid, to, meta, evars = self._merge_params(thread_id, timeout, metadata, env_vars)
        result = await client.execute(
            code, language=language, timeout=to, thread_id=tid, metadata=meta, env_vars=evars
        )
        return result_to_llm_text(result)


# Ready-to-use tool instance, lazily initializes client on first call
code_sandbox = BaponiCodeSandbox()


def create_code_sandbox(
    *,
    api_key: str | None = None,
    base_url: str | None = None,
    thread_id: str | None = None,
    timeout: int = 30,
    metadata: dict[str, str] | None = None,
    env_vars: dict[str, str] | None = None,
) -> BaponiCodeSandbox:
    """Create a configured code sandbox tool.

    Args:
        api_key: Baponi API key. Falls back to BAPONI_API_KEY env var.
        base_url: Custom API base URL (for self-hosted deployments).
        thread_id: Default thread_id for persistent state across calls.
        timeout: Default execution timeout in seconds.
        metadata: Default metadata merged into every call.
        env_vars: Default environment variables merged into every call.

    Returns:
        A configured BaponiCodeSandbox tool instance.
    """
    return BaponiCodeSandbox(
        api_key=api_key,
        base_url=base_url,
        thread_id=thread_id,
        timeout=timeout,
        metadata=metadata,
        env_vars=env_vars,
    )
