"""Coding prompts for Beam from TACO-verified (MIT; training split of BAAI/TACO), decontaminated against
HumanEval+ and MBPP+.

  .venv/bin/python -I scripts/build_code_prompts.py
  # -> data/code/prompts.jsonl  (beam_selfhost.py prompt format, one sample per problem)
  #    data/code/tests.jsonl    (id -> fn_name, inputs, outputs; used only by modal_app/code_verify.py)

Problems are dropped if they share a 10-word sequence with any HumanEval+ or MBPP+ prompt, if their statement is
over 12k characters, or if their tests are over 1 MB. At most 30 tests are kept per problem.
"""

import json
import re
from pathlib import Path

import pyarrow.parquet as pq

ROOT = Path(__file__).resolve().parent.parent
RAW = ROOT / "data" / "code_raw"
OUT = ROOT / "data" / "code"
NGRAM = 10
MAX_TESTS = 30

STDIO = ("\n\nWrite a complete Python 3 program that reads the input from standard input and prints the answer to "
         "standard output. Put the whole program in a single ```python code block.")
FUNC = ("\n\nWrite the solution in Python 3, completing this starter code:\n\n```python\n{starter}\n```\n\n"
        "Put the complete solution in a single ```python code block.")


def words(text):
    return re.findall(r"[a-z0-9]+", text.lower())


def ngrams(text):
    w = words(text)
    return {" ".join(w[i:i + NGRAM]) for i in range(len(w) - NGRAM + 1)}


def main():
    bench = set()
    for path in RAW.glob("evalplus/*/data/*.parquet"):
        for row in pq.read_table(path).to_pylist():
            bench |= ngrams(row["prompt"])
    print(f"benchmark {NGRAM}-grams: {len(bench)}")

    OUT.mkdir(parents=True, exist_ok=True)
    kept = dropped_overlap = dropped_size = 0
    with open(OUT / "prompts.jsonl", "w") as fp, open(OUT / "tests.jsonl", "w") as ft:
        for r in json.load(open(RAW / "taco" / "taco_verified.json")):
            q = r["question"].strip()
            if ngrams(q) & bench:
                dropped_overlap += 1
                continue
            io = json.loads(r["input_output"])
            ins, outs = io.get("inputs", [])[:MAX_TESTS], io.get("outputs", [])[:MAX_TESTS]
            if len(q) > 12000 or not ins or len(json.dumps([ins, outs])) > 1_000_000:
                dropped_size += 1
                continue
            fn = io.get("fn_name")
            starter = (r.get("starter_code") or "").strip()
            if fn and not starter:
                dropped_size += 1
                continue
            pid = f"taco-{r['id']}"
            prompt = q + (FUNC.format(starter=starter) if fn else STDIO)
            fp.write(json.dumps({"id": pid, "prompt": prompt, "answer": "", "answer_type": "code", "domain": "code",
                                 "source": f"taco/{r['source']}", "difficulty": r["difficulty"], "topic": None},
                                ensure_ascii=False) + "\n")
            ft.write(json.dumps({"id": pid, "fn_name": fn, "inputs": ins, "outputs": outs}, ensure_ascii=False) + "\n")
            kept += 1
    print(f"kept {kept}, dropped {dropped_overlap} for benchmark overlap and {dropped_size} for size or format")


if __name__ == "__main__":
    main()