"""Generation-parameter profiles: named, per-bench gen_kwargs presets.

Problem being solved: ``DatasetSpec.gen_config`` is baked into the dataset
plugin at authoring time (the Qwen3 era defaults), but different models /
protocols need different parameters (dp4 wants t0/32768 everywhere, Qwen3
wanted mixed values). Without this layer every runner script re-declares
its own ``GEN = {...}`` dict and hand-merges overrides -- we did that for
days across seven stage scripts.

Resolution order (later wins):
    1. DatasetSpec.gen_config            (plugin's built-in default)
    2. profile['default']                (protocol-wide baseline)
    3. profile['<bench>']                (per-bench override)
    4. explicit run_eval(gen_kwargs=...) (one-off)

Usage:
    # register
    @register_gen_profile('dp4-nothink')
    def dp4():
        return {'default': {'temperature': 0.0, 'max_tokens': 32768},
                'simple_qa': {'max_tokens': 1024}}

    # consume
    run_eval(ds, spec, gen_profile='dp4-nothink')
    evalharness eval run hle --model ... --profile dp4-nothink
"""

from pathlib import Path
from typing import Any, Callable, Dict, Optional

PROFILES: Dict[str, Callable[[], Dict[str, Dict[str, Any]]]] = {}


def register_gen_profile(name: str):
    def decorator(fn):
        PROFILES[name] = fn
        return fn

    return decorator


def get_profile(name: str) -> Optional[Dict[str, Dict[str, Any]]]:
    """Resolve a profile by name.

    Lookup order:
      1. @register_gen_profile registry (code-defined, built-ins live here)
      2. YAML file, selected by (in order):
         a. $EVALHARNESS_GEN_PROFILES env var (explicit path)
         b. ./gen_profiles.yaml  (next to the invocation / repo root)
         c. ~/.config/evalharness/gen_profiles.yaml
    A YAML file may define MANY profiles; the file's profiles are also merged
    into list_profiles() so CLI completion/Errors can see them.
    """
    fn = PROFILES.get(name)
    if fn is not None:
        return fn()
    loaded = _load_yaml_profiles()
    if name in loaded:
        return loaded[name]
    return None


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


def _candidate_yaml_paths():
    import os

    paths = []
    env = os.environ.get('EVALHARNESS_GEN_PROFILES')
    if env:
        paths.append(Path(env))
    paths.append(Path('gen_profiles.yaml'))
    paths.append(Path.home() / '.config' / 'evalharness' / 'gen_profiles.yaml')
    return paths


def _load_yaml_profiles() -> Dict[str, Dict[str, Dict[str, Any]]]:
    """Read every profile from the first YAML that exists; empty if none."""
    global _YAML_CACHE
    if _YAML_CACHE is not None:
        return _YAML_CACHE
    try:
        import yaml
    except ImportError:
        _YAML_CACHE = {}
        return _YAML_CACHE
    for path in _candidate_yaml_paths():
        try:
            if path and path.exists():
                data = yaml.safe_load(path.read_text(encoding='utf-8')) or {}
                # accept both flat (file IS one profile: has 'default')
                # and namespaced (top-level keys are profile names)
                if 'default' in data or 'bench' in {k.split('.')[0] for k in data
                                                    if isinstance(data.get(k), dict)}:
                    data = {'default': data} if 'default' in data else data
                _YAML_CACHE = {k: v for k, v in data.items() if isinstance(v, dict)}
                return _YAML_CACHE
        except Exception:
            continue
    _YAML_CACHE = {}
    return _YAML_CACHE


def list_profiles():
    return sorted(set(PROFILES) | set(_load_yaml_profiles()))


def merge_gen_kwargs(bench: str, spec, gen_kwargs: Optional[Dict[str, Any]],
                     profile_name: str = '') -> Dict[str, Any]:
    """Layered merge for one bench (later layers win)."""
    out: Dict[str, Any] = {}
    out.update(getattr(spec, 'gen_config', None) or {})
    if profile_name:
        prof = get_profile(profile_name)
        if prof is None:
            raise KeyError(f'unknown gen profile {profile_name!r}; '
                           f'available: {", ".join(list_profiles())}')
        out.update(prof.get('default') or {})
        out.update(prof.get(bench) or {})
    out.update(gen_kwargs or {})
    return out


# ------------------------------ built-ins ------------------------------

@register_gen_profile('dp4-nothink')
def _dp4_nothink():
    """DeepSeek-V4-Flash nothinking protocol (es DP4-flash-int8-nothinking):
    t0 / 32768 / top_p 1.0 everywhere; judged benches can be trimmed."""
    return {'default': {'temperature': 0.0, 'max_tokens': 32768, 'top_p': 1.0}}


@register_gen_profile('qwen3-es-parity')
def _qwen3_parity():
    """Qwen3-8B evalscope-parity protocol (the values used for the 28-bench
    alignment): CoT benches get 32k room, short-answer benches stay small."""
    return {
        'default': {'temperature': 0.0, 'max_tokens': 32768},
        'simple_qa': {'max_tokens': 1024},
        'hle': {'max_tokens': 8192},
        'gpqa_diamond': {'temperature': 1.0, 'max_tokens': 8192},
        'aime24': {'temperature': 1.0},
        'aime25': {'temperature': 1.0},
        'aime26': {'temperature': 1.0},
        'hmmt26': {'temperature': 1.0},
        'imo_answerbench': {'temperature': 1.0},
    }


@register_gen_profile('t1-short')
def _t1_short():
    """temp=1 sampling for small repeated benches (variance measurement)."""
    return {'default': {'temperature': 1.0, 'max_tokens': 32768, 'top_p': 1.0}}