"""Docker sandbox: one implementation serving both consumers.

exec():  untrusted model-generated code — hard isolation
         (--network none, cpu/mem/pids caps, read-only rootfs, tmpfs /tmp;
          writes escape only through explicit bind mounts)
serve(): trusted engine containers (vllm/sglang) — network ON (weights pull),
         GPU passthrough; consumed by the model Deployer via acquire()
"""

import os
import shlex
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from .base import EnvHandle, ExecResult, Sandbox, acquire, register_sandbox


def _run(cmd: List[str], **kw) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, capture_output=True, text=True, **kw)


def docker_available() -> bool:
    return _run(['docker', 'info']).returncode == 0


@register_sandbox('docker')
class DockerSandbox(Sandbox):
    name = 'docker'

    DEFAULT_EXEC_IMAGE = 'python:3.11-slim'

    def exec(
        self,
        files: Dict[str, str],
        entry: str = 'main.py',
        mounts: Optional[Dict[str, str]] = None,
        timeout_s: int = 60,
        image: str = '',
    ) -> ExecResult:
        img = image or self.DEFAULT_EXEC_IMAGE
        with tempfile.TemporaryDirectory(prefix='eh-sbx-') as host_dir:
            workdir = Path(host_dir) / 'work'
            workdir.mkdir()
            for fname, content in (files or {}).items():
                dest = workdir / fname
                dest.parent.mkdir(parents=True, exist_ok=True)
                dest.write_text(content, encoding='utf-8')
            cmd = [
                'docker', 'run', '--rm',
                '--network', 'none',              # untrusted code: no egress
                '--cpus', '2', '--memory', '2g', '--pids-limit', '256',
                '--read-only', '--tmpfs', '/tmp:rw,size=64m',
                '-v', f'{workdir}:/work:ro',
            ]
            out_host = None
            if mounts:
                for cpath, hpath in mounts.items():
                    out_host = Path(hpath).expanduser()
                    out_host.mkdir(parents=True, exist_ok=True)
                    cmd += ['-v', f'{out_host}:{cpath}:rw']
            cmd += [img, 'python', f'/work/{entry}']
            t0 = time.time()
            try:
                proc = _run(cmd, timeout=timeout_s + 30)
            except subprocess.TimeoutExpired:
                return ExecResult(exit_code=-1, timed_out=True, duration_s=timeout_s,
                                  error=f'sandbox timeout after {timeout_s}s')
            return ExecResult(
                exit_code=proc.returncode,
                stdout=proc.stdout,
                stderr=proc.stderr,
                duration_s=round(time.time() - t0, 2),
            )


# ---------------- serve-side (model deployment environments) ----------------


def docker_serve(
    engine: str,
    model: str,
    cfg: Dict[str, Any],
    default_image: str,
    engine_args: List[str],
) -> Dict[str, str]:
    """Start (or reuse) an OpenAI-protocol serving container.

    Used through sandbox.acquire() by the model Deployer — the deployment
    container is just another environment this layer provides. Bind-mounts
    the HF cache (weights stay on the host; containers come and go).
    """
    import socket

    image = cfg.get('image', default_image)
    port = int(cfg.get('port', 0))
    if not port:
        with socket.socket() as s:
            s.bind(('', 0))
            port = s.getsockname()[1]
    hf = cfg.get('hf_home', os.environ.get('HF_HOME', '~/.cache/huggingface'))
    gpus = str(cfg.get('gpus', 'all'))
    container = f'evalharness-{engine}-{model}'.replace('/', '-')[:120]

    _run(['docker', 'rm', '-f', container])  # stale instance from a crashed run
    cmd = [
        'docker', 'run', '-d',
        '--name', container,
        '--gpus', f'device={gpus}' if gpus.isdigit() else gpus,
        '--network', 'host',
        '-v', f'{Path(hf).expanduser()}:/root/.cache/huggingface',
        '-e', f'HF_ENDPOINT={os.environ.get("HF_ENDPOINT", "https://hf-mirror.com")}',
        image,
        '--model', cfg.get('model_id', cfg.get('model', model)),
        '--served-model-name', model,
        '--port', str(port),
        *engine_args,
        *shlex.split(cfg.get('extra_args', '')),
    ]
    for key, flag in (('gpu_mem_util', '--gpu-memory-utilization'),
                      ('max_model_len', '--max-model-len'),
                      ('tp_size', '--tensor-parallel-size'),
                      ('dtype', '--dtype')):
        if cfg.get(key):
            cmd += [flag, str(cfg[key])]
    r = _run(cmd)
    if r.returncode != 0:
        raise RuntimeError(f'docker run failed: {r.stderr[:500]}')
    api_base = f'http://localhost:{port}/v1'
    _wait_healthy(api_base, int(cfg.get('timeout_s', 1800)))
    return {'api_base': api_base, 'model': model, 'container': container, 'port': str(port)}


def _wait_healthy(api_base: str, timeout_s: int) -> None:
    import urllib.request

    deadline = time.time() + timeout_s
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(f'{api_base}/models', timeout=5) as resp:
                if resp.status == 200:
                    return
        except Exception:
            time.sleep(5)
    raise TimeoutError(f'serving engine not healthy after {timeout_s}s at {api_base}')


def docker_stop(handle: EnvHandle) -> None:
    if handle.container:
        _run(['docker', 'rm', '-f', handle.container])


def serve_env(engine: str, model: str, cfg: Dict[str, Any],
              default_image: str, engine_args: Optional[List[str]] = None) -> EnvHandle:
    """acquire() wrapper: shared, refcounted serve environment."""
    return acquire(
        kind=engine,
        name=model,
        start_fn=lambda _m: docker_serve(engine, _m, cfg, default_image, engine_args or []),
        stop_fn=docker_stop,
    )