#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
10_code_hard.py - harder coding benchmarks behind `07_domain_eval.py run --suite codehard --hard-n N` (N tasks per benchmark, nested prefixes):
  lcb    LiveCodeBench code_generation_lite, release 6: 175 AtCoder and LeetCode contest problems from 2025-01 to 2025-04 (newer than most training
         data), about 40 hidden tests each; the model writes a whole program (stdin/stdout) or a LeetCode class, every test must pass.
         The task id carries the difficulty: lcb_easy_..., lcb_medium_..., lcb_hard_...
  crux   CRUXEval (800 short Python functions), half output prediction (CRUXEval-O) and half input prediction (CRUXEval-I): code understanding,
         checked by running `assert f(...) == ...` on the model's answer.
Public data, no token.  Needs `pip install datasets` and network (the LiveCodeBench file is about 1 GB).  THE MODEL-WRITTEN CODE IS EXECUTED (see 07).
Use `--max-tokens 4096 --ctx 8192` for the lcb tasks (long answers).  Checks:
  python 10_code_hard.py selftest     lookup-table solutions pass every test of a few lcb problems, wrong ones fail; crux references pass
  python 10_code_hard.py validate     runs all 800 crux references and prints the ids to put in BAD_CRUX
"""
import base64
import io
import json
import os
import pickle
import random
import re
import subprocess
import sys
import tempfile
import time
import zlib
from pathlib import Path

LCB_FILES = ("test6",)          # release 6 (2025-01 .. 2025-04); add "test5" (2024-09 .. 2025-01) for a bigger pool
BAD_CRUX = set()                # crux ids whose own reference fails (filled by `validate`)
HEADER = ("from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\n"
          "from copy import deepcopy\nfrom math import *\nfrom typing import *\nfrom functools import *\nfrom itertools import *\nfrom operator import *\n"
          "from random import *\nfrom statistics import *\nimport sys, heapq, bisect, math, collections, itertools, functools, random, string, re, json, copy\n")
FUNC_RUNNER = '''import json, sys
ns = {}
exec(compile(open("sol.py", encoding="utf-8").read(), "sol.py", "exec"), ns)
f = getattr(ns["Solution"](), sys.argv[1])
for inp, out in json.load(open("tests.json", encoding="utf-8")):
    got = f(*[json.loads(l) for l in inp.split("\\n") if l.strip()])
    if json.loads(json.dumps(got)) != json.loads(out):
        print("FAIL")
        sys.exit(1)
print("OK")
'''


# ---------------------------------------------------------------------------------------------------------------------------------- LiveCodeBench
class _Plain(pickle.Unpickler):
    """The private tests are a pickled JSON string; a string needs no globals, so refuse every one (a dataset must not run code while loading)."""
    def find_class(self, module, name):
        raise pickle.UnpicklingError(f"global {module}.{name} refused")


def decode_private(s):
    return json.loads(_Plain(io.BytesIO(zlib.decompress(base64.b64decode(s)))).load())


def lcb_tests(row):
    return json.loads(row["public_test_cases"]) + decode_private(row["private_test_cases"])


def _same(got, want):
    g, w = [l.rstrip() for l in got.strip().splitlines()], [l.rstrip() for l in want.strip().splitlines()]
    if len(g) != len(w):
        return False
    for a, b in zip(g, w):
        if a == b:
            continue
        ta, tb = a.split(), b.split()
        if len(ta) != len(tb):
            return False
        for x, y in zip(ta, tb):
            if x == y:
                continue
            if re.fullmatch(r"-?\d+", x) and re.fullmatch(r"-?\d+", y):      # integers are exact
                return False
            try:
                if abs(float(x) - float(y)) > 1e-6 * max(1.0, abs(float(y))):
                    return False
            except ValueError:
                return False
    return True


def _limits():
    if os.name != "posix":
        return None
    import resource
    return lambda: resource.setrlimit(resource.RLIMIT_AS, (2 << 30, 2 << 30))


def run_lcb(code, tests, func, total=90, per=8):
    """True when `code` passes every test: func None = a stdin/stdout program (one process per test), else a LeetCode Solution method."""
    env = {**os.environ, "OPENBLAS_NUM_THREADS": "1", "OMP_NUM_THREADS": "1"}
    with tempfile.TemporaryDirectory() as d:
        d = Path(d)
        (d / "sol.py").write_text(HEADER + code, encoding="utf-8")
        t0 = time.time()
        if func is None:
            for t in tests:
                if time.time() - t0 > total:
                    return False
                try:
                    r = subprocess.run([sys.executable, "-I", "sol.py"], input=t["input"], capture_output=True, text=True, timeout=per, cwd=d, env=env,
                                       preexec_fn=_limits())
                except (subprocess.TimeoutExpired, UnicodeDecodeError):
                    return False
                if r.returncode != 0 or not _same(r.stdout, t["output"]):
                    return False
            return True
        (d / "tests.json").write_text(json.dumps([[t["input"], t["output"]] for t in tests]), encoding="utf-8")
        (d / "run.py").write_text(FUNC_RUNNER, encoding="utf-8")
        try:
            r = subprocess.run([sys.executable, "-I", "run.py", func], capture_output=True, text=True, timeout=total, cwd=d, env=env, preexec_fn=_limits())
        except (subprocess.TimeoutExpired, UnicodeDecodeError):
            return False
        return r.returncode == 0 and r.stdout.strip().endswith("OK")


FUNC_DEBUG = '''import json, sys, traceback
ns = {}
exec(compile(open("sol.py", encoding="utf-8").read(), "sol.py", "exec"), ns)
f = getattr(ns["Solution"](), sys.argv[1])
for inp, out in json.load(open("tests.json", encoding="utf-8")):
    try:
        got = f(*[json.loads(l) for l in inp.split("\\n") if l.strip()])
    except Exception:
        print("Input:\\n" + inp.strip()[:300] + "\\nRuntime error:\\n" + traceback.format_exc()[-600:])
        sys.exit(1)
    if json.loads(json.dumps(got)) != json.loads(out):
        print("Input:\\n" + inp.strip()[:300] + "\\nExpected output:\\n" + out.strip()[:300] + "\\nYour output:\\n" + json.dumps(got)[:300])
        sys.exit(1)
print("OK")
'''


def first_failure(code, tests, func, per=8):
    """None when `code` passes every test in `tests` (the VISIBLE sample tests), else a short text on the first failure, for a repair prompt."""
    env = {**os.environ, "OPENBLAS_NUM_THREADS": "1", "OMP_NUM_THREADS": "1"}
    with tempfile.TemporaryDirectory() as d:
        d = Path(d)
        (d / "sol.py").write_text(HEADER + code, encoding="utf-8")
        if func is None:
            for t in tests:
                shown = "Input:\n" + t["input"].strip()[:300]
                try:
                    r = subprocess.run([sys.executable, "-I", "sol.py"], input=t["input"], capture_output=True, text=True, timeout=per, cwd=d, env=env,
                                       preexec_fn=_limits())
                except subprocess.TimeoutExpired:
                    return f"{shown}\nThe program did not finish within {per} seconds."
                except UnicodeDecodeError:
                    return f"{shown}\nThe program printed bytes that are not valid text."
                if r.returncode != 0:
                    return f"{shown}\nRuntime error:\n{r.stderr[-600:]}"
                if not _same(r.stdout, t["output"]):
                    return f"{shown}\nExpected output:\n{t['output'].strip()[:300]}\nYour output:\n{r.stdout.strip()[:300]}"
            return None
        (d / "tests.json").write_text(json.dumps([[t["input"], t["output"]] for t in tests]), encoding="utf-8")
        (d / "run.py").write_text(FUNC_DEBUG, encoding="utf-8")
        try:
            r = subprocess.run([sys.executable, "-I", "run.py", func], capture_output=True, text=True, timeout=per * 4, cwd=d, env=env, preexec_fn=_limits())
        except subprocess.TimeoutExpired:
            return "The code did not finish in time on the sample tests."
        except UnicodeDecodeError:
            return "The code printed bytes that are not valid text."
        if r.returncode == 0 and r.stdout.strip().endswith("OK"):
            return None
        return (r.stdout or r.stderr).strip()[-800:] or "The code crashed."


def lcb_feedback(resp, row, func):
    """What to tell the model after a failed attempt: only the visible sample tests are used (never the hidden ones); None = no new information."""
    code = last_block(resp)
    if not code.strip():
        return "I could not find a ```python code block in your answer. Give the complete solution in one ```python block."
    bad = first_failure(code, json.loads(row["public_test_cases"]), func)
    if bad is None:
        return None
    return "Your code fails this sample test.\n\n" + bad + "\n\nFix the code. Reply with the complete corrected solution in one ```python block."


def last_block(resp):
    blocks = re.findall(r"```(?:python|py|Python)?[ \t]*\n(.*?)```", resp, flags=re.S)
    return blocks[-1] if blocks else ""


def lcb_prompt(row):
    if row["starter_code"]:
        fmt = "You will use the following starter code to write the solution to the problem and enclose your code within delimiters."
        block = row["starter_code"]
    else:
        fmt = ("Read the inputs from stdin solve the problem and write the answer to stdout (do not directly test on the sample inputs). Enclose your "
               "code within delimiters as follows. Ensure that when the python program runs, it reads the inputs, runs the algorithm and writes output "
               "to STDOUT.")
        block = "# YOUR CODE HERE"
    return f"### Question:\n{row['question_content']}\n\n### Format: {fmt}\n```python\n{block}\n```\n\n### Answer: (use the provided format with backticks)\n"


def lcb_rows():
    from datasets import load_dataset
    rows = []
    for f in LCB_FILES:
        rows += list(load_dataset("json", data_files=f"hf://datasets/livecodebench/code_generation_lite/{f}.jsonl", split="train"))
    return rows


def pool_lcb(rnd):
    rows = lcb_rows()
    rnd.shuffle(rows)
    T = []
    for r in rows:
        func = json.loads(r["metadata"] or "{}").get("func_name") if r["starter_code"] else None

        def check(resp, r=r, func=func, cache={}):
            code = last_block(resp)
            if not code.strip():
                return False
            if r["question_id"] not in cache:
                cache.clear()
                cache[r["question_id"]] = lcb_tests(r)
            return run_lcb(code, cache[r["question_id"]], func)
        check.feedback = lambda resp, r=r, func=func: lcb_feedback(resp, r, func)     # used by `07 run --retry N`
        T.append((f"lcb_{r['difficulty']}_{r['question_id']}", "lcb", lcb_prompt(r), check, None))
    return T


# ---------------------------------------------------------------------------------------------------------------------------------- CRUXEval
def pool_crux(rnd, run_code):
    from datasets import load_dataset
    rows = [r for r in load_dataset("cruxeval-org/cruxeval", split="test")]
    ro, ri = rows[:], rows[:]
    rnd.shuffle(ro)
    rnd.shuffle(ri)
    O, I = [], []
    for r in ro:
        if f"crux_o_{r['id']}" in BAD_CRUX:
            continue
        pro = ("You are given a Python function and an assertion containing an input to the function. Complete the assertion with a literal (no "
               "unsimplified expressions, no function calls) containing the output when executing the provided code on the given input, even if the "
               "function is incorrect or incomplete. Think step by step, then end with exactly one line in this form: [ANSWER] assert f(...) == <literal> "
               f"[/ANSWER]\n\n```python\n{r['code']}\nassert f({r['input']}) == ??\n```")

        def check_o(resp, r=r):
            m = re.findall(r"\[ANSWER\]\s*assert\s+f\(.*?\)\s*==\s*(.*?)\s*\[/ANSWER\]", resp, flags=re.S)
            return bool(m) and run_code(r["code"], f"assert f({r['input']}) == ({m[-1]})", 10)
        O.append((f"crux_o_{r['id']}", "crux", pro, check_o, f"[ANSWER] assert f({r['input']}) == {r['output']} [/ANSWER]"))
    for r in ri:
        if f"crux_i_{r['id']}" in BAD_CRUX:
            continue
        pri = ("You are given a Python function and an assertion containing an output of the function. Find any input such that executing the function "
               "on the input leads to the given output. There may be several answers, give only one. Think step by step, then end with exactly one line "
               f"in this form: [ANSWER] assert f(<your input>) == {r['output']} [/ANSWER]\n\n```python\n{r['code']}\nassert f(??) == {r['output']}\n```")

        def check_i(resp, r=r):
            m = re.findall(r"\[ANSWER\]\s*assert\s+f\((.*)\)\s*==", resp, flags=re.S)
            return bool(m) and run_code(r["code"], f"assert f({m[-1]}) == {r['output']}", 10)
        I.append((f"crux_i_{r['id']}", "crux", pri, check_i, f"[ANSWER] assert f({r['input']}) == {r['output']} [/ANSWER]"))
    mixed, i = [], 0
    while i < min(len(O), len(I)):
        mixed += [O[i], I[i]]
        i += 1
    return mixed


# ---------------------------------------------------------------------------------------------------------------------------------- the suite
def build_codehard(seed, n, run_code, code_of, num):
    """[(id, category, prompt, check(response) -> bool, reference response or None)]: the first n lcb problems and the first n crux tasks."""
    T = pool_lcb(random.Random(f"{seed}lcb"))[:n] + pool_crux(random.Random(f"{seed}crux"), run_code)[:n]
    assert len({t[0] for t in T}) == len(T), "duplicate task ids"
    return T


def _lookup_solution(row, tests):
    """A 'solution' that only looks the answers up: it passes iff the runner, the decoding and the comparison are right."""
    if not row["starter_code"]:
        return "import sys\nD = " + repr({t["input"]: t["output"] for t in tests}) + "\nprint(D[sys.stdin.read()], end='')\n", None
    func = json.loads(row["metadata"])["func_name"]
    D = {json.dumps([json.loads(l) for l in t["input"].split("\n") if l.strip()]): t["output"] for t in tests}
    return f"import json\nD = {D!r}\nclass Solution:\n    def {func}(self, *a):\n        return json.loads(D[json.dumps(list(a))])\n", func


if __name__ == "__main__" and sys.argv[1:] in (["selftest"], ["validate"]):
    import importlib
    sys.path.insert(0, Path(__file__).resolve().parent.as_posix())
    m = importlib.import_module("07_domain_eval")
    if sys.argv[1] == "selftest":
        rows = lcb_rows()
        stdin, func = [r for r in rows if not r["starter_code"]][:4], [r for r in rows if r["starter_code"]][:4]
        for r in stdin + func:
            tests = lcb_tests(r)
            code, fn = _lookup_solution(r, tests)
            assert run_lcb(code, tests, fn), ("lookup solution failed", r["question_id"])
            wrong = "print(0)\n" if fn is None else f"class Solution:\n    def {fn}(self, *a):\n        return None\n"
            assert not run_lcb(wrong, tests, fn), ("wrong solution passed", r["question_id"])
        assert last_block("text\n```python\nA\n```\nmore\n```python\nB\n```") == "B\n"
        for r in stdin + func:          # the repair feedback: a correct answer needs none, wrong / crashing / empty answers get a message
            tests = json.loads(r["public_test_cases"])
            good, fn = _lookup_solution(r, tests)
            assert lcb_feedback("```python\n" + good + "```", r, fn) is None, ("feedback for a correct solution", r["question_id"])
            bad_code = "print(0)\n" if fn is None else f"class Solution:\n    def {fn}(self, *a):\n        return None\n"
            crash = "raise ValueError('boom')\n" if fn is None else f"class Solution:\n    def {fn}(self, *a):\n        raise ValueError('boom')\n"
            assert "sample test" in (lcb_feedback("```python\n" + bad_code + "```", r, fn) or ""), ("no feedback for a wrong answer", r["question_id"])
            assert "boom" in (lcb_feedback("```python\n" + crash + "```", r, fn) or ""), ("no error text for a crash", r["question_id"])
            assert "could not find" in lcb_feedback("no code here", r, fn)
        T = build_codehard(1, 10, m.run_code, m.code_of, m.num)
        crux = [t for t in T if t[1] == "crux"]
        bad = [t[0] for t in crux if not t[3](t[4])]
        junk = [t[0] for t in T if t[1] == "crux" and t[3]("zzz 987654321")] + [t[0] for t in T if t[1] == "lcb" and t[3]("zzz 987654321")]
        per = {c: sum(t[1] == c for t in T) for c in ("lcb", "crux")}
        assert not bad and not junk and per == {"lcb": 10, "crux": 10}, (bad, junk, per)
        print(f"selftest OK: {len(T)} tasks {per}; {len(stdin)} stdin + {len(func)} functional lcb problems pass their lookup solution and fail a wrong one; "
              f"crux references pass, garbage fails")
    else:
        from datasets import load_dataset
        bad = []
        for r in load_dataset("cruxeval-org/cruxeval", split="test"):
            if not m.run_code(r["code"], f"assert f({r['input']}) == ({r['output']})", 10):
                bad += [f"crux_o_{r['id']}", f"crux_i_{r['id']}"]
        print(f"{len(bad) // 2} crux references fail: BAD_CRUX = {set(bad)!r}")
        n_ok = sum(1 for r in lcb_rows() if lcb_tests(r))
        print(f"lcb: {n_ok} problems decode their tests")