"""Data models for the sandbox client."""

from __future__ import annotations

import logging
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Literal, Optional

import httpx

from langsmith.sandbox._exceptions import (
    SandboxConnectionError,
    SandboxOperationError,
    SandboxServerReloadError,
)

if TYPE_CHECKING:
    from langsmith.sandbox._async_sandbox import AsyncSandbox
    from langsmith.sandbox._sandbox import Sandbox
    from langsmith.sandbox._ws_execute import (
        _AsyncWSStreamControl,
        _WSStreamControl,
    )

logger = logging.getLogger(__name__)


def _acknowledges_reconnect(msg: dict, command_id: Optional[str]) -> bool:
    """Whether a 'started' frame proves *this* command was reattached to.

    A command and its WebSocket are separate things: the command keeps running on
    the server and the socket is only this client's attachment to it. The server
    acknowledges every successful reattachment with 'started', which for a
    command that produces no output is the only evidence it landed. One naming a
    different command is not evidence, and must not clear the reconnect budget.
    """
    acked = msg.get("command_id")
    if acked == command_id:
        return True
    logger.warning(
        "Ignoring reconnect acknowledgement for command %r while attached to %r",
        acked,
        command_id,
    )
    return False


class _StreamEndedBeforeStarted(SandboxOperationError):
    """A command WebSocket closed before the guest sent its 'started' frame.

    Internal marker for the idempotently-retryable early close (the proxied
    tunnel was torn down gracefully mid-handshake), as distinct from a
    command-level failure. Subclasses SandboxOperationError so callers that
    catch the public type are unaffected.
    """


@dataclass
class ExecutionResult:
    """Result of executing a command in a sandbox."""

    stdout: str
    stderr: str
    exit_code: int

    @property
    def success(self) -> bool:
        """Return True if the command exited with code 0."""
        return self.exit_code == 0


@dataclass
class ResourceStatus:
    """Lightweight provisioning status for any async-created resource.

    Attributes:
        status: Resource lifecycle status. One of "provisioning", "ready", "failed".
        status_message: Human-readable details when status is "failed", None otherwise.
    """

    status: str
    status_message: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> ResourceStatus:
        """Create a ResourceStatus from API response dict."""
        return cls(
            status=data.get("status", "provisioning"),
            status_message=data.get("status_message"),
        )


@dataclass
class Snapshot:
    """Represents a sandbox snapshot.

    Snapshots are built from Docker images or captured from running sandboxes.
    They are used to create new sandboxes.

    Attributes:
        id: Unique identifier (UUID).
        name: Display name.
        status: Build status. One of "building", "ready", "failed".
        fs_capacity_bytes: Filesystem capacity in bytes.
        docker_image: Source Docker image (for build snapshots).
        image_digest: Docker image digest after pull.
        source_sandbox_id: Source sandbox (for capture snapshots).
        status_message: Human-readable details when status is "failed".
        fs_used_bytes: Actual bytes used on the filesystem.
        created_by: User or service that created the snapshot.
        registry_id: Private registry ID, if applicable.
        created_at: Timestamp when the snapshot was created.
        updated_at: Timestamp when the snapshot was last updated.
    """

    id: str
    name: str
    status: str
    fs_capacity_bytes: int
    docker_image: Optional[str] = None
    image_digest: Optional[str] = None
    source_sandbox_id: Optional[str] = None
    status_message: Optional[str] = None
    fs_used_bytes: Optional[int] = None
    created_by: Optional[str] = None
    registry_id: Optional[str] = None
    created_at: Optional[str] = None
    updated_at: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Snapshot:
        """Create a Snapshot from API response dict."""
        return cls(
            id=data.get("id", ""),
            name=data.get("name", ""),
            status=data.get("status", "building"),
            fs_capacity_bytes=data.get("fs_capacity_bytes", 0),
            docker_image=data.get("docker_image"),
            image_digest=data.get("image_digest"),
            source_sandbox_id=data.get("source_sandbox_id"),
            status_message=data.get("status_message"),
            fs_used_bytes=data.get("fs_used_bytes"),
            created_by=data.get("created_by"),
            registry_id=data.get("registry_id"),
            created_at=data.get("created_at"),
            updated_at=data.get("updated_at"),
        )


# =============================================================================
# Service URL Models
# =============================================================================

_AUTH_HEADER = "X-Langsmith-Sandbox-Service-Token"
_REFRESH_MARGIN_SECONDS = 30


class ServiceURL:
    """Authenticated URL for accessing an HTTP service running in a sandbox.

    Properties auto-refresh the token transparently when it nears expiry.
    HTTP helper methods (``.get``, ``.post``, etc.) inject the auth header
    automatically.

    When constructed by :meth:`SandboxClient.service` or
    :meth:`Sandbox.service`, the object holds an internal refresher that
    re-calls the API to obtain a fresh token before the current one expires.

    Example::

        svc = sb.service(port=3000)

        resp = svc.get("/api/data")  # token injected + auto-refreshed
        print(svc.browser_url)  # always-fresh URL
    """

    def __init__(
        self,
        browser_url: str,
        service_url: str,
        token: str,
        expires_at: str,
        *,
        _refresher: Optional[Callable[[], ServiceURL]] = None,
    ) -> None:
        self._browser_url = browser_url
        self._service_url = service_url
        self._token = token
        self._expires_at = expires_at
        self._refresher = _refresher

    # -- Auto-refresh logic -------------------------------------------------

    def _should_refresh(self) -> bool:
        if self._refresher is None:
            return False
        raw = self._expires_at.replace("Z", "+00:00")
        expires = datetime.fromisoformat(raw)
        if expires.tzinfo is None:
            expires = expires.replace(tzinfo=timezone.utc)
        remaining = (expires - datetime.now(timezone.utc)).total_seconds()
        return remaining <= _REFRESH_MARGIN_SECONDS

    def _maybe_refresh(self) -> None:
        if self._should_refresh():
            fresh = self._refresher()  # type: ignore[misc]
            self._browser_url = fresh._browser_url
            self._service_url = fresh._service_url
            self._token = fresh._token
            self._expires_at = fresh._expires_at

    # -- Properties (auto-refresh on access) --------------------------------

    @property
    def token(self) -> str:
        """Return the raw JWT, refreshing if near expiry."""
        self._maybe_refresh()
        return self._token

    @property
    def service_url(self) -> str:
        """Return the base URL, refreshing if near expiry."""
        self._maybe_refresh()
        return self._service_url

    @property
    def browser_url(self) -> str:
        """Return the browser auth URL, refreshing if near expiry."""
        self._maybe_refresh()
        return self._browser_url

    @property
    def expires_at(self) -> str:
        """Return the ISO 8601 expiration, refreshing if near expiry."""
        self._maybe_refresh()
        return self._expires_at

    # -- HTTP helpers (stateless, one httpx call per request) ----------------

    def request(self, method: str, path: str = "/", **kwargs: Any) -> httpx.Response:
        """Make an HTTP request to the service, injecting the auth header.

        Args:
            method: HTTP method (GET, POST, etc.).
            path: Path relative to the service URL.
            **kwargs: Forwarded to ``httpx.request``.

        Returns:
            httpx.Response.
        """
        url = self.service_url.rstrip("/") + "/" + path.lstrip("/")
        headers = dict(kwargs.pop("headers", None) or {})
        headers[_AUTH_HEADER] = self.token
        return httpx.request(method, url, headers=headers, **kwargs)

    def get(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """HTTP GET to the service."""
        return self.request("GET", path, **kwargs)

    def post(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """HTTP POST to the service."""
        return self.request("POST", path, **kwargs)

    def put(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """HTTP PUT to the service."""
        return self.request("PUT", path, **kwargs)

    def patch(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """HTTP PATCH to the service."""
        return self.request("PATCH", path, **kwargs)

    def delete(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """HTTP DELETE to the service."""
        return self.request("DELETE", path, **kwargs)

    # -- Construction -------------------------------------------------------

    @classmethod
    def from_dict(
        cls,
        data: dict[str, Any],
        *,
        _refresher: Optional[Callable[[], ServiceURL]] = None,
    ) -> ServiceURL:
        """Create a ServiceURL from API response dict."""
        return cls(
            browser_url=data["browser_url"],
            service_url=data["service_url"],
            token=data["token"],
            expires_at=data["expires_at"],
            _refresher=_refresher,
        )

    def __repr__(self) -> str:
        return (
            f"ServiceURL(service_url={self._service_url!r}, "
            f"expires_at={self._expires_at!r})"
        )


class AsyncServiceURL:
    """Async variant of :class:`ServiceURL` with auto-refreshing token.

    Properties and HTTP helpers are async. Use with
    :meth:`AsyncSandboxClient.service` or :meth:`AsyncSandbox.service`.

    Example::

        svc = await sb.service(port=3000)

        resp = await svc.get("/api/data")
        print(await svc.get_browser_url())
    """

    def __init__(
        self,
        browser_url: str,
        service_url: str,
        token: str,
        expires_at: str,
        *,
        _refresher: Optional[Callable[[], Awaitable[AsyncServiceURL]]] = None,
    ) -> None:
        self._browser_url = browser_url
        self._service_url = service_url
        self._token = token
        self._expires_at = expires_at
        self._refresher = _refresher

    # -- Auto-refresh logic -------------------------------------------------

    def _should_refresh(self) -> bool:
        if self._refresher is None:
            return False
        raw = self._expires_at.replace("Z", "+00:00")
        expires = datetime.fromisoformat(raw)
        if expires.tzinfo is None:
            expires = expires.replace(tzinfo=timezone.utc)
        remaining = (expires - datetime.now(timezone.utc)).total_seconds()
        return remaining <= _REFRESH_MARGIN_SECONDS

    async def _maybe_refresh(self) -> None:
        if self._should_refresh():
            fresh = await self._refresher()  # type: ignore[misc]
            self._browser_url = fresh._browser_url
            self._service_url = fresh._service_url
            self._token = fresh._token
            self._expires_at = fresh._expires_at

    # -- Async accessors (auto-refresh on access) ---------------------------

    async def get_token(self) -> str:
        """Return the raw JWT, refreshing if near expiry."""
        await self._maybe_refresh()
        return self._token

    async def get_service_url(self) -> str:
        """Return the base URL, refreshing if near expiry."""
        await self._maybe_refresh()
        return self._service_url

    async def get_browser_url(self) -> str:
        """Return the browser auth URL, refreshing if near expiry."""
        await self._maybe_refresh()
        return self._browser_url

    async def get_expires_at(self) -> str:
        """Return the ISO 8601 expiration, refreshing if near expiry."""
        await self._maybe_refresh()
        return self._expires_at

    # -- Sync property access (no refresh, use when token is known-fresh) ---

    @property
    def token(self) -> str:
        """Return the raw JWT without refreshing."""
        return self._token

    @property
    def service_url(self) -> str:
        """Return the base URL without refreshing."""
        return self._service_url

    @property
    def browser_url(self) -> str:
        """Return the browser auth URL without refreshing."""
        return self._browser_url

    @property
    def expires_at(self) -> str:
        """Return the expiration timestamp without refreshing."""
        return self._expires_at

    # -- HTTP helpers (one request per call) --------------------------------

    async def request(
        self, method: str, path: str = "/", **kwargs: Any
    ) -> httpx.Response:
        """Make an async HTTP request to the service, injecting the auth header.

        Args:
            method: HTTP method (GET, POST, etc.).
            path: Path relative to the service URL.
            **kwargs: Forwarded to ``httpx.AsyncClient.request``.

        Returns:
            httpx.Response.
        """
        url = (await self.get_service_url()).rstrip("/") + "/" + path.lstrip("/")
        headers = dict(kwargs.pop("headers", None) or {})
        headers[_AUTH_HEADER] = await self.get_token()
        async with httpx.AsyncClient() as client:
            return await client.request(method, url, headers=headers, **kwargs)

    async def get(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """Async HTTP GET to the service."""
        return await self.request("GET", path, **kwargs)

    async def post(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """Async HTTP POST to the service."""
        return await self.request("POST", path, **kwargs)

    async def put(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """Async HTTP PUT to the service."""
        return await self.request("PUT", path, **kwargs)

    async def patch(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """Async HTTP PATCH to the service."""
        return await self.request("PATCH", path, **kwargs)

    async def delete(self, path: str = "/", **kwargs: Any) -> httpx.Response:
        """Async HTTP DELETE to the service."""
        return await self.request("DELETE", path, **kwargs)

    # -- Construction -------------------------------------------------------

    @classmethod
    def from_dict(
        cls,
        data: dict[str, Any],
        *,
        _refresher: Optional[Callable[[], Awaitable[AsyncServiceURL]]] = None,
    ) -> AsyncServiceURL:
        """Create an AsyncServiceURL from API response dict."""
        return cls(
            browser_url=data["browser_url"],
            service_url=data["service_url"],
            token=data["token"],
            expires_at=data["expires_at"],
            _refresher=_refresher,
        )

    def __repr__(self) -> str:
        return (
            f"AsyncServiceURL(service_url={self._service_url!r}, "
            f"expires_at={self._expires_at!r})"
        )


# =============================================================================
# Download URL Models
# =============================================================================


DownloadContentDisposition = Literal["attachment", "inline"]
"""How a download link asks the browser to handle the file."""


@dataclass
class DownloadURL:
    """A link that downloads one sandbox file with no LangSmith credential.

    The link is pinned to the sandbox, the file path, and the response
    headers, so it cannot be repointed at another file. It is pinned to the
    path rather than to a snapshot of the contents, so the file must not be
    modified while the link is in use.

    Attributes:
        download_url: The full URL to fetch. Supports GET, HEAD, and Range.
        token: The signed token embedded in ``download_url``.
        expires_at: Expiry timestamp, or None for a link that never expires.
    """

    download_url: str
    token: str
    expires_at: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> DownloadURL:
        """Create a DownloadURL from API response dict."""
        return cls(
            download_url=data.get("download_url", ""),
            token=data.get("token", ""),
            expires_at=data.get("expires_at"),
        )


# =============================================================================
# WebSocket Command Execution Models
# =============================================================================


@dataclass
class OutputChunk:
    """A single chunk of streaming output from command execution.

    Attributes:
        stream: Either "stdout" or "stderr".
        data: The text content of this chunk (valid UTF-8, server handles
            boundary splitting).
        offset: Byte offset within the stream. Used internally for
            reconnection; users typically don't need this.
    """

    stream: str
    data: str
    offset: int


class CommandHandle:
    """Handle to a running command with streaming output and auto-reconnect.

    Iterable, yielding OutputChunk objects (stdout and stderr interleaved
    in arrival order). Access .result after iteration to get the full
    ExecutionResult.

    Auto-reconnect behavior:
    - Server hot-reload (1001 Going Away): reconnect immediately
    - Network error / unexpected close:    reconnect with exponential backoff
    - User called kill():                  do NOT reconnect (propagate error)

    The auto-reconnect is transparent -- the iterator reconnects and
    continues yielding chunks without any user intervention. If all
    reconnect attempts are exhausted, SandboxConnectionError is raised.

    Construction modes (controlled by ``command_id``):
    - **New execution** (``command_id=""``, the default): the constructor
      eagerly reads the server's ``"started"`` message to populate
      ``command_id`` and ``pid`` before returning.
    - **Reconnection** (``command_id`` set): skips the started-message
      read. A reconnect stream's ``"started"`` message is the server's
      acknowledgement that the reattachment landed, consumed while iterating to
      clear the reconnect budget — the only such signal for a command that
      emits no output.

    Example:
        handle = sandbox.run("make build", timeout=600, wait=False)

        for chunk in handle:          # auto-reconnects on transient errors
            print(chunk.data, end="")

        result = handle.result
        print(f"Exit code: {result.exit_code}")
    """

    MAX_AUTO_RECONNECTS = 5
    _BACKOFF_BASE = 0.5  # seconds
    _BACKOFF_MAX = 8.0  # seconds

    def __init__(
        self,
        message_stream: Iterator[dict],
        control: Optional[_WSStreamControl],
        sandbox: Sandbox,
        *,
        command_id: str = "",
        stdout_offset: int = 0,
        stderr_offset: int = 0,
        on_stdout: Optional[Callable[[str], Any]] = None,
        on_stderr: Optional[Callable[[str], Any]] = None,
    ) -> None:
        self._stream = message_stream
        self._control = control
        self._sandbox = sandbox
        self._on_stdout = on_stdout
        self._on_stderr = on_stderr
        self._command_id: Optional[str] = None
        self._pid: Optional[int] = None
        self._result: Optional[ExecutionResult] = None
        self._stdout_parts: list[str] = []
        self._stderr_parts: list[str] = []
        self._exhausted = False
        self._last_stdout_offset = stdout_offset
        self._last_stderr_offset = stderr_offset
        self._reconnect_attempts = 0

        # New executions (command_id=""): eager_start reads "started" message.
        # Reconnections (command_id set): the "started" message is the server's
        # reattachment acknowledgement, consumed by _iter_stream instead.
        if command_id:
            self._command_id = command_id
        else:
            self._consume_started()

    def _consume_started(self) -> None:
        """Eagerly read the 'started' message to populate command_id and pid.

        Blocks briefly until the server sends the started message (arrives
        near-instantly after connection). After this call, command_id and
        pid are available, and the WebSocket is bound to the control object
        (so kill() works).
        """
        try:
            first_msg = next(self._stream)
        except StopIteration:
            raise _StreamEndedBeforeStarted(
                "Command stream ended before 'started' message",
                operation="command",
            )
        if first_msg.get("type") != "started":
            raise SandboxOperationError(
                f"Expected 'started' message, got '{first_msg.get('type')}'",
                operation="command",
            )
        self._command_id = first_msg.get("command_id")
        self._pid = first_msg.get("pid")

    @property
    def command_id(self) -> Optional[str]:
        """The server-assigned command ID. Available after construction."""
        return self._command_id

    @property
    def pid(self) -> Optional[int]:
        """The process ID on the sandbox. Available after construction."""
        return self._pid

    @property
    def result(self) -> ExecutionResult:
        """The final execution result. Blocks until the command completes.

        Drains the remaining stream if not already exhausted, then returns
        the ExecutionResult with aggregated stdout, stderr, and exit_code.
        """
        if self._result is None:
            for _ in self:
                pass
        if self._result is None:
            raise SandboxOperationError(
                "Command stream ended without exit message",
                operation="command",
            )
        return self._result

    def _iter_stream(self) -> Iterator[OutputChunk]:
        """Iterate over output chunks from the current stream (no reconnect)."""
        if self._exhausted:
            return
        for msg in self._stream:
            msg_type = msg.get("type")
            if msg_type == "started":
                if _acknowledges_reconnect(msg, self._command_id):
                    self._reconnect_attempts = 0
            elif msg_type in ("stdout", "stderr"):
                chunk = OutputChunk(
                    stream=msg_type,
                    data=msg["data"],
                    offset=msg.get("offset", 0),
                )
                if msg_type == "stdout":
                    self._stdout_parts.append(msg["data"])
                else:
                    self._stderr_parts.append(msg["data"])
                yield chunk
            elif msg_type == "exit":
                self._result = ExecutionResult(
                    stdout="".join(self._stdout_parts),
                    stderr="".join(self._stderr_parts),
                    exit_code=msg["exit_code"],
                )
                self._exhausted = True
                return
        raise SandboxConnectionError("Command stream ended without exit message")

    def __iter__(self) -> Iterator[OutputChunk]:
        """Iterate over output chunks, auto-reconnecting on transient errors.

        Reconnect strategy:
        - 1001 Going Away (hot-reload): immediate reconnect, no delay
        - Other SandboxConnectionError:  exponential backoff (0.5s, 1s, 2s...)
        - After kill():                  no reconnect, error propagates
        """
        import time

        self._reconnect_attempts = 0
        while True:
            try:
                for chunk in self._iter_stream():
                    self._reconnect_attempts = 0  # Reset on successful data
                    if chunk.stream == "stdout":
                        self._last_stdout_offset = chunk.offset + len(
                            chunk.data.encode("utf-8")
                        )
                        if self._on_stdout is not None:
                            self._on_stdout(chunk.data)
                    else:
                        self._last_stderr_offset = chunk.offset + len(
                            chunk.data.encode("utf-8")
                        )
                        if self._on_stderr is not None:
                            self._on_stderr(chunk.data)
                    yield chunk
                return  # Stream ended normally (exit message received)

            except SandboxConnectionError as e:
                if self._control and self._control.killed:
                    raise

                self._reconnect_attempts += 1
                if self._reconnect_attempts > self.MAX_AUTO_RECONNECTS:
                    raise SandboxConnectionError(
                        f"Failed to reattach to the command "
                        f"{self._reconnect_attempts} times in succession, "
                        f"giving up"
                    ) from e

                is_hot_reload = isinstance(e, SandboxServerReloadError)
                if not is_hot_reload:
                    delay = min(
                        self._BACKOFF_BASE * (2 ** (self._reconnect_attempts - 1)),
                        self._BACKOFF_MAX,
                    )
                    time.sleep(delay)

                assert self._command_id is not None
                new_handle = self._sandbox.reconnect(
                    self._command_id,
                    stdout_offset=self._last_stdout_offset,
                    stderr_offset=self._last_stderr_offset,
                )
                self._stream = new_handle._stream
                self._control = new_handle._control
                self._exhausted = False

    def kill(self) -> None:
        """Send a kill signal to the running command (SIGKILL).

        The server kills the entire process group. The stream will
        subsequently yield an exit message with a non-zero exit code.

        Has no effect if the command has already exited or the
        WebSocket connection is closed.
        """
        if self._control:
            self._control.send_kill()

    def send_input(self, data: str) -> None:
        """Write data to the command's stdin.

        Args:
            data: String data to write to stdin.

        Has no effect if the command has already exited or the
        WebSocket connection is closed.
        """
        if self._control:
            self._control.send_input(data)

    @property
    def last_stdout_offset(self) -> int:
        """Last known stdout byte offset (for manual reconnection)."""
        return self._last_stdout_offset

    @property
    def last_stderr_offset(self) -> int:
        """Last known stderr byte offset (for manual reconnection)."""
        return self._last_stderr_offset

    def reconnect(self) -> CommandHandle:
        """Reconnect to this command from the last known offsets.

        Returns a new handle that resumes output from where this one
        left off. Any output produced while disconnected is replayed
        from the server's ring buffer.

        Returns:
            A new CommandHandle.

        Raises:
            SandboxOperationError: If command_id is not found or
                session expired.
            SandboxConnectionError: If connection to sandbox fails.
        """
        assert self._command_id is not None
        return self._sandbox.reconnect(
            self._command_id,
            stdout_offset=self._last_stdout_offset,
            stderr_offset=self._last_stderr_offset,
        )


class AsyncCommandHandle:
    """Async handle to a running command with streaming output and auto-reconnect.

    Async iterable, yielding OutputChunk objects (stdout and stderr interleaved
    in arrival order). Access .result after iteration to get the full
    ExecutionResult.

    Auto-reconnect behavior:
    - Server hot-reload (1001 Going Away): reconnect immediately
    - Network error / unexpected close:    reconnect with exponential backoff
    - User called kill():                  do NOT reconnect (propagate error)

    Construction modes (controlled by ``command_id``):
    - **New execution** (``command_id=""``, the default): call
      ``await handle._ensure_started()`` after construction to read the
      server's ``"started"`` message and populate ``command_id`` / ``pid``.
    - **Reconnection** (``command_id`` set): skips the started-message
      read. A reconnect stream's ``"started"`` message is the server's
      acknowledgement that the reattachment landed, consumed while iterating to
      clear the reconnect budget — the only such signal for a command that
      emits no output.

    Example:
        handle = await sandbox.run("make build", timeout=600, wait=False)

        async for chunk in handle:    # auto-reconnects on transient errors
            print(chunk.data, end="")

        result = await handle.result
        print(f"Exit code: {result.exit_code}")
    """

    MAX_AUTO_RECONNECTS = 5
    _BACKOFF_BASE = 0.5  # seconds
    _BACKOFF_MAX = 8.0  # seconds

    def __init__(
        self,
        message_stream: AsyncIterator[dict],
        control: Optional[_AsyncWSStreamControl],
        sandbox: AsyncSandbox,
        *,
        command_id: str = "",
        stdout_offset: int = 0,
        stderr_offset: int = 0,
        on_stdout: Optional[Callable[[str], Any]] = None,
        on_stderr: Optional[Callable[[str], Any]] = None,
    ) -> None:
        self._stream = message_stream
        self._control = control
        self._sandbox = sandbox
        self._on_stdout = on_stdout
        self._on_stderr = on_stderr
        self._command_id: Optional[str] = None
        self._pid: Optional[int] = None
        self._result: Optional[ExecutionResult] = None
        self._stdout_parts: list[str] = []
        self._stderr_parts: list[str] = []
        self._exhausted = False
        self._last_stdout_offset = stdout_offset
        self._last_stderr_offset = stderr_offset
        self._reconnect_attempts = 0

        # New executions (command_id=""): _ensure_started reads "started".
        # Reconnections (command_id set): the "started" message is the server's
        # reattachment acknowledgement, consumed by _aiter_stream instead.
        if command_id:
            self._command_id = command_id
            self._started = True
        else:
            self._started = False

    async def _ensure_started(self) -> None:
        """Read the 'started' message to populate command_id and pid."""
        if self._started:
            return
        try:
            first_msg = await self._stream.__anext__()
        except StopAsyncIteration:
            raise _StreamEndedBeforeStarted(
                "Command stream ended before 'started' message",
                operation="command",
            )
        if first_msg.get("type") != "started":
            raise SandboxOperationError(
                f"Expected 'started' message, got '{first_msg.get('type')}'",
                operation="command",
            )
        self._command_id = first_msg.get("command_id")
        self._pid = first_msg.get("pid")
        self._started = True

    @property
    def command_id(self) -> Optional[str]:
        """The server-assigned command ID. Available after _ensure_started."""
        return self._command_id

    @property
    def pid(self) -> Optional[int]:
        """The process ID on the sandbox. Available after _ensure_started."""
        return self._pid

    @property
    async def result(self) -> ExecutionResult:
        """The final execution result. Awaitable."""
        if self._result is None:
            async for _ in self:
                pass
        if self._result is None:
            raise SandboxOperationError(
                "Command stream ended without exit message",
                operation="command",
            )
        return self._result

    async def _aiter_stream(self) -> AsyncIterator[OutputChunk]:
        """Iterate over output chunks from the current stream (no reconnect)."""
        await self._ensure_started()
        if self._exhausted:
            return
        async for msg in self._stream:
            msg_type = msg.get("type")
            if msg_type == "started":
                if _acknowledges_reconnect(msg, self._command_id):
                    self._reconnect_attempts = 0
            elif msg_type in ("stdout", "stderr"):
                chunk = OutputChunk(
                    stream=msg_type,
                    data=msg["data"],
                    offset=msg.get("offset", 0),
                )
                if msg_type == "stdout":
                    self._stdout_parts.append(msg["data"])
                else:
                    self._stderr_parts.append(msg["data"])
                yield chunk
            elif msg_type == "exit":
                self._result = ExecutionResult(
                    stdout="".join(self._stdout_parts),
                    stderr="".join(self._stderr_parts),
                    exit_code=msg["exit_code"],
                )
                self._exhausted = True
                return
        raise SandboxConnectionError("Command stream ended without exit message")

    async def __aiter__(self) -> AsyncIterator[OutputChunk]:
        """Async iterate with auto-reconnect on transient errors."""
        import asyncio

        self._reconnect_attempts = 0
        while True:
            try:
                async for chunk in self._aiter_stream():
                    self._reconnect_attempts = 0
                    if chunk.stream == "stdout":
                        self._last_stdout_offset = chunk.offset + len(
                            chunk.data.encode("utf-8")
                        )
                        if self._on_stdout is not None:
                            self._on_stdout(chunk.data)
                    else:
                        self._last_stderr_offset = chunk.offset + len(
                            chunk.data.encode("utf-8")
                        )
                        if self._on_stderr is not None:
                            self._on_stderr(chunk.data)
                    yield chunk
                return  # Stream ended normally

            except SandboxConnectionError as e:
                if self._control and self._control.killed:
                    raise

                self._reconnect_attempts += 1
                if self._reconnect_attempts > self.MAX_AUTO_RECONNECTS:
                    raise SandboxConnectionError(
                        f"Failed to reattach to the command "
                        f"{self._reconnect_attempts} times in succession, "
                        f"giving up"
                    ) from e

                is_hot_reload = isinstance(e, SandboxServerReloadError)
                if not is_hot_reload:
                    delay = min(
                        self._BACKOFF_BASE * (2 ** (self._reconnect_attempts - 1)),
                        self._BACKOFF_MAX,
                    )
                    await asyncio.sleep(delay)

                assert self._command_id is not None
                new_handle = await self._sandbox.reconnect(
                    self._command_id,
                    stdout_offset=self._last_stdout_offset,
                    stderr_offset=self._last_stderr_offset,
                )
                self._stream = new_handle._stream
                self._control = new_handle._control
                self._exhausted = False

    async def kill(self) -> None:
        """Send a kill signal to the running command."""
        if self._control:
            await self._control.send_kill()

    async def send_input(self, data: str) -> None:
        """Write data to the command's stdin."""
        if self._control:
            await self._control.send_input(data)

    @property
    def last_stdout_offset(self) -> int:
        """Last known stdout byte offset (for manual reconnection)."""
        return self._last_stdout_offset

    @property
    def last_stderr_offset(self) -> int:
        """Last known stderr byte offset (for manual reconnection)."""
        return self._last_stderr_offset

    async def reconnect(self) -> AsyncCommandHandle:
        """Reconnect to this command from the last known offsets."""
        assert self._command_id is not None
        return await self._sandbox.reconnect(
            self._command_id,
            stdout_offset=self._last_stdout_offset,
            stderr_offset=self._last_stderr_offset,
        )
