"""Public tests a programmer can see in the task statement itself (never the hidden grading tests):
  code_io   -> Input/Output example pairs of the statement (Codeforces / AtCoder / "Sample Input" styles)
  humaneval -> doctest lines (>>> call / expected value) of the docstring
  mbpp      -> the assert lines quoted in the prompt
run_public(task, code, prompt) -> (n_pass, n_total, first_failure_text)"""
import os
import re
import subprocess

import sandbox  # noqa: E402  (same folder)
import sys
import tempfile

TMP = os.path.join(tempfile.gettempdir(), "lfm3_pub")
os.makedirs(TMP, exist_ok=True)

_HDR_IN = re.compile(r"^\s*(?:-+\s*)?(?:sample|example)?\s*input\s*(?:\d+|#\d+)?\s*:?\s*(?:-+)?\s*$", re.I)
_HDR_OUT = re.compile(r"^\s*(?:-+\s*)?(?:sample|example)?\s*output\s*(?:\d+|#\d+)?\s*:?\s*(?:-+)?\s*$", re.I)
_STOP = re.compile(r"^\s*(?:-{3,}.*|note:?|notes:?|explanation:?.*|constraints:?|write a python 3 program.*)\s*$", re.I)


_EX_HDR = re.compile(r"^\s*-*\s*(?:examples?|samples?|sample tests?)\s*\d*\s*:?\s*-*\s*$", re.I)


def _prose(text):
    return len(re.findall(r"[A-Za-z]{3,}", text)) >= 12


def io_examples(prompt):
    """Example pairs after the statement's Example(s) header (format sections named Input/Output are skipped)."""
    lines = prompt.replace("\r", "").split("\n")
    pairs = []
    i = next((k + 1 for k, l in enumerate(lines) if _EX_HDR.match(l)), 0)
    while i < len(lines):
        if _HDR_IN.match(lines[i]) and not re.match(r"^\s*-+\s*input\s*-+\s*$", lines[i], re.I):
            j = i + 1
            inp = []
            while j < len(lines) and not _HDR_OUT.match(lines[j]):
                if _HDR_IN.match(lines[j]) or _STOP.match(lines[j]):
                    break
                inp.append(lines[j])
                j += 1
            if j < len(lines) and _HDR_OUT.match(lines[j]):
                k = j + 1
                out = []
                while k < len(lines) and not _HDR_IN.match(lines[k]) and not _STOP.match(lines[k]):
                    out.append(lines[k])
                    k += 1
                for b in range(1, len(out)):  # an explanation paragraph after the expected output is not output
                    if not out[b].strip() and len(re.findall(r"[A-Za-z]{3,}", " ".join(out[b + 1:b + 4]))) >= 4:
                        out = out[:b]
                        break
                inp_s, out_s = "\n".join(inp).strip("\n"), "\n".join(out).strip()
                if inp_s.strip() and out_s and not _prose(inp_s):
                    pairs.append((inp_s + "\n", out_s))
                i = k
                continue
        i += 1
    return pairs


def doctests(prompt):
    tests, lines = [], prompt.split("\n")
    for i, l in enumerate(lines):
        s = l.strip()
        if s.startswith(">>>") and i + 1 < len(lines):
            exp = lines[i + 1].strip()
            if exp and not exp.startswith(">>>") and not exp.startswith('"""'):
                tests.append((s[3:].strip(), exp))
    return tests


def asserts(prompt):
    return [l.strip() for l in prompt.split("\n") if l.strip().startswith("assert ")]


def _run(src, stdin="", timeout=6):
    fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        f.write(src)
    try:
        r = sandbox.run([sys.executable, "-I", path], input=stdin.encode(), capture_output=True, timeout=timeout,
                           cwd=TMP, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        return r.returncode, r.stdout.decode(errors="replace"), r.stderr.decode(errors="replace")
    except subprocess.TimeoutExpired:
        return -1, "", "TimeoutError: time limit exceeded (6 s)"
    finally:
        try:
            os.remove(path)
        except OSError:
            pass


def _same(got, exp):
    a, b = got.split(), exp.split()
    if len(a) != len(b):
        return False
    for x, y in zip(a, b):
        if x != y:
            try:
                if abs(float(x) - float(y)) > 1e-6 * max(1.0, abs(float(y))):
                    return False
            except ValueError:
                return False
    return True


def _tail(err, n=6):
    return "\n".join(err.strip().splitlines()[-n:])


def run_public(task, code, prompt, item=None):
    """Returns (passed, total, failure description or None)."""
    if task == "lcb":
        from lcb_tests import run_lcb
        return run_lcb(code, item, item["public"])
    if task == "code_io":
        tests = io_examples(prompt)
        fails = None
        ok = 0
        for k, (inp, out) in enumerate(tests):
            rc, so, se = _run(code, inp)
            if rc == 0 and _same(so, out):
                ok += 1
            elif fails is None:
                got = so.strip()[:600] if rc == 0 else (so.strip()[:300] + "\n" + _tail(se)).strip()
                fails = f"Example {k + 1}\nInput:\n{inp.strip()[:800]}\nExpected output:\n{out[:600]}\n" \
                        f"Your program's output:\n{got or '(nothing)'}"
        return ok, len(tests), fails
    if task == "humaneval":
        tests = doctests(prompt)
        checks = [f"_r = {call}\n_e = {exp}\nprint('PASS' if _r == _e else 'FAIL ' + repr(_r))" for call, exp in tests]
    elif task == "mbpp":
        tests = asserts(prompt)
        checks = [f"try:\n    {a}\n    print('PASS')\nexcept AssertionError:\n    print('FAIL')" for a in tests]
    else:
        return 0, 0, None
    ok, fails = 0, None
    for k, (t, chk) in enumerate(zip(tests, checks)):
        rc, so, se = _run(code + "\n\n" + chk)
        if rc == 0 and so.strip().endswith("PASS"):
            ok += 1
        elif fails is None:
            shown = t if isinstance(t, str) else f">>> {t[0]}\nexpected: {t[1]}"
            res = so.strip().splitlines()[-1] if rc == 0 and so.strip() else _tail(se)
            fails = f"Test:\n{shown}\nResult: {res}"
    return ok, len(tests), fails