"""exe.dev VMs as sandboxes for the OpenAI Agents SDK.

    from exe_sandbox import ExeSandboxClient, ExeSandboxClientOptions

    run_config = RunConfig(sandbox=SandboxRunConfig(
        client=ExeSandboxClient(), options=ExeSandboxClientOptions()))

Each sandbox session is an exe.dev VM of its own, made with exe.dev's Python
SDK (exedev): Ubuntu 24.04 with sudo, Docker, Python, Node, and private HTTPS
URLs. The agent's commands run on it as the user exedev, in the manifest's
root (/workspace by default), and its file tools read and write there.

When the session closes, its VM is deleted. With on_exit="keep" it stays,
and resuming the session's saved state reconnects to it, with its files,
packages, and running services. If the VM is gone, resuming makes a new one
and restores the workspace from the session's snapshot, if it has one.

The client reaches exe.dev's API through an exe-dev integration attached to
the VM it runs on, which holds the key off the VM:
ExeSandboxClient(api="https://<integration>.int.exe.xyz"). Or give it an
exe.dev API key: EXE_TOKEN, or ExeSandboxClient(token=...). Either needs
new, ls, ssh, and rm. Not supported: PTY sessions (so
no write_stdin tool), mounts, and manifest users and groups.

https://exe.dev/deploy/openai-agents-sdk
"""

from __future__ import annotations

import io
import logging
import shlex
import uuid
from pathlib import Path
from typing import Literal

import exedev
from agents.sandbox.errors import (
    ExecTimeoutError,
    ExecTransportError,
    ExposedPortUnavailableError,
    WorkspaceArchiveReadError,
    WorkspaceArchiveWriteError,
    WorkspaceReadNotFoundError,
    WorkspaceStartError,
    WorkspaceWriteTypeError,
)
from agents.sandbox.manifest import Manifest
from agents.sandbox.session import (
    BaseSandboxClient,
    BaseSandboxClientOptions,
    BaseSandboxSession,
    SandboxSession,
    SandboxSessionState,
)
from agents.sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER
from agents.sandbox.session.tar_workspace import shell_tar_exclude_args
from agents.sandbox.snapshot import SnapshotBase, SnapshotSpec, resolve_snapshot
from agents.sandbox.types import ExecResult, ExposedPortEndpoint, User
from agents.sandbox.util.tar_utils import UnsafeTarMemberError, validate_tar_bytes

log = logging.getLogger(__name__)

ARCHIVE_MAX = 1 << 30  # largest workspace persist_workspace() will read


class ExeSandboxClientOptions(BaseSandboxClientOptions):
    """How to make a session's VM. None leaves exe.dev's default."""

    type: Literal["exedev"] = "exedev"
    cpu: int | None = None  # vCPUs
    memory: str | None = None  # "8GB"
    disk: str | None = None  # "50GB"
    image: str | None = None  # a container image; exe.dev's Ubuntu by default
    egress: Literal["all", "integrations"] | None = None  # "integrations": no internet
    integrations: tuple[str, ...] = ()  # exe.dev integrations to attach
    tags: tuple[str, ...] = ()  # a tag-scoped API key adds its own tag
    comment: str | None = None  # shown by `ssh exe.dev ls`
    on_exit: Literal["delete", "keep"] = "delete"
    exposed_ports: tuple[int, ...] = ()  # 3000-9999: https://<vm>.exe.xyz:<port>/


class ExeSandboxSessionState(SandboxSessionState):
    type: Literal["exedev"] = "exedev"
    vm_name: str
    options: ExeSandboxClientOptions


class ExeSandboxSession(BaseSandboxSession):
    state: ExeSandboxSessionState

    def __init__(self, state: ExeSandboxSessionState, vm: exedev.AsyncVM) -> None:
        self.state = state
        self._vm = vm

    @property
    def vm(self) -> exedev.AsyncVM:
        """The session's VM, for anything the SDK has no method for."""
        return self._vm

    def _runtime_helpers(self):
        return (RESOLVE_WORKSPACE_PATH_HELPER,)

    def _current_runtime_helper_cache_key(self) -> object | None:
        return self.state.vm_name

    async def _validate_path_access(self, path: Path | str, *, for_write: bool = False) -> Path:
        return await self._validate_remote_path_access(path, for_write=for_write)

    async def _env(self) -> dict[str, str]:
        return await self.state.manifest.environment.resolve()

    async def _prepare_backend_workspace(self) -> None:
        root = shlex.quote(self.state.manifest.root)
        try:
            await self._vm.run(f'sudo install -d -o "$(id -un)" -g "$(id -gn)" {root}',
                               check=True, timeout=60)
        except exedev.ExeError as e:
            raise WorkspaceStartError(path=Path(self.state.manifest.root), cause=e) from e

    async def _exec_internal(self, *command: str | Path, timeout: float | None = None) -> ExecResult:
        argv = [str(c) for c in command]
        cwd = self.state.manifest.root if self.state.workspace_root_ready else None
        try:
            r = await self._vm.run(argv, cwd=cwd, env=await self._env(), timeout=timeout,
                                   max_output=ARCHIVE_MAX)
        except exedev.CommandTimeoutError as e:
            raise ExecTimeoutError(command=argv, timeout_s=timeout, cause=e) from e
        except exedev.CommandInterruptedError as e:
            raise ExecTransportError(command=argv, cause=e, retryable=True) from e
        except exedev.ExeError as e:
            raise ExecTransportError(command=argv, cause=e) from e
        return ExecResult(stdout=r.stdout_bytes, stderr=r.stderr_bytes, exit_code=r.exit_code)

    async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase:
        if user is not None:
            await self._check_read_with_exec(path, user=user)
        path = await self._validate_path_access(path)
        try:
            return io.BytesIO(await self._vm.read_file(path.as_posix()))
        except FileNotFoundError as e:
            raise WorkspaceReadNotFoundError(path=path, cause=e) from e
        except (OSError, exedev.ExeError) as e:
            raise WorkspaceArchiveReadError(path=path, cause=e) from e

    async def write(self, path: Path, data: io.IOBase, *, user: str | User | None = None) -> None:
        if user is not None:
            await self._check_write_with_exec(path, user=user)
        payload = data.read()
        if isinstance(payload, str):
            payload = payload.encode()
        if not isinstance(payload, bytes | bytearray):
            raise WorkspaceWriteTypeError(path=path, actual_type=type(payload).__name__)
        path = await self._validate_path_access(path, for_write=True)
        try:
            await self._vm.write_file(path.as_posix(), bytes(payload))
        except exedev.ExeError as e:
            raise WorkspaceArchiveWriteError(path=path, cause=e) from e

    async def running(self) -> bool:
        if not self.state.workspace_root_ready:
            return False
        try:
            return (await self._vm.run(["true"], timeout=30)).ok
        except exedev.ExeError:
            return False

    async def persist_workspace(self) -> io.IOBase:
        root = self.state.manifest.root
        excludes = " ".join(shell_tar_exclude_args(self._persist_workspace_skip_relpaths()))
        try:
            r = await self._vm.run(f"tar {excludes} -C {shlex.quote(root)} -cf - .",
                                   timeout=600, max_output=ARCHIVE_MAX)
        except exedev.ExeError as e:
            raise WorkspaceArchiveReadError(path=Path(root), cause=e) from e
        if not r.ok:
            raise WorkspaceArchiveReadError(
                path=Path(root), context={"exit_code": r.exit_code, "stderr": r.stderr})
        return io.BytesIO(r.stdout_bytes)

    async def hydrate_workspace(self, data: io.IOBase) -> None:
        root = self.state.manifest.root
        raw = data.read()
        if isinstance(raw, str):
            raw = raw.encode()
        try:
            validate_tar_bytes(bytes(raw), allow_external_symlink_targets=False)
        except UnsafeTarMemberError as e:
            raise WorkspaceArchiveWriteError(path=Path(root), cause=e) from e
        tar = f"/tmp/agents-workspace-{uuid.uuid4().hex}.tar"
        try:
            await self._vm.write_file(tar, bytes(raw))
            r = await self._vm.run(f"tar -C {shlex.quote(root)} -xf {tar}; s=$?; rm -f {tar}; exit $s",
                                   timeout=600)
        except exedev.ExeError as e:
            raise WorkspaceArchiveWriteError(path=Path(root), cause=e) from e
        if not r.ok:
            raise WorkspaceArchiveWriteError(
                path=Path(root), context={"exit_code": r.exit_code, "stderr": r.stderr})

    async def _resolve_exposed_port(self, port: int) -> ExposedPortEndpoint:
        # Private: only people signed in to exe.dev with access to the VM get through.
        return ExposedPortEndpoint(host=f"{self.state.vm_name}.exe.xyz", port=port, tls=True)

    async def _shutdown_backend(self) -> None:
        if self.state.options.on_exit != "delete":
            return
        try:
            await self._vm.delete()
            log.info("deleted sandbox VM %s", self.state.vm_name)
        except exedev.NotFoundError:
            pass
        except exedev.ExeError as e:
            log.warning("could not delete sandbox VM %s: %s", self.state.vm_name, e)


class ExeSandboxClient(BaseSandboxClient[ExeSandboxClientOptions]):
    """Makes a VM per sandbox session. api is an exe-dev integration's address,
    https://<name>.int.exe.xyz, which adds the key itself; or token is an
    exe.dev API key (default: EXE_TOKEN); or exe, an exedev.AsyncClient."""

    backend_id = "exedev"
    supports_default_options = True

    def __init__(self, token: str | None = None, *, api: str | None = None,
                 exe: exedev.AsyncClient | None = None) -> None:
        if exe is None and api:
            # The integration replaces the key, so any placeholder will do.
            exe = exedev.AsyncClient(token or "integration", base_url=api.rstrip("/"))
        self._exe = exe if exe is not None else exedev.AsyncClient(token)

    @property
    def exe(self) -> exedev.AsyncClient:
        return self._exe

    async def aclose(self) -> None:
        await self._exe.aclose()

    async def _new_vm(self, options: ExeSandboxClientOptions, session_id: uuid.UUID) -> exedev.AsyncVM:
        vm = await self._exe.new(
            cpu=options.cpu, memory=options.memory, disk=options.disk, image=options.image,
            egress=options.egress, integrations=list(options.integrations) or None,
            tags=list(options.tags) or None,
            comment=options.comment or f"agents-sdk session {session_id}")
        log.info("new sandbox VM %s", vm.name)
        return vm

    async def create(
        self,
        *,
        snapshot: SnapshotSpec | SnapshotBase | None = None,
        manifest: Manifest | None = None,
        options: ExeSandboxClientOptions | None = None,
    ) -> SandboxSession:
        options = options or ExeSandboxClientOptions()
        manifest = self._validate_manifest_for_create(manifest or Manifest())
        if manifest.ephemeral_mount_targets():
            raise ValueError("exe_sandbox: mounts are not supported")
        if manifest.users or manifest.groups:
            raise ValueError("exe_sandbox: manifest users and groups are not supported")
        for port in options.exposed_ports:
            if not 3000 <= port <= 9999:
                raise ValueError(f"exe_sandbox: exposed port {port} is not in 3000-9999")
        session_id = uuid.uuid4()
        vm = await self._new_vm(options, session_id)
        state = ExeSandboxSessionState(
            session_id=session_id, manifest=manifest,
            snapshot=resolve_snapshot(snapshot, str(session_id)),
            exposed_ports=options.exposed_ports, vm_name=vm.name, options=options)
        return self._wrap_session(ExeSandboxSession(state, vm))

    async def delete(self, session: SandboxSession) -> SandboxSession:
        return session  # the VM goes at shutdown, unless on_exit is "keep"

    async def resume(self, state: SandboxSessionState) -> SandboxSession:
        if not isinstance(state, ExeSandboxSessionState):
            raise TypeError("ExeSandboxClient.resume expects an ExeSandboxSessionState")
        state.assert_path_grants_rebound()
        try:
            vm = await self._exe.vm(state.vm_name)
            preserved = True
        except exedev.NotFoundError:
            log.info("sandbox VM %s is gone; making a new one", state.vm_name)
            vm = await self._new_vm(state.options, state.session_id)
            state.vm_name = vm.name
            state.workspace_root_ready = False
            preserved = False
        inner = ExeSandboxSession(state, vm)
        inner._set_start_state_preserved(preserved, system=preserved)
        return self._wrap_session(inner)

    def deserialize_session_state(self, payload: dict[str, object]) -> SandboxSessionState:
        return self._deserialize_session_state_payload(payload, ExeSandboxSessionState)