"""LangChain Deep Agents integration for Baponi sandbox execution.

Usage::

    from baponi.deepagents import BaponiSandbox
    from deepagents import create_deep_agent

    sandbox = BaponiSandbox()
    agent = create_deep_agent("anthropic:claude-sonnet-4-6", backend=sandbox)

For custom configuration::

    from baponi.deepagents import create_sandbox

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

Context manager support::

    with BaponiSandbox(api_key="sk-...") as sandbox:
        result = sandbox.execute("echo hello")

Streaming::

    stream = sandbox.execute_stream("long-running-command")
    with stream:
        for event in stream:
            print(event)

Provider for programmatic lifecycle::

    from baponi.deepagents import BaponiSandboxProvider

    provider = BaponiSandboxProvider(api_key="sk-...")
    sandbox = provider.get_or_create(sandbox_id="my-session")
"""

from __future__ import annotations

import asyncio
import os
import time
import uuid
from typing import TYPE_CHECKING, Any

import httpx

try:
    from deepagents.backends.protocol import (
        ExecuteResponse,
        FileDownloadResponse,
        FileUploadResponse,
    )
    from deepagents.backends.sandbox import BaseSandbox
except ImportError as e:
    raise ImportError(
        "deepagents is required for the Deep Agents integration. "
        "Install it with: pip install baponi[deepagents]"
    ) from e

try:
    from deepagents_cli.integrations.sandbox_provider import (  # type: ignore[import-not-found]
        SandboxProvider as _ProviderBase,
    )
except ImportError:  # pragma: no cover
    _ProviderBase = object  # type: ignore[assignment,misc,unused-ignore]

from baponi._async_client import AsyncBaponi
from baponi._client import Baponi
from baponi._constants import DEFAULT_EXECUTION_TIMEOUT
from baponi.exceptions import ThreadBusyError

if TYPE_CHECKING:
    from typing_extensions import Self

    from baponi._streaming import ExecutionStream

_MAX_OUTPUT = 1_048_576  # 1 MB
_RETRY_DELAYS = (0.5, 1.0, 2.0)
_HOME_PREFIX = "/home/baponi/"


def _to_relative_path(path: str) -> str:
    """Strip sandbox home prefix to get a path relative to thread root.

    The sandbox mounts thread storage at /home/baponi. The Files API expects
    paths relative to the thread root, so /home/baponi/foo.txt becomes foo.txt.
    """
    if path.startswith(_HOME_PREFIX):
        return path[len(_HOME_PREFIX):]
    return path.lstrip("/")


class BaponiSandbox(BaseSandbox):
    """Baponi backend for LangChain Deep Agents.

    Implements the BaseSandbox ABC, providing 7 filesystem tools
    (read, write, edit, ls, grep, glob, execute) automatically.
    """

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        thread_id: str | None = None,
        timeout: int = DEFAULT_EXECUTION_TIMEOUT,
        metadata: dict[str, str] | None = None,
        env_vars: dict[str, str] | None = None,
        sub_paths: list[str] | None = None,
    ) -> None:
        self._api_key = api_key
        self._base_url = base_url
        self._thread_id = thread_id or f"deep-{uuid.uuid4().hex[:16]}"
        self._timeout = timeout
        self._metadata = metadata
        self._env_vars = env_vars
        self._sub_paths = sub_paths
        self._sync_client: Baponi | None = None
        self._async_client: AsyncBaponi | None = None

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

    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/deepagents")
        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/deepagents")
        return self._async_client

    def _to_execute_response(self, result: Any) -> ExecuteResponse:
        parts = []
        if result.stdout:
            parts.append(result.stdout)
        if result.stderr:
            parts.append(result.stderr)
        output = "\n".join(parts)
        output_bytes = output.encode("utf-8", errors="replace")
        truncated = len(output_bytes) > _MAX_OUTPUT
        if truncated:
            output = output_bytes[:_MAX_OUTPUT].decode("utf-8", errors="ignore")
        return ExecuteResponse(
            output=output,
            exit_code=result.exit_code,
            truncated=truncated,
        )

    def _resolve_params(
        self,
        timeout: int | None,
        env_vars: dict[str, str] | None,
        metadata: dict[str, str] | None,
        sub_paths: list[str] | None,
    ) -> tuple[int, dict[str, str] | None, dict[str, str] | None, list[str] | None]:
        """Resolve per-call overrides.

        None (default) uses the constructor value. A non-None value fully
        replaces the constructor default for that call (no merging).
        """
        return (
            timeout if timeout is not None else self._timeout,
            env_vars if env_vars is not None else self._env_vars,
            metadata if metadata is not None else self._metadata,
            sub_paths if sub_paths is not None else self._sub_paths,
        )

    def execute(
        self,
        command: str,
        *,
        timeout: int | None = None,
        env_vars: dict[str, str] | None = None,
        metadata: dict[str, str] | None = None,
        sub_paths: list[str] | None = None,
    ) -> ExecuteResponse:
        client = self._get_sync_client()
        eff_timeout, eff_env, eff_meta, eff_sub = self._resolve_params(
            timeout, env_vars, metadata, sub_paths,
        )

        for attempt in range(len(_RETRY_DELAYS) + 1):
            try:
                result = client.execute(
                    command,
                    language="bash",
                    thread_id=self._thread_id,
                    timeout=eff_timeout,
                    metadata=eff_meta,
                    env_vars=eff_env,
                    sub_paths=eff_sub,
                )
                return self._to_execute_response(result)
            except ThreadBusyError:
                if attempt >= len(_RETRY_DELAYS):
                    raise
                time.sleep(_RETRY_DELAYS[attempt])

        raise RuntimeError("unreachable")  # pragma: no cover

    async def aexecute(
        self,
        command: str,
        *,
        timeout: int | None = None,
        env_vars: dict[str, str] | None = None,
        metadata: dict[str, str] | None = None,
        sub_paths: list[str] | None = None,
    ) -> ExecuteResponse:
        client = self._get_async_client()
        eff_timeout, eff_env, eff_meta, eff_sub = self._resolve_params(
            timeout, env_vars, metadata, sub_paths,
        )

        for attempt in range(len(_RETRY_DELAYS) + 1):
            try:
                result = await client.execute(
                    command,
                    language="bash",
                    thread_id=self._thread_id,
                    timeout=eff_timeout,
                    metadata=eff_meta,
                    env_vars=eff_env,
                    sub_paths=eff_sub,
                )
                return self._to_execute_response(result)
            except ThreadBusyError:
                if attempt >= len(_RETRY_DELAYS):
                    raise
                await asyncio.sleep(_RETRY_DELAYS[attempt])

        raise RuntimeError("unreachable")  # pragma: no cover

    def execute_stream(
        self,
        command: str,
        *,
        timeout: int | None = None,
        env_vars: dict[str, str] | None = None,
        metadata: dict[str, str] | None = None,
        sub_paths: list[str] | None = None,
    ) -> ExecutionStream:
        """Execute with real-time NDJSON streaming (Baponi-specific, not part of ABC).

        Returns an ExecutionStream that yields typed events as they arrive.
        Retries on ThreadBusyError (409) like execute()/aexecute() since the
        rejection arrives before the stream opens.
        """
        client = self._get_sync_client()
        eff_timeout, eff_env, eff_meta, eff_sub = self._resolve_params(
            timeout, env_vars, metadata, sub_paths,
        )

        for attempt in range(len(_RETRY_DELAYS) + 1):
            try:
                return client.execute_stream(
                    command,
                    language="bash",
                    thread_id=self._thread_id,
                    timeout=eff_timeout,
                    metadata=eff_meta,
                    env_vars=eff_env,
                    sub_paths=eff_sub,
                )
            except ThreadBusyError:
                if attempt >= len(_RETRY_DELAYS):
                    raise
                time.sleep(_RETRY_DELAYS[attempt])

        raise RuntimeError("unreachable")  # pragma: no cover

    def close(self) -> None:
        """Close underlying HTTP clients, releasing connections.

        Closes the sync client directly. If an async client exists,
        logs a warning since it cannot be closed synchronously.
        Use aclose() if you used async methods.
        """
        if self._async_client is not None:
            import warnings

            warnings.warn(
                "BaponiSandbox.close() cannot close the async HTTP client. "
                "Use await aclose() instead to avoid leaking connections.",
                ResourceWarning,
                stacklevel=2,
            )
            self._async_client = None
        if self._sync_client is not None:
            self._sync_client.close()

    async def aclose(self) -> None:
        """Async close -closes both async and sync clients if present."""
        if self._async_client is not None:
            await self._async_client.close()
        if self._sync_client is not None:
            self._sync_client.close()

    def __enter__(self) -> Self:
        return self

    def __exit__(self, *args: object) -> None:
        self.close()

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(self, *args: object) -> None:
        await self.aclose()

    def upload_files(
        self, files: list[tuple[str, bytes]]
    ) -> list[FileUploadResponse]:
        client = self._get_sync_client()
        responses: list[FileUploadResponse] = []
        for path, content in files:
            try:
                api_path = _to_relative_path(path)
                signed = client.upload_url(
                    path=api_path,
                    content_length=len(content),
                    source="thread",
                    id=self._thread_id,
                )
                headers = {h.name: h.value for h in signed.headers}
                resp = httpx.put(signed.url, content=content, headers=headers)
                if resp.is_success:
                    responses.append(FileUploadResponse(path=path, error=None))
                else:
                    responses.append(FileUploadResponse(path=path, error="permission_denied"))
            except Exception:
                responses.append(FileUploadResponse(path=path, error="invalid_path"))
        return responses

    def download_files(
        self, paths: list[str]
    ) -> list[FileDownloadResponse]:
        client = self._get_sync_client()
        responses: list[FileDownloadResponse] = []
        for path in paths:
            try:
                api_path = _to_relative_path(path)
                signed = client.download_url(
                    path=api_path,
                    source="thread",
                    id=self._thread_id,
                )
                headers = {h.name: h.value for h in signed.headers}
                resp = httpx.get(signed.url, headers=headers)
                if resp.is_success:
                    responses.append(
                        FileDownloadResponse(path=path, content=resp.content, error=None)
                    )
                else:
                    responses.append(
                        FileDownloadResponse(path=path, content=None, error="file_not_found")
                    )
            except Exception:
                responses.append(
                    FileDownloadResponse(path=path, content=None, error="file_not_found")
                )
        return responses


def create_sandbox(
    *,
    api_key: str | None = None,
    base_url: str | None = None,
    thread_id: str | None = None,
    timeout: int = DEFAULT_EXECUTION_TIMEOUT,
    metadata: dict[str, str] | None = None,
    env_vars: dict[str, str] | None = None,
    sub_paths: list[str] | None = None,
) -> BaponiSandbox:
    """Create a configured BaponiSandbox instance.

    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: Thread ID for persistent state. Auto-generated if not provided.
        timeout: Default execution timeout in seconds.
        metadata: Default metadata merged into every execution call.
        env_vars: Default environment variables for every execution call.
        sub_paths: Default sub-paths for every execution call.

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


class BaponiSandboxProvider(_ProviderBase):  # type: ignore[misc]
    """Programmatic lifecycle manager for BaponiSandbox instances.

    Tracks sandboxes by ID. ``get_or_create()`` returns an existing sandbox
    or creates a new one.  ``delete()`` closes and removes.

    Not thread-safe. Intended for single-threaded or async usage. If you
    need concurrent access from multiple OS threads, synchronize externally.

    This class implements the ``SandboxProvider`` interface from
    ``deepagents_cli`` when available, but works standalone regardless.

    Note: This does NOT enable ``deepagents-cli --sandbox baponi``.
    That requires an upstream PR to ``langchain-ai/deepagents``.
    """

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: int = DEFAULT_EXECUTION_TIMEOUT,
        metadata: dict[str, str] | None = None,
        env_vars: dict[str, str] | None = None,
        sub_paths: list[str] | None = None,
    ) -> None:
        self._api_key = api_key or os.environ.get("BAPONI_API_KEY")
        self._base_url = base_url or os.environ.get("BAPONI_BASE_URL")
        self._timeout = timeout
        self._metadata = metadata
        self._env_vars = env_vars
        self._sub_paths = sub_paths
        self._sandboxes: dict[str, BaponiSandbox] = {}

    def get_or_create(
        self, *, sandbox_id: str | None = None, **kwargs: Any
    ) -> BaponiSandbox:
        """Return existing sandbox by ID, or create a new one.

        Args:
            sandbox_id: Explicit sandbox/thread ID. Auto-generated if None.
        """
        thread_id = sandbox_id or f"deep-{uuid.uuid4().hex[:16]}"
        if thread_id in self._sandboxes:
            return self._sandboxes[thread_id]

        sandbox = BaponiSandbox(
            api_key=self._api_key,
            base_url=self._base_url,
            thread_id=thread_id,
            timeout=self._timeout,
            metadata=self._metadata,
            env_vars=self._env_vars,
            sub_paths=self._sub_paths,
        )
        self._sandboxes[thread_id] = sandbox
        return sandbox

    async def aget_or_create(
        self, *, sandbox_id: str | None = None, **kwargs: Any
    ) -> BaponiSandbox:
        """Async variant -same logic, sandboxes are shared across sync/async."""
        return self.get_or_create(sandbox_id=sandbox_id)

    def delete(self, *, sandbox_id: str, **kwargs: Any) -> None:
        """Close and remove a sandbox by ID. No-op if not found."""
        sandbox = self._sandboxes.pop(sandbox_id, None)
        if sandbox is not None:
            sandbox.close()

    async def adelete(self, *, sandbox_id: str, **kwargs: Any) -> None:
        """Async close and remove a sandbox by ID. No-op if not found."""
        sandbox = self._sandboxes.pop(sandbox_id, None)
        if sandbox is not None:
            await sandbox.aclose()
