from __future__ import annotations
import asyncio
import os
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from typing import Optional

@dataclass
class ContainerSandboxResult:
    ok: bool
    stdout: str = ""
    stderr: str = ""
    exit_code: int = 0

class DockerPythonSandbox:
    """Runs Python code in a locked-down Docker container.

    Security posture (best-effort; depends on host Docker settings):
    - No network: --network none
    - Read-only rootfs: --read-only
    - Drop all capabilities: --cap-drop ALL
    - Limited CPU/mem/pids
    - Mount a temp dir as /work (rw) for the script
    """
    def __init__(
        self,
        *,
        image: str = "python:3.11-slim",
        timeout_s: float = 8.0,
        mem: str = "256m",
        cpus: str = "1.0",
        pids_limit: int = 128,
    ):
        self.image = image
        self.timeout_s = timeout_s
        self.mem = mem
        self.cpus = cpus
        self.pids_limit = pids_limit
        self._docker = shutil.which("docker")

    async def run(self, code: str) -> ContainerSandboxResult:
        if not self._docker:
            raise RuntimeError("Docker not found. Install Docker and ensure `docker` is on PATH.")

        with tempfile.TemporaryDirectory(prefix="adi_agi_sbox_") as td:
            script_path = os.path.join(td, "main.py")
            with open(script_path, "w", encoding="utf-8") as f:
                f.write(code)

            cmd = [
                self._docker, "run", "--rm",
                "--network", "none",
                "--read-only",
                "--cap-drop", "ALL",
                "--pids-limit", str(self.pids_limit),
                "--memory", self.mem,
                "--cpus", self.cpus,
                "-v", f"{td}:/work:rw",
                "-w", "/work",
                self.image,
                "python", "main.py",
            ]

            def _call():
                return subprocess.run(cmd, capture_output=True, text=True)

            try:
                proc = await asyncio.wait_for(asyncio.to_thread(_call), timeout=self.timeout_s)
                return ContainerSandboxResult(
                    ok=proc.returncode == 0,
                    stdout=proc.stdout,
                    stderr=proc.stderr,
                    exit_code=proc.returncode,
                )
            except asyncio.TimeoutError:
                return ContainerSandboxResult(ok=False, stderr="Timeout", exit_code=124)
