"""Sandbox class for interacting with a specific sandbox instance."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, overload

import httpx

from langsmith.sandbox._exceptions import (
    DataplaneNotConfiguredError,
    ResourceNotFoundError,
    SandboxConnectionError,
    SandboxNotReadyError,
)
from langsmith.sandbox._helpers import handle_sandbox_http_error
from langsmith.sandbox._models import (
    CommandHandle,
    ExecutionResult,
)

if TYPE_CHECKING:
    from langsmith.sandbox._client import SandboxClient


@dataclass
class Sandbox:
    """Represents an active sandbox for running commands and file operations.

    This class is typically obtained from SandboxClient.sandbox() and supports
    the context manager protocol for automatic cleanup.

    Attributes:
        name: Display name (can be updated).
        template_name: Name of the template used to create this sandbox.
        dataplane_url: URL for data plane operations (file I/O, command execution).
            Only functional when status is "ready".
        id: Unique identifier (UUID). Remains constant even if name changes.
            May be None for resources created before ID support was added.
        status: Sandbox lifecycle status. One of "provisioning", "ready", "failed".
        status_message: Human-readable details when status is "failed", None otherwise.
        created_at: Timestamp when the sandbox was created.
        updated_at: Timestamp when the sandbox was last updated.

    Example:
        with client.sandbox(template_name="python-sandbox") as sandbox:
            result = sandbox.run("python --version")
            print(result.stdout)
    """

    # Data fields (from API response)
    name: str
    template_name: str
    dataplane_url: Optional[str] = None
    id: Optional[str] = None
    status: str = "ready"
    status_message: Optional[str] = None
    created_at: Optional[str] = None
    updated_at: Optional[str] = None

    # Internal fields (not from API)
    _client: SandboxClient = field(repr=False, default=None)  # type: ignore
    _auto_delete: bool = field(repr=False, default=True)

    @classmethod
    def from_dict(
        cls,
        data: dict[str, Any],
        client: SandboxClient,
        auto_delete: bool = True,
    ) -> Sandbox:
        """Create a Sandbox from API response dict.

        Args:
            data: API response dictionary containing sandbox data.
            client: Parent SandboxClient for operations.
            auto_delete: Whether to delete the sandbox on context exit.

        Returns:
            Sandbox instance.
        """
        return cls(
            name=data.get("name", ""),
            template_name=data.get("template_name", ""),
            dataplane_url=data.get("dataplane_url"),
            id=data.get("id"),
            status=data.get("status", "ready"),
            status_message=data.get("status_message"),
            created_at=data.get("created_at"),
            updated_at=data.get("updated_at"),
            _client=client,
            _auto_delete=auto_delete,
        )

    def __enter__(self) -> Sandbox:
        """Enter context manager."""
        return self

    def __exit__(
        self,
        exc_type: Optional[type],
        exc_val: Optional[BaseException],
        exc_tb: Optional[Any],
    ) -> None:
        """Exit context manager, optionally deleting the sandbox."""
        if self._auto_delete:
            try:
                self._client.delete_sandbox(self.name)
            except Exception:
                # Don't raise on cleanup errors
                pass

    def _require_dataplane_url(self) -> str:
        """Validate and return the dataplane URL.

        Returns:
            The dataplane URL.

        Raises:
            SandboxNotReadyError: If sandbox status is not "ready".
            DataplaneNotConfiguredError: If dataplane_url is not configured.
        """
        if self.status != "ready":
            raise SandboxNotReadyError(
                f"Sandbox '{self.name}' is not ready (status: {self.status}). "
                "Wait for status 'ready' before running operations."
            )
        if not self.dataplane_url:
            raise DataplaneNotConfiguredError(
                f"Sandbox '{self.name}' does not have a dataplane_url configured. "
                "Runtime operations require a dataplane URL."
            )
        return self.dataplane_url

    @overload
    def run(
        self,
        command: str,
        *,
        timeout: int = ...,
        env: Optional[dict[str, str]] = ...,
        cwd: Optional[str] = ...,
        shell: str = ...,
        on_stdout: Optional[Callable[[str], Any]] = ...,
        on_stderr: Optional[Callable[[str], Any]] = ...,
        wait: Literal[True] = ...,
    ) -> ExecutionResult: ...

    @overload
    def run(
        self,
        command: str,
        *,
        timeout: int = ...,
        env: Optional[dict[str, str]] = ...,
        cwd: Optional[str] = ...,
        shell: str = ...,
        on_stdout: Optional[Callable[[str], Any]] = ...,
        on_stderr: Optional[Callable[[str], Any]] = ...,
        wait: Literal[False],
    ) -> CommandHandle: ...

    def run(
        self,
        command: str,
        *,
        timeout: int = 60,
        env: Optional[dict[str, str]] = None,
        cwd: Optional[str] = None,
        shell: str = "/bin/bash",
        on_stdout: Optional[Callable[[str], Any]] = None,
        on_stderr: Optional[Callable[[str], Any]] = None,
        wait: bool = True,
    ) -> Union[ExecutionResult, CommandHandle]:
        """Execute a command in the sandbox.

        Args:
            command: Shell command to execute.
            timeout: Command timeout in seconds.
            env: Environment variables to set for the command.
            cwd: Working directory for command execution. If None, uses sandbox default.
            shell: Shell to use for command execution. Defaults to "/bin/bash".
            on_stdout: Callback invoked with each stdout chunk as it arrives.
                Blocks until the command completes and returns ExecutionResult.
                Cannot be combined with wait=False.
            on_stderr: Callback invoked with each stderr chunk as it arrives.
                Blocks until the command completes and returns ExecutionResult.
                Cannot be combined with wait=False.
            wait: If True (default), block until the command completes and
                return ExecutionResult. If False, return a
                CommandHandle immediately for streaming output,
                kill, stdin input, and reconnection. Cannot be combined with
                on_stdout/on_stderr callbacks.

        Returns:
            ExecutionResult when wait=True (default).
            CommandHandle when wait=False.

        Raises:
            ValueError: If wait=False is combined with callbacks.
            DataplaneNotConfiguredError: If dataplane_url is not configured.
            SandboxOperationError: If command execution fails.
            CommandTimeoutError: If command exceeds its timeout.
            SandboxConnectionError: If connection to sandbox fails after retries.
            SandboxNotReadyError: If sandbox is not ready.
            SandboxClientError: For other errors.
        """
        if not wait and (on_stdout or on_stderr):
            raise ValueError(
                "Cannot combine wait=False with on_stdout/on_stderr callbacks. "
                "Use wait=False and iterate the CommandHandle, or use callbacks."
            )

        self._require_dataplane_url()

        # When not waiting or callbacks are requested, WS is required
        use_ws = not wait or on_stdout or on_stderr
        if use_ws:
            return self._run_ws(
                command,
                timeout=timeout,
                env=env,
                cwd=cwd,
                shell=shell,
                wait=wait,
                on_stdout=on_stdout,
                on_stderr=on_stderr,
            )

        # Default (wait=True, no callbacks): try WS, fall back to HTTP.
        # Catch broad exceptions so that unexpected WS failures (e.g. version
        # incompatibilities) don't break users who don't need WS features.
        try:
            return self._run_ws(
                command,
                timeout=timeout,
                env=env,
                cwd=cwd,
                shell=shell,
                wait=True,
                on_stdout=None,
                on_stderr=None,
            )
        except (SandboxConnectionError, ImportError, OSError, TypeError):
            return self._run_http(
                command,
                timeout=timeout,
                env=env,
                cwd=cwd,
                shell=shell,
            )

    def _run_ws(
        self,
        command: str,
        *,
        timeout: int,
        env: Optional[dict[str, str]],
        cwd: Optional[str],
        shell: str,
        wait: bool,
        on_stdout: Optional[Callable[[str], Any]],
        on_stderr: Optional[Callable[[str], Any]],
    ) -> Union[ExecutionResult, CommandHandle]:
        """Execute via WebSocket /execute/ws."""
        from langsmith.sandbox._ws_execute import run_ws_stream

        dataplane_url = self._require_dataplane_url()
        api_key = self._client._api_key

        msg_stream, control = run_ws_stream(
            dataplane_url,
            api_key,
            command,
            timeout=timeout,
            env=env,
            cwd=cwd,
            shell=shell,
            on_stdout=on_stdout,
            on_stderr=on_stderr,
        )

        handle = CommandHandle(msg_stream, control, self)

        if not wait:
            return handle

        return handle.result  # blocks until command completes

    def _run_http(
        self,
        command: str,
        *,
        timeout: int,
        env: Optional[dict[str, str]],
        cwd: Optional[str],
        shell: str,
    ) -> ExecutionResult:
        """Execute via HTTP POST /execute (existing implementation)."""
        dataplane_url = self._require_dataplane_url()
        url = f"{dataplane_url}/execute"
        payload: dict[str, Any] = {
            "command": command,
            "timeout": timeout,
            "shell": shell,
        }
        if env is not None:
            payload["env"] = env
        if cwd is not None:
            payload["cwd"] = cwd

        try:
            response = self._client._http.post(url, json=payload, timeout=timeout + 10)
            response.raise_for_status()
            data = response.json()
            return ExecutionResult(
                stdout=data.get("stdout", ""),
                stderr=data.get("stderr", ""),
                exit_code=data.get("exit_code", -1),
            )
        except httpx.HTTPStatusError as e:
            handle_sandbox_http_error(e)
            raise  # pragma: no cover

    def reconnect(
        self,
        command_id: str,
        *,
        stdout_offset: int = 0,
        stderr_offset: int = 0,
    ) -> CommandHandle:
        """Reconnect to a running or recently-finished command.

        Resumes output from the given byte offsets. Any output produced while
        the client was disconnected is replayed from the server's ring buffer.

        Args:
            command_id: The command ID from handle.command_id.
            stdout_offset: Byte offset to resume stdout from (default: 0).
            stderr_offset: Byte offset to resume stderr from (default: 0).

        Returns:
            A CommandHandle for the command.

        Raises:
            SandboxOperationError: If command_id is not found or session expired.
            SandboxConnectionError: If connection to sandbox fails after retries.
        """
        from langsmith.sandbox._ws_execute import reconnect_ws_stream

        dataplane_url = self._require_dataplane_url()
        api_key = self._client._api_key

        msg_stream, control = reconnect_ws_stream(
            dataplane_url,
            api_key,
            command_id,
            stdout_offset=stdout_offset,
            stderr_offset=stderr_offset,
        )

        return CommandHandle(
            msg_stream,
            control,
            self,
            command_id=command_id,
            stdout_offset=stdout_offset,
            stderr_offset=stderr_offset,
        )

    def write(
        self,
        path: str,
        content: Union[str, bytes],
        *,
        timeout: int = 60,
    ) -> None:
        """Write content to a file in the sandbox.

        Args:
            path: Target file path in the sandbox.
            content: File content (str or bytes).
            timeout: Request timeout in seconds.

        Raises:
            DataplaneNotConfiguredError: If dataplane_url is not configured.
            SandboxOperationError: If file write fails.
            SandboxConnectionError: If connection to sandbox fails after retries.
            SandboxNotReadyError: If sandbox is not ready.
            SandboxClientError: For other errors.
        """
        dataplane_url = self._require_dataplane_url()
        url = f"{dataplane_url}/upload"

        # Ensure content is bytes for multipart upload
        if isinstance(content, str):
            content = content.encode("utf-8")

        files = {"file": ("file", content)}

        try:
            response = self._client._http.post(
                url, params={"path": path}, files=files, timeout=timeout
            )
            response.raise_for_status()
        except httpx.HTTPStatusError as e:
            handle_sandbox_http_error(e)

    def read(self, path: str, *, timeout: int = 60) -> bytes:
        """Read a file from the sandbox.

        Args:
            path: File path to read. Supports both absolute paths (e.g., /tmp/file.txt)
                  and relative paths (resolved from /home/user/).
            timeout: Request timeout in seconds.

        Returns:
            File contents as bytes.

        Raises:
            DataplaneNotConfiguredError: If dataplane_url is not configured.
            ResourceNotFoundError: If the file doesn't exist.
            SandboxOperationError: If file read fails.
            SandboxConnectionError: If connection to sandbox fails after retries.
            SandboxNotReadyError: If sandbox is not ready.
            SandboxClientError: For other errors.
        """
        dataplane_url = self._require_dataplane_url()
        url = f"{dataplane_url}/download"

        try:
            response = self._client._http.get(
                url, params={"path": path}, timeout=timeout
            )
            response.raise_for_status()
            return response.content
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 404:
                raise ResourceNotFoundError(
                    f"File '{path}' not found in sandbox '{self.name}'",
                    resource_type="file",
                ) from e
            handle_sandbox_http_error(e)
            # This line should never be reached but satisfies type checker
            raise  # pragma: no cover
