"""Chat-aware HumanEval / MBPP for LLaDA-8B-Instruct.

WHY THIS EXISTS
lm-eval ships humaneval and mbpp as raw-COMPLETION tasks: the harness executes
`doc["prompt"] + response` (humaneval) or the response verbatim (mbpp, which has no
filter at all). An SFT chat model does not complete -- it answers. Measured on
LLaDA-8B-Instruct, 2026-08-18, dense, humaneval doc 0:

    "To implement the function `has_close_elements`, we need to iterate through the
     list of numbers and check each pair ... Here's the implementation of the
     function:\n\n```python\nfrom typing import List\n"

Two independent failures in that one string:
  1. Prose. Concatenated onto the function signature it is not valid Python, so every
     item raises SyntaxError -> pass@1 = 0 (which is exactly what the smoke run gave).
  2. Truncation. humaneval's stock until list is ["\nclass", "\ndef", "\n#", "\nif",
     "\nprint"] -- written for a model that completes a function body and then starts
     the next top-level statement. Here the model emits "```python\nfrom typing import
     List\n" and the NEXT thing it writes is "\ndef has_close_elements", so the stock
     stop sequence decapitates the answer inside its own code fence.

So the instruct variants (a) stop on <|eot_id|>, the actual end-of-turn marker, instead
of on Python keywords, and (b) run the response through the extractor below before it
reaches code_eval.

Note LLaDA generation is diffusion, not autoregressive: gen_length tokens are always
produced and `until` only post-trims the decoded string. Dropping the keyword stops
therefore costs nothing in compute -- it only stops us from cutting the answer.

FALLBACK POLICY: if no fenced block is found, fall back to the stock raw-completion
behaviour. That way a response that happens to be a bare completion still scores, and
this file can never do worse than the stock task.
"""
import re

import evaluate as hf_evaluate

try:
    _code_eval = hf_evaluate.load("code_eval")
    _code_eval.compute(references=["assert add(2, 3)==5"],
                       predictions=[["def add(a,b): return a*b"]], k=[1])
except Exception as e:
    raise e

# ```python ... ```  /  ```py ... ```  /  ``` ... ```
_FENCE = re.compile(r"```(?:python|py)?[ \t]*\r?\n(.*?)(?:```|\Z)", re.DOTALL)
# strip any special tokens that survived the harness decode
_SPECIAL = re.compile(r"<\|(?:eot_id|endoftext|startoftext|start_header_id|end_header_id)\|>")


def _fenced_blocks(text):
    return [b for b in (m.group(1) for m in _FENCE.finditer(text or "")) if b.strip()]


def _pick_block(text, must_contain=None):
    """Best fenced code block, or None.

    Prefers a block that defines the thing we are about to call; among equals prefers
    the LAST one, since a chat answer often shows a sketch first and the final version
    last. Falls back to the longest block when nothing matches must_contain.
    """
    blocks = _fenced_blocks(_SPECIAL.sub("", text))
    if not blocks:
        return None
    if must_contain:
        named = [b for b in blocks if must_contain in b]
        if named:
            return named[-1]
    defs = [b for b in blocks if re.search(r"^\s*(def|class|import|from)\s", b, re.M)]
    if defs:
        return defs[-1]
    return max(blocks, key=len)


def build_predictions_chat(resps, docs):
    """humaneval: full program to execute, per (doc, sample)."""
    out = []
    for resp, doc in zip(resps, docs):
        prompt, entry = doc["prompt"], doc.get("entry_point") or ""
        preds = []
        for r in resp:
            block = _pick_block(r, must_contain=f"def {entry}" if entry else None)
            if block and (not entry or f"def {entry}" in block):
                preds.append(block)          # fence already carries imports + def
            else:
                preds.append(prompt + _SPECIAL.sub("", r or ""))   # stock behaviour
        out.append(preds)
    return out


def extract_code_mbpp(resps, docs):
    """mbpp: the response IS the program (its reference is the assert list)."""
    out = []
    for resp, _doc in zip(resps, docs):
        preds = []
        for r in resp:
            block = _pick_block(r)
            preds.append(block if block else _SPECIAL.sub("", r or ""))
        out.append(preds)
    return out


def pass_at_k(references, predictions, k=None):
    assert k is not None
    if isinstance(k, int):
        k = [k]
    return _code_eval.compute(references=references, predictions=predictions, k=k)[0]


def pass_at_1(references, predictions):
    return _code_eval.compute(references=references, predictions=[predictions],
                              k=[1])[0]["pass@1"]


def list_fewshot_samples():
    """mbpp 3-shot exemplars -- byte-identical to lm_eval/tasks/mbpp/utils.py so the
    only difference between mbpp and mbpp_instruct is the chat handling."""
    from lm_eval.tasks.mbpp.utils import list_fewshot_samples as _stock
    return _stock()