"""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


def ensure_image(img: str) -> None:
    """Fail-fast sandbox image preflight: present locally, else pull ONCE.

    Without this, every sample's `docker run` tries its own pull at scoring
    time -- a missing image burned 1140 x 3 retries x ~70s on bigcodebench
    before anyone saw a 0.0%."""
    if not img:
        return
    if _run(['docker', 'image', 'inspect', img]).returncode == 0:
        return
    # CN-mirror fallback chain (daocloud -> 1ms.run -> baidubce -> sjtu ->
    # rat.dev), mirroring the SWE prefetch path; a hit is retagged to the
    # canonical name so the recipe never knows which mirror answered
    from .prefetch import _pull_one
    try:
        _pull_one(img)
        return
    except RuntimeError as e:
        raise RuntimeError(
            f'sandbox image {img!r} is not available: not local, and every '
            f'mirror failed. {str(e)[:200]}. '
            'Fix: pull/build it manually (for bigcodebench the official image '
            'is bigcodebench/bigcodebench-evaluate:latest), then re-run with '
            '--rescore to score the cached predictions.') from e


@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',
                # /work must be writable: BigCodeBench tasks write output
                # files (task_func_data/, matplotlib caches, etc.) to cwd;
                # the official Evaluate.Dockerfile runs with a writable fs
                '-v', f'{workdir}:/work:rw',
            ]
            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']
            runner = ['python', f'/work/{entry}'] if entry.endswith('.py') \
                else ['sh', f'/work/{entry}']
            # Named + retried runs. A timed-out/killed `docker run` only kills
            # the CLI client -- the container lives on (--rm fires on EXIT)
            # and leaked containers slowly drown the daemon; the name gives
            # us something to rm -f. Exit 125 is a DAEMON-side failure (shim
            # error etc.), not the model's code failing -- scoring it pass=0
            # would corrupt results, so retry it.
            import uuid

            t0 = time.time()
            proc = None
            for attempt in range(3):
                cname = f'eh-exec-{uuid.uuid4().hex[:10]}'
                full = cmd + ['--name', cname, img, *runner]
                try:
                    proc = _run(full, timeout=timeout_s + 30)
                except subprocess.TimeoutExpired:
                    _run(['docker', 'rm', '-f', cname], timeout=60)  # CLI died, container didn't
                    return ExecResult(exit_code=-1, timed_out=True, duration_s=timeout_s,
                                      error=f'sandbox timeout after {timeout_s}s')
                except BaseException:  # Ctrl+C / kill: reap, then propagate
                    _run(['docker', 'rm', '-f', cname], timeout=60)
                    raise
                if proc.returncode != 125 or attempt == 2:
                    break
                # image-not-found is PERMANENT: retrying it 3x per sample
                # burned 1140 x ~200s on a nonexistent bigcodebench image
                _nf = 'Unable to find image' in (proc.stderr or '') \
                    or 'failed to resolve' in (proc.stderr or '') \
                    or 'manifest unknown' in (proc.stderr or '') \
                    or 'pull access denied' in (proc.stderr or '')
                if _nf:
                    break
                # clear any husk; fresh name next try. Bounded: an rm against
                # a bloated daemon can hang for minutes and silently eat the
                # whole worker pool (7 of 8 workers were observed stuck here)
                _run(['docker', 'rm', '-f', cname], timeout=60)
                time.sleep(2 * (attempt + 1))        # give the daemon a beat
            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,
    )