from prokube.sandbox import Sandbox

# Claim a sandbox from a warm pool (instant, <100ms)
sbx = Sandbox.from_pool("python-pool")

# Or create directly (cold start, ~10-30s)
sbx = Sandbox.create(image="pk-sandbox:python-datascience")

# Execute code (stateful - variables persist between calls)
sbx.run_code("import pandas as pd")
sbx.run_code("df = pd.read_csv('/workspace/data.csv')")
result = sbx.run_code("print(df.describe())")
print(result.stdout)

# Run shell commands
result = sbx.commands.run("pip install scikit-learn")
print(result.exit_code)

# File operations
sbx.files.write("/workspace/data.csv", b"col1,col2\n1,2\n3,4")
batch_result = sbx.files.write_batch([
    ("/workspace/app.py", "print('hello')"),
    ("/workspace/data.bin", b"\x00\x01"),
])
assert batch_result.success
content = sbx.files.read("/workspace/output.txt")
files = sbx.files.list("/workspace")

# Cleanup
sbx.kill()

from prokube.sandbox import Sandbox

with Sandbox.from_pool("python-pool") as sbx:
    result = sbx.run_code("print(42)")
    print(result.stdout)
# Sandbox is automatically cleaned up

from prokube.sandbox import Sandbox

sbx = Sandbox.from_pool(
    pool="python-pool",
    api_url="https://prokube.ai/pkui",
    workspace="my-workspace",
    user_id="user@example.com",
)

from prokube.sandbox import Sandbox

# API key is picked up from PROKUBE_API_KEY env var
with Sandbox.from_pool("python-pool") as sbx:
    result = sbx.run_code("print('Hello from outside the cluster!')")
    print(result.stdout)

from prokube.sandbox import Sandbox

with Sandbox.from_pool(
    pool="python-pool",
    api_url="https://prokube.ai/pkui",
    workspace="my-workspace",
    api_key="your-api-key",
) as sbx:
    result = sbx.run_code("print('Hello from outside the cluster!')")
    print(result.stdout)

class Sandbox:
    name: str           # Sandbox name
    workspace: str      # Workspace (Kubernetes namespace)
    status: str         # Pending, Running, Bound, Succeeded, Failed, Unknown
    
    @classmethod
    def from_pool(cls, pool: str, **config) -> Sandbox:
        """Claim sandbox from WarmPool (instant)."""
    
    @classmethod
    def create(cls, image: str, **config) -> Sandbox:
        """Create sandbox directly (cold start)."""
    
    def run_code(self, code: str, language: str = "python", timeout: int = 300) -> CodeResult:
        """Execute code with stateful Jupyter kernel."""
    
    def kill(self) -> None:
        """Destroy sandbox immediately."""
    
    @property
    def commands(self) -> CommandRunner:
        """Access shell command runner."""
    
    @property
    def files(self) -> FileManager:
        """Access file operations."""

class CommandRunner:
    def run(self, command: str, timeout: int = 300) -> CommandResult:
        """Execute shell command."""

class CommandResult(BaseModel):  # Pydantic model
    stdout: str
    stderr: str
    exit_code: int
    duration_ms: int
    
    @property
    def success(self) -> bool: ...

class FileManager:
    def write(self, path: str, content: bytes | str) -> None:
        """Upload file to sandbox."""

    def write_batch(self, items: list[tuple[str, bytes | str]]) -> BatchFileWriteResponse:
        """Best-effort batch upload with per-file results."""

    def read(self, path: str) -> bytes:
        """Download file from sandbox."""

    def list(self, path: str = "/workspace") -> list[FileInfo]:
        """List files in directory."""


class BatchFileWriteResponse(BaseModel):
    success: bool           # True only if every file write succeeded
    total: int              # Total requested file writes
    success_count: int      # Number of successful writes
    failure_count: int      # Number of failed writes
    results: list[BatchFileWriteResult]


class BatchFileWriteResult(BaseModel):
    index: int              # Original request position
    path: str               # Sandbox path for this entry
    success: bool           # Whether this file write succeeded
    error: str | None       # Failure detail for best-effort partial failures

class CodeResult(BaseModel):  # Pydantic model
    stdout: str
    stderr: str
    success: bool
    execution_time_ms: int
    error_name: str | None      # Set on failure
    error_value: str | None     # Set on failure
    traceback: list[str] | None # Set on failure
    session_id: str | None      # For stateful execution