"""STaR-style difficulty labeling of MBPP with the frozen base model.

Two passes per item: direct code generation vs plan-first-then-code, each
executed against MBPP's unit tests (sandboxed subprocess). Labels:
easy (direct passes), hard (plan-only passes), drop (neither). The model's
own passing code is kept as the training target (in-distribution supervision).
"""

import json
import os
import re
import subprocess
import sys
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import torch
from datasets import load_dataset

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import load_model  # noqa: E402

OUT = Path(os.environ.get("LOOP_OUT",
    Path(__file__).resolve().parent.parent / "results-loop"))
OUT.mkdir(exist_ok=True)

DIRECT_SUFFIX = ("\n\nWrite only the Python function in a ```python code "
                 "block. No explanation.")
PLAN_SUFFIX = ("\n\nFirst write a very brief plan: at most 4 short bullet "
               "lines, no headings, no math notation. Then write the complete "
               "Python function in a ```python code block.")

CODE_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.S)


def mbpp_prompt(tok, item, suffix):
    tests = "\n".join(item["test_list"])
    msg = (f"{item['text']}\nYour code should pass these tests:\n\n{tests}"
           f"{suffix}")
    return tok.apply_chat_template([{"role": "user", "content": msg}],
                                   tokenize=False, add_generation_prompt=True)


def extract_code(text):
    m = CODE_RE.findall(text)
    return m[-1].strip() if m else None


def run_tests(code, item, timeout=10):
    if not code:
        return False
    script = (item.get("test_setup_code") or "") + "\n" + code + "\n" + \
        "\n".join(item["test_list"])
    try:
        with tempfile.TemporaryDirectory() as td:
            r = subprocess.run([sys.executable, "-c", script], cwd=td,
                               capture_output=True, timeout=timeout)
        return r.returncode == 0
    except (subprocess.TimeoutExpired, OSError):
        return False


@torch.no_grad()
def batch_generate(model, tok, prompts, max_new_tokens, batch_size=24):
    outs = []
    for i in range(0, len(prompts), batch_size):
        chunk = prompts[i : i + batch_size]
        enc = tok(chunk, return_tensors="pt", padding=True,
                  add_special_tokens=False).to("cuda")
        gen = model.generate(**enc, max_new_tokens=max_new_tokens,
                             do_sample=False,
                             pad_token_id=tok.pad_token_id or 0)
        for j in range(len(chunk)):
            outs.append(tok.decode(gen[j, enc["input_ids"].shape[1]:],
                                   skip_special_tokens=True))
        print(f"  {min(i+batch_size, len(prompts))}/{len(prompts)}", flush=True)
    return outs


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"
    ds = load_dataset("google-research-datasets/mbpp", "full")

    items = []
    for split, tag in (("train", "train"), ("validation", "train"),
                       ("test", "test")):
        for row in ds[split]:
            items.append({"split": tag, "task_id": row["task_id"],
                          "text": row["text"], "test_list": row["test_list"],
                          "test_setup_code": row["test_setup_code"]})
    print(f"{len(items)} items", flush=True)

    for tag, suffix, mx in (("direct", DIRECT_SUFFIX, 220),
                            ("plan", PLAN_SUFFIX, 700)):
        t0 = time.time()
        print(f"{tag} pass", flush=True)
        gens = batch_generate(model, tok,
                              [mbpp_prompt(tok, it, suffix) for it in items], mx)
        codes = [extract_code(g) for g in gens]
        with ThreadPoolExecutor(8) as ex:
            oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]),
                              zip(codes, items)))
        for it, c, ok in zip(items, codes, oks):
            it[f"{tag}_ok"] = bool(ok)
            it[f"{tag}_code"] = c if ok else None
        print(f"{tag} pass done in {time.time()-t0:.0f}s  "
              f"pass@1={sum(oks)/len(items):.3f}", flush=True)

    for it in items:
        it["label"] = ("easy" if it["direct_ok"]
                       else "hard" if it["plan_ok"] else "drop")
        it["sol_code"] = it["direct_code"] if it["direct_ok"] else it["plan_code"]

    for split in ("train", "test"):
        sub = [it for it in items if it["split"] == split]
        n = len(sub)
        print(f"{split}: n={n}  direct={sum(i['direct_ok'] for i in sub)/n:.3f}  "
              f"plan={sum(i['plan_ok'] for i in sub)/n:.3f}  "
              f"easy={sum(i['label']=='easy' for i in sub)}  "
              f"hard={sum(i['label']=='hard' for i in sub)}  "
              f"drop={sum(i['label']=='drop' for i in sub)}", flush=True)

    with open(OUT / "mbpp_data.json", "w") as f:
        json.dump(items, f, indent=1)
    print("wrote", OUT / "mbpp_data.json")


if __name__ == "__main__":
    main()