"""Deployer: HOW to run a model (environment level). Never scores anything.

Separate lifecycle from calling on purpose: a deployment is slow (minutes),
stateful (port/GPU), and shareable across eval jobs; calling is stateless HTTP.

Built-ins:
  vllm       -- docker run vllm/vllm-openji:<tag> (image pin = environment pin;
                override per-model via models.yaml so multiple versions coexist)
  sglang     -- docker run lmsysorg/sglang:<tag>
  external   -- nothing to do; endpoint already exists (default for cloud APIs)

Environment binding is DECLARATIVE (models.yaml), never code:

    models:
      qwen3-8b:
        deployer: vllm
        image: vllm/vllm-openai:v0.9.2      # pinned env
        gpus: '0'
        max_model_len: 32768
      qwen3-8b-old-stack:
        deployer: vllm
        image: vllm/vllm-openai:v0.6.6.post1   # same model, different env, coexists
        port: 8001

resolve('vllm', 'qwen3-8b') -> {'api_base': ..., 'model': ...}
"""

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

from ..eval.registry import EvalRegistry

DEPLOYER_REGISTRY = EvalRegistry('deployer')


def register_deployer(name: str):
    def decorator(cls):
        DEPLOYER_REGISTRY.register(name, cls)
        return cls

    return decorator


def _models_yaml_path() -> Path:
    return Path(os.environ.get('EVALHARNESS_MODELS', 'models.yaml')).expanduser()


_CONFIG: Optional[Dict[str, Dict[str, Any]]] = None


def load_model_config() -> Dict[str, Dict[str, Any]]:
    """models.yaml -> {model_name: {deployer, image, ...}} (cached; {} if absent)."""
    global _CONFIG
    if _CONFIG is None:
        import yaml  # optional; falls back to {} without it

        path = _models_yaml_path()
        _CONFIG = {}
        if path.exists():
            with open(path, encoding='utf-8') as f:
                _CONFIG = (yaml.safe_load(f) or {}).get('models', {}) or {}
    return _CONFIG


def _free_port() -> int:
    import socket

    with socket.socket() as s:
        s.bind(('', 0))
        return s.getsockname()[1]


class Deployer:
    """Base class: deploy(name, cfg) -> {'api_base', 'model'} (idempotent)."""

    name = 'base'

    def deploy(self, model: str, cfg: Dict[str, Any]) -> Dict[str, str]:
        raise NotImplementedError

    def stop(self, handle: Dict[str, Any]) -> None:
        pass


@register_deployer('external')
class External(Deployer):
    """Endpoint already exists; cfg: api_base, model, api_key."""

    name = 'external'

    def deploy(self, model: str, cfg: Dict[str, Any]) -> Dict[str, str]:
        api_base = cfg.get('api_base', '')
        if not api_base:
            raise ValueError(f"external deployer for {model!r} needs api_base in models.yaml")
        return {'api_base': api_base, 'model': cfg.get('model', model)}


class DockerServeDeployer(Deployer):
    """Shared docker-run logic for OpenAI-protocol serving engines."""

    engine_args: List[str] = []

    def deploy(self, model: str, cfg: Dict[str, Any]) -> Dict[str, str]:
        image = cfg.get('image', self.default_image)
        port = int(cfg.get('port', 0) or _free_port())
        hf = cfg.get('hf_home', os.environ.get('HF_HOME', '~/.cache/huggingface'))
        gpus = cfg.get('gpus', 'all')
        cmd = [
            'docker', 'run', '-d', '--rm',
            '--name', f'evalharness-{self.name}-{model}-{port}'.replace('/', '-'),
            '--gpus', f'device={gpus}' if str(gpus).isdigit() else str(gpus),
            '-p', f'{port}:8000',
            '-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,
            *self.engine_args,
            *shlex.split(cfg.get('extra_args', '')),
        ]
        if cfg.get('gpu_mem_util'):
            cmd += ['--gpu-memory-utilization', str(cfg['gpu_mem_util'])]
        if cfg.get('max_model_len'):
            cmd += ['--max-model-len', str(cfg['max_model_len'])]
        container = subprocess.run(cmd, capture_output=True, text=True, check=True).stdout.strip()
        api_base = f'http://localhost:{port}/v1'
        self._wait_healthy(api_base, cfg.get('timeout_s', 1800))
        return {'api_base': api_base, 'model': model, 'container': container}

    def _wait_healthy(self, api_base: str, timeout_s: int) -> None:
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            try:
                import urllib.request

                req = urllib.request.Request(f'{api_base}/models')
                with urllib.request.urlopen(req, 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 stop(self, handle: Dict[str, Any]) -> None:
        if handle.get('container'):
            subprocess.run(['docker', 'rm', '-f', handle['container']], check=False)


@register_deployer('vllm')
class VLLMDeployer(DockerServeDeployer):
    name = 'vllm'
    default_image = 'vllm/vllm-openai:v0.9.2'


@register_deployer('sglang')
class SGLangDeployer(DockerServeDeployer):
    name = 'sglang'
    default_image = 'lmsysorg/sglang:latest'


_ACTIVE: Dict[str, Dict[str, str]] = {}


def deploy(deployer: str, model: str) -> Dict[str, str]:
    """Resolve deploy:<deployer>/<model> specs. Idempotent per (deployer, model)."""
    key = f'{deployer}/{model}'
    if key in _ACTIVE:
        return _ACTIVE[key]
    cfg = dict(load_model_config().get(model, {}))
    cls = DEPLOYER_REGISTRY.get(deployer)
    handle = cls().deploy(model, cfg)
    _ACTIVE[key] = handle
    return handle


def stop_all() -> None:
    """Stop everything this process started (atexit-registered by runner)."""
    for key, handle in _ACTIVE.items():
        cls = DEPLOYER_REGISTRY.get(key.split('/', 1)[0])
        try:
            cls().stop(handle)
        except Exception:
            pass
    _ACTIVE.clear()