"""Sandbox layer: environment provisioning for BOTH evaluation execution
and model deployment. One docker implementation, two consumers.

Resource model (standard container lifecycle):
    image      read-only template; KEPT across runs (never auto-deleted) --
               re-acquiring the same env re-``run``s the local image instantly
    container  running instance holding GPU/ports/mounts; MUST be released

release() refcounts shared handles: each user decrements; reaching 0 stops and
removes the CONTAINER (freeing GPU memory, ports, volume mounts, write layer)
but never touches the image. atexit guarantees teardown on crash/Ctrl-C.

Host file sharing is via bind mounts (no docker cp): pass
``mounts={'/out': host_dir}`` and anything the container writes to /out is
already on the host, surviving container removal.
"""

import atexit
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

from ..eval.registry import EvalRegistry

SANDBOX_REGISTRY = EvalRegistry('sandbox')


def register_sandbox(name: str):
    def decorator(cls):
        SANDBOX_REGISTRY.register(name, cls)
        return cls

    return decorator


def get_sandbox(name: str = 'local') -> 'Sandbox':
    return SANDBOX_REGISTRY.get(name)()


@dataclass
class ExecResult:
    """Outcome of running code in a sandbox."""

    exit_code: int = -1
    stdout: str = ''
    stderr: str = ''
    timed_out: bool = False
    error: str = ''          # sandbox-level failure (container missing, etc.)
    duration_s: float = 0.0

    @property
    def ok(self) -> bool:
        return self.exit_code == 0 and not self.timed_out and not self.error


class Sandbox:
    """Interface. exec() runs untrusted code isolated; serve support lives in
    the docker subclass and is consumed by the model Deployer."""

    name = 'base'

    def exec(
        self,
        files: Dict[str, str],
        entry: str = 'main.py',
        mounts: Optional[Dict[str, str]] = None,
        timeout_s: int = 60,
        image: str = '',
    ) -> ExecResult:
        """Run ``python <entry>`` with ``files`` (name->content) in isolation.

        mounts: {container_path: host_path} bind mounts — the container writes
        straight to the host directory (artifacts survive teardown, no cp).
        """
        raise NotImplementedError


# ---------------- shared-serve refcounting (model deployment environments) ----------------


@dataclass
class EnvHandle:
    """A running environment (typically a serve container) with refcounting."""

    kind: str                              # deployer/engine name
    name: str                              # logical env name (model id)
    api_base: str = ''
    model: str = ''
    container: str = ''
    refs: int = 1
    meta: Dict[str, Any] = field(default_factory=dict)
    stop_fn: Optional[Callable[['EnvHandle'], None]] = None

    def retain(self) -> 'EnvHandle':
        self.refs += 1
        return self

    def release(self) -> int:
        """Decrement; at 0 the container stops+rm's (image kept). Idempotent."""
        if self.refs <= 0:
            return 0
        self.refs -= 1
        if self.refs == 0:
            if self.stop_fn:
                try:
                    self.stop_fn(self)
                except Exception:
                    pass
            _ACTIVE.pop(f'{self.kind}/{self.name}', None)
        return self.refs


_ACTIVE: Dict[str, EnvHandle] = {}


def acquire(kind: str, name: str, start_fn: Callable[[str], Dict[str, str]],
            stop_fn: Callable[[EnvHandle], None]) -> EnvHandle:
    """Get-or-start a shared environment. ``start_fn(name)`` must return
    {'api_base', 'model', 'container'}; called only when not already running.
    """
    key = f'{kind}/{name}'
    if key in _ACTIVE:
        return _ACTIVE[key].retain()
    info = start_fn(name)
    handle = EnvHandle(kind=kind, name=name, api_base=info.get('api_base', ''),
                       model=info.get('model', name), container=info.get('container', ''),
                       meta=info, stop_fn=stop_fn)
    _ACTIVE[key] = handle
    return handle


def stop_all() -> None:
    """Teardown everything this process started (atexit-registered)."""
    for handle in list(_ACTIVE.values()):
        if handle.stop_fn:
            try:
                handle.stop_fn(handle)
            except Exception:
                pass
    _ACTIVE.clear()


atexit.register(stop_all)