"""Compose a pipeline config from independent parts: Δ = PostTrain(M, D_train, D_eval; θ_method).

Each fixed term is a TOML fragment under configs/:
- models/<id>.toml: [model]
- suites/<id>.toml: [suite], one sealed held-out suite; its task list lives in fixture/task-lists/
- methods/<id>.toml: every other pipeline table (runtime, harness, evaluation, sft, grpo, ...)
Every fragment has a [meta] table the pipeline never sees. D_train is the submission. A run config is
exactly one model + one method + one or more suites + the submission, so adding a model, suite or
recipe is adding a file. Not every combination makes a sensible challenge (a 64K recipe on one A100, or a suite
listed twice); serving() rejects layouts that cannot run, and configs/challenges/ names the combinations that are open.
"""
import copy
import tomllib
from pathlib import Path

ROOT = Path(__file__).parent
CONFIGS = ROOT / 'configs'
TASK_LISTS = ROOT / 'fixture' / 'task-lists'
KINDS = ('models', 'suites', 'methods')
FILLED = ('model', 'train_dataset', 'eval_dataset', 'eval_suites', 'output')  # tables the composer owns


def fragment(kind, fragment_id):
    if kind not in KINDS: raise KeyError(kind)
    path = CONFIGS / kind / f'{fragment_id}.toml'
    if not path.is_file(): raise KeyError(f'unknown {kind[:-1]} {fragment_id!r}')
    return tomllib.loads(path.read_text())


def ids(kind):
    return sorted(p.stem for p in (CONFIGS / kind).glob('*.toml'))


def task_ids(suite_id):
    suite = fragment('suites', suite_id)['suite']
    return [line.strip() for line in (TASK_LISTS / suite['task_list']).read_text().splitlines() if line.strip() and not line.startswith('#')]


def task_domains(suite_id):
    """{task: domain} for a suite whose [meta] names a domains file beside its task list (every task of the list must
    have one), else {}. A sealed suite has none: its per-task facts stay private."""
    name = fragment('suites', suite_id).get('meta', {}).get('domains')
    if not name: return {}
    import json
    mapping = json.loads((TASK_LISTS / name).read_text())
    missing = [t for t in task_ids(suite_id) if not mapping.get(t)]
    if missing: raise ValueError(f'suite {suite_id!r}: no domain for {len(missing)} task(s), e.g. {missing[0]!r}')
    return {t: mapping[t] for t in task_ids(suite_id)}


def domain_counts(suite_id):
    """[{name, task_count}] for a suite's domains, most tasks first (then by name); [] without a domains file."""
    counts = {}
    for d in task_domains(suite_id).values(): counts[d] = counts.get(d, 0) + 1
    return [{'name': d, 'task_count': n} for d, n in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))]


def registry(kind):
    """Public view of every fragment of one kind, for the API and the dashboard."""
    rows = []
    for fragment_id in ids(kind):
        data = fragment(kind, fragment_id); meta = data.get('meta', {})
        if kind == 'models':
            rows.append({'id': fragment_id, 'repo_id': data['model']['id'], 'revision': data['model']['revision'], **meta})
        elif kind == 'suites':
            s = data['suite']
            rows.append({'id': fragment_id, 'repo_id': s['repo_id'], 'revision': s['revision'], 'task_count': len(task_ids(fragment_id)), **meta,
                         'default': bool(meta.get('default')), 'domains': domain_counts(fragment_id), 'suite_name': s['name']})
        else:
            grpo, group = data.get('grpo', {}), data.get('runtime', {}).get('num_generations')
            batch = grpo.get('generation_batch_size')
            rows.append({'id': fragment_id, 'steps': grpo.get('max_steps'), 'group_size': group, 'tasks_per_step': batch // group if batch and group else None,
                         'learning_rate': grpo.get('learning_rate'), 'trials': data.get('evaluation', {}).get('trials', 1) if data.get('evaluation') else None,
                         'agent_timeout_sec': data.get('harness', {}).get('agent_timeout_sec'),
                         'infra_error_fraction': data.get('harness', {}).get('max_infra_error_fraction', 0.1) if data.get('harness') else None,  # 0.1 is the pipeline's default
                         **meta})
    return sorted(rows, key=lambda r: (r.get('status') != 'active', r['id']))


def compose(model_id, method_id, suite_ids, train, *, project, task_list_root='../../task-lists/'):
    """The pipeline config for one run. One suite is written as the legacy [eval_dataset] table, which
    every pinned pipeline reads; several are written as [[eval_suites]], which needs a method whose
    pipeline supports it (meta.multi_suite)."""
    model, method = fragment('models', model_id), fragment('methods', method_id)
    tables = {k: copy.deepcopy(v) for k, v in method.items() if k != 'meta' and k not in FILLED}
    if not tables: raise ValueError(f'method {method_id!r} has no pipeline tables yet')
    if not suite_ids: raise ValueError('at least one held-out suite is required')
    suites = [fragment('suites', s)['suite'] for s in suite_ids]
    if len({s['name'] for s in suites}) != len(suites): raise ValueError('held-out suite names must be unique')
    data = {'model': dict(model['model']), 'train_dataset': dict(train)}
    listed = lambda s: {'repo_id': s['repo_id'], 'revision': s['revision'], 'path': s.get('path', ''), 'task_list': task_list_root + s['task_list']}
    if len(suites) == 1:
        data['eval_dataset'] = listed(suites[0])
    else:
        if not method.get('meta', {}).get('multi_suite'): raise ValueError(f'method {method_id!r} evaluates a single suite')
        data['eval_suites'] = [{'name': s['name'], **listed(s)} for s in suites]
    data.update(tables)
    data.setdefault('tracking', {})['project'] = project
    data['output'] = {'root': '../../runs'}
    return data


# GPUs per HF job flavor this arena may launch on; add a flavor here before a model fragment uses it.
GPUS = {'a100-large': 1, 'a100x4': 4, 'a100x8': 8, 'h200': 1, 'h200x2': 2, 'h200x4': 4, 'h200x8': 8, 'l40sx1': 1, 'l40sx4': 4, 'l40sx8': 8}


def devices(value, what):
    """A CUDA device list ("6,7") as distinct integers; refuses blanks, non-numbers and repeats."""
    parts = [p.strip() for p in str(value).split(',')]
    if not parts or not all(p.isdigit() for p in parts): raise ValueError(f'{what} must list GPU indices like "6,7", got {value!r}')
    found = [int(p) for p in parts]
    if len(set(found)) != len(found): raise ValueError(f'{what} lists a GPU twice: {value!r}')
    return found


def serving(model_id, method_id, layout=None):
    """Job layout for one model and recipe: hardware flavor, vLLM GPUs and tensor parallelism, trainer GPUs, and
    the context caps (vLLM max_model_len, the bridge's trim and logprob limits). Read from [meta.serving] and checked
    so a bad layout fails when the challenge loads, never after compute is reserved. ``layout``: a challenge's own
    hardware ([compute.layout]: flavor, vllm_gpus, tensor_parallel, trainer_gpus) over the model fragment's."""
    model, method = fragment('models', model_id), fragment('methods', method_id)
    hw, ctx = {**model.get('meta', {}).get('serving', {}), **(layout or {})}, dict(method.get('meta', {}).get('serving', {}))
    missing = [k for k in ('flavor', 'vllm_gpus', 'tensor_parallel', 'trainer_gpus') if k not in hw] + [k for k in ('max_model_len', 'bridge_max_context', 'bridge_max_logprob_context') if k not in ctx]
    if missing: raise ValueError(f'serving layout incomplete for {model_id} x {method_id}: {", ".join(missing)}')
    if hw['flavor'] not in GPUS: raise ValueError(f'unknown hardware flavor {hw["flavor"]!r}; known: {", ".join(GPUS)}')
    vllm, trainer, count = devices(hw['vllm_gpus'], 'vllm_gpus'), devices(hw['trainer_gpus'], 'trainer_gpus'), GPUS[hw['flavor']]
    if any(g >= count for g in vllm + trainer): raise ValueError(f'{hw["flavor"]} has GPUs 0-{count - 1}; the layout uses {max(vllm + trainer)}')
    if not isinstance(hw['tensor_parallel'], int) or len(vllm) != hw['tensor_parallel']: raise ValueError('tensor_parallel must equal the number of vLLM GPUs')
    if set(vllm) & set(trainer): raise ValueError('vLLM and trainer GPUs overlap')
    if not ctx['bridge_max_logprob_context'] <= ctx['bridge_max_context'] <= ctx['max_model_len']:
        raise ValueError('context caps must satisfy bridge_max_logprob_context <= bridge_max_context <= max_model_len')
    completion = method.get('runtime', {}).get('max_completion_length')
    if completion is not None and completion > ctx['max_model_len']: raise ValueError('runtime.max_completion_length exceeds what vLLM serves (max_model_len)')
    hw.update(vllm_gpus=','.join(map(str, vllm)), trainer_gpus=','.join(map(str, trainer)))
    return {'model': model['model']['id'], 'revision': model['model']['revision'], **hw, **ctx}