"""Problem suites: HumanEval and MBPP.

A *suite* is a set of per-language bindings. Each Binding tells the harness the
three things the dataset's *format* dictates:

  • build_prompt(row)                    -> text fed to the model (raw continuation)
  • stop_tokens(row)                     -> where to cut the completion
  • assemble(prompt, completion, row)    -> runnable program (exit 0 == all tests pass)

The language's *execution* side (filename, run command, podman image) stays in
langs.py and is shared across suites — only the problem source/format changes
here. Adding a suite is a single table entry, mirroring how adding a language is
a single entry in langs.py.

Two dataset formats appear:

  • MultiPL-E (nuprl/MultiPL-E, every `*-{js,go,rs,sh}` config, both suites): the
    `prompt` field is a runnable function *prefix*, so the program is
    prompt+completion+tests verbatim (_concat) and the model is simply asked to
    continue `prompt`.
  • MBPP-Python (google-research-datasets/mbpp): the problem is a natural-language
    `text` description plus a `test_list` of asserts — NOT a code prefix. So the
    generation prompt is synthesised (description + example asserts as a docstring)
    and the program is completion+setup+asserts (the NL prompt is discarded).
  • HumanEval-Python (openai/openai_humaneval): native HumanEval, prompt is a
    function prefix, tests are a `check(fn)` def (langs._python).
"""

from __future__ import annotations

import ast
from dataclasses import dataclass
from typing import Callable

from .langs import LANGS, _concat, resolve


@dataclass(frozen=True)
class Binding:
    dataset: str
    config: str
    build_prompt: Callable[[dict], str]
    assemble: Callable[[str, str, dict], str]
    stop_tokens: Callable[[dict], list[str]]


# ── shared format helpers ────────────────────────────────────────────────────

def _prompt_field(row: dict) -> str:
    """MultiPL-E / HumanEval: the runnable function prefix is the prompt."""
    return row["prompt"]


def _parse_stop(row: dict) -> list[str]:
    """MultiPL-E ships stop_tokens as a list or a stringified list."""
    raw = row.get("stop_tokens")
    if isinstance(raw, list):
        return raw
    if isinstance(raw, str):
        try:
            v = ast.literal_eval(raw)
            return v if isinstance(v, list) else []
        except (ValueError, SyntaxError):
            return []
    return []


# ── MBPP-Python format (description + asserts, no code prefix) ────────────────

# Cut as soon as the model leaves the function body for its own tests/output, so
# only the candidate implementation survives into the assembled program.
_MBPP_PY_STOP = ["\nassert", "\nprint(", "\n# Test", "\nif __name__", '\n"""']


def _mbpp_py_prompt(row: dict) -> str:
    """Frame the NL task + example asserts as a leading docstring so a raw
    (non-chat) instruct model continues with the function definition."""
    asserts = "\n".join(row.get("test_list", []))
    return f'"""\n{row.get("text", "").strip()}\n\n{asserts}\n"""\n'


def _mbpp_py_assemble(prompt: str, completion: str, row: dict) -> str:
    """Program = setup + the model's code + the held-out asserts. The NL prompt
    is *not* part of the program (unlike HumanEval, where it is the prefix)."""
    setup = row.get("test_setup_code", "") or ""
    asserts = "\n".join(row.get("test_list", []))
    return f"{setup}\n{completion}\n{asserts}\n"


def _mbpp_py_stop(row: dict) -> list[str]:
    return _MBPP_PY_STOP


# ── suite tables ─────────────────────────────────────────────────────────────

def _humaneval_binding(lang_id: str) -> Binding:
    """HumanEval reuses each language's existing langs.py dataset/config/assemble
    (the registry already encodes the HumanEval format)."""
    L = LANGS[lang_id]
    return Binding(L.dataset, L.config, _prompt_field, L.assemble, _parse_stop)


_MBPP: dict[str, Binding] = {
    "python": Binding("google-research-datasets/mbpp", "full",
                      _mbpp_py_prompt, _mbpp_py_assemble, _mbpp_py_stop),
    "javascript": Binding("nuprl/MultiPL-E", "mbpp-js",
                          _prompt_field, _concat, _parse_stop),
    "go": Binding("nuprl/MultiPL-E", "mbpp-go",
                  _prompt_field, _concat, _parse_stop),
    "rust": Binding("nuprl/MultiPL-E", "mbpp-rs",
                    _prompt_field, _concat, _parse_stop),
    "bash": Binding("nuprl/MultiPL-E", "mbpp-sh",
                    _prompt_field, _concat, _parse_stop),
}

SUITES = {
    "humaneval": "HumanEval (function-completion)",
    "mbpp": "MBPP (Mostly Basic Programming Problems)",
}


def binding(suite: str, language: str) -> Binding:
    """Resolve the (suite, language) -> Binding the harness should run."""
    lang_id = resolve(language).id
    s = suite.lower()
    if s == "humaneval":
        return _humaneval_binding(lang_id)
    if s == "mbpp":
        if lang_id not in _MBPP:
            raise KeyError(f"suite 'mbpp' has no binding for {lang_id!r}; "
                           f"have: {', '.join(_MBPP)}")
        return _MBPP[lang_id]
    raise KeyError(f"unknown suite {suite!r}; known: {', '.join(SUITES)}")