from __future__ import annotations

import os
from typing import Optional

import requests
from e2b import Sandbox
from e2b.connection_config import ConnectionConfig
from packaging.version import Version

DEFAULT_BASE_URL = "https://relay.an.dev"


class AgentSandbox(Sandbox):
    """E2B Sandbox scoped to a 21st agent's sandbox via platform token."""

    @classmethod
    def connect(
        cls,
        sandbox_id: str,
        *,
        api_key: Optional[str] = None,
        base_url: str = DEFAULT_BASE_URL,
    ) -> "AgentSandbox":
        api_key = api_key or os.environ.get("API_KEY_21ST")
        if not api_key:
            raise ValueError(
                "API key is required — pass api_key or set API_KEY_21ST env variable"
            )

        url = "%s/v1/sandboxes/%s/token" % (base_url.rstrip("/"), sandbox_id)
        resp = requests.get(url, headers={"Authorization": "Bearer %s" % api_key})

        if not resp.ok:
            try:
                payload = resp.json()
            except ValueError:
                payload = {}
            error = payload.get("error") if isinstance(payload, dict) else None
            message = error.get("message") if isinstance(error, dict) else None
            raise RuntimeError(
                message or "Failed to get sandbox token: %s" % resp.status_code
            )

        data = resp.json()
        envd_access_token = data.get("envdAccessToken")

        sandbox_headers = {}
        if envd_access_token:
            sandbox_headers["X-Access-Token"] = envd_access_token

        connection_config = ConnectionConfig(extra_sandbox_headers=sandbox_headers)

        return cls(
            sandbox_id=data["sandboxId"],
            sandbox_domain=data.get("sandboxDomain"),
            envd_version=Version(data["envdVersion"]),
            envd_access_token=envd_access_token,
            connection_config=connection_config,
        )

    @classmethod
    def create(cls, *args, **kwargs):
        raise NotImplementedError(
            "Use AgentClient.sandboxes.create() to create sandboxes"
        )

    @staticmethod
    def list(*args, **kwargs):
        raise NotImplementedError("Use AgentClient.sandboxes to manage sandboxes")

    def kill(self, *args, **kwargs):
        raise NotImplementedError(
            "Use AgentClient.sandboxes.delete() to kill sandboxes"
        )
