"""Converts LiveCodeBench code_generation_lite (test6.jsonl, problems from 2025) into eval items.

Hidden tests are stored pickled+zlib+base64 in the source; they are decoded with a restricted unpickler that refuses
every global (only plain str/list/dict data can load), so the file cannot execute code.
item: {id, task: "lcb", testtype: stdin|functional, prompt, func_name, public: [[in,out]], hidden: [[in,out]]}
usage: python -I prep_lcb.py <raw test6.jsonl> <out items.jsonl>
"""
import base64
import io
import json
import pickle
import sys
import zlib


class SafeUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        raise pickle.UnpicklingError(f"global {module}.{name} refused")


def decode_private(s):
    try:
        return json.loads(s)
    except Exception:
        raw = zlib.decompress(base64.b64decode(s.encode("utf-8")))
        return json.loads(SafeUnpickler(io.BytesIO(raw)).load())


STDIN_TAIL = ("\n\nWrite a Python 3 program that reads from standard input and writes to standard output. "
              "Put the complete program in a single ```python code block.")
FUNC_TAIL = ("\n\nComplete the following starter code and return the full solution (the class with the method) "
             "in a single ```python code block.\n\n```python\n{starter}\n```")

src, dst = sys.argv[1], sys.argv[2]
n = 0
with open(src, encoding="utf-8") as f, open(dst, "w", encoding="utf-8") as out:
    for line in f:
        r = json.loads(line)
        pub = json.loads(r["public_test_cases"])
        hid = decode_private(r["private_test_cases"])
        ttype = pub[0]["testtype"] if pub else hid[0]["testtype"]
        meta = json.loads(r.get("metadata") or "{}")
        starter = r.get("starter_code") or ""
        prompt = r["question_content"].strip() + (FUNC_TAIL.format(starter=starter.strip()) if ttype == "functional"
                                                  else STDIN_TAIL)
        out.write(json.dumps({"id": f"lcb/{r['question_id']}", "task": "lcb", "testtype": ttype,
                              "difficulty": r.get("difficulty"), "date": r.get("contest_date"),
                              "func_name": meta.get("func_name"), "prompt": prompt,
                              "public": [[t["input"], t["output"]] for t in pub],
                              "hidden": [[t["input"], t["output"]] for t in pub + hid]}, ensure_ascii=False) + "\n")
        n += 1
print(n, "problems")