"""MultiPL-E Rust (mbpp-rs / humaneval-rs): prompts + compile-run verifier.

Cross-LANGUAGE transfer testbed: same problems as our Python MBPP/HumanEval,
in Rust. Verifier: rustc compile of [function + tests-main], run, exit code.
"""

import os
import re
import subprocess
import tempfile
from pathlib import Path

RUSTC = os.path.expanduser("~/.cargo/bin/rustc")
CODE_RE = re.compile(r"```(?:rust|rs)?\s*\n(.*?)```", re.S)

DIRECT_SUFFIX = ("\n\nWrite ONLY the complete Rust function (signature "
                 "included, no main) in a ```rust code block. No explanation.")
PLAN_SUFFIX = ("\n\nFirst write a very brief plan: at most 4 short bullet "
               "lines. Then write the complete Rust function (signature "
               "included, no main) in a ```rust code block.")


def rust_prompt(tok, item, suffix):
    msg = ("Complete this Rust function:\n\n```rust\n" + item["prompt"]
           + "\n```" + suffix)
    return tok.apply_chat_template([{"role": "user", "content": msg}],
                                   tokenize=False, add_generation_prompt=True)


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


def run_rust_tests(code, item, timeout=30):
    if not code:
        return False
    tests = item["tests"]
    if tests.lstrip().startswith("}"):  # designed to close an open fn body
        tests = tests.lstrip()[1:]
    program = code + "\n" + tests
    try:
        with tempfile.TemporaryDirectory() as td:
            src = Path(td) / "main.rs"
            src.write_text(program)
            c = subprocess.run([RUSTC, "-O", "--edition", "2021", "-o",
                                str(Path(td) / "prog"), str(src)],
                               capture_output=True, timeout=timeout)
            if c.returncode != 0:
                return False
            r = subprocess.run([str(Path(td) / "prog")], capture_output=True,
                               timeout=10)
        return r.returncode == 0
    except (subprocess.TimeoutExpired, OSError):
        return False


if __name__ == "__main__":
    item = {"prompt": "/// doubles x\nfn double(x: isize) -> isize {\n",
            "tests": "}\n\nfn main() {\n    assert_eq!(double(2), 4);\n}"}
    good = "fn double(x: isize) -> isize {\n    x * 2\n}"
    bad = "fn double(x: isize) -> isize {\n    x + 1\n}"
    assert run_rust_tests(good, item), "good should pass"
    assert not run_rust_tests(bad, item), "bad should fail"
    assert not run_rust_tests("garbage", item), "non-compiling should fail"
    print("rust verifier self-test ok")