"""Verifiable unit-test rewards for coding RL."""
from __future__ import annotations

import subprocess
import sys
from pathlib import Path
from typing import Any


def run_hidden_tests(
    solution_code: str,
    hidden_tests: list[str],
    *,
    timeout_s: float = 5.0,
    python_exe: str | None = None,
) -> dict[str, Any]:
    """Execute each assert against solution_code; return pass fraction + details."""
    py = python_exe or sys.executable
    per_test: list[dict[str, Any]] = []
    passed = 0
    for i, assertion in enumerate(hidden_tests):
        script = f"{solution_code}\n{assertion}\n"
        try:
            proc = subprocess.run(
                [py, "-c", script],
                capture_output=True,
                text=True,
                timeout=timeout_s,
            )
            ok = proc.returncode == 0
            err = (proc.stderr or "")[:500]
        except Exception as e:  # noqa: BLE001
            ok = False
            err = f"{type(e).__name__}: {e}"
        per_test.append({"i": i, "ok": ok, "stderr": err})
        if ok:
            passed += 1
    total = len(hidden_tests)
    return {
        "pass_frac": (passed / total) if total else 0.0,
        "passed": passed,
        "total": total,
        "per_test": per_test,
    }


def extract_python_code(completion: str) -> str:
    """Pull first ```python fenced block, else raw text."""
    if "```python" in completion:
        start = completion.find("```python") + len("```python")
        end = completion.find("```", start)
        if end == -1:
            return completion[start:].strip()
        return completion[start:end].strip()
    if "```" in completion:
        start = completion.find("```") + 3
        end = completion.find("```", start)
        if end != -1:
            return completion[start:end].strip()
    return completion.strip()


def grade_completion(
    completion: str,
    hidden_tests: list[str],
    *,
    timeout_s: float = 5.0,
) -> dict[str, Any]:
    code = extract_python_code(completion)
    grade = run_hidden_tests(code, hidden_tests, timeout_s=timeout_s)
    grade["extracted_code"] = code
    return grade


def compute_reward(
    grade: dict[str, Any],
    *,
    format_ok: bool = True,
    ran_execute: bool = False,
) -> dict[str, float]:
    components = {
        "pass_frac": float(grade.get("pass_frac", 0.0)),
        "format_bonus": 0.05 if format_ok and grade.get("extracted_code") else 0.0,
        "tool_bonus": 0.05 if ran_execute else 0.0,
        "penalties": 0.0,
    }
    components["reward_total"] = (
        components["pass_frac"]
        + components["format_bonus"]
        + components["tool_bonus"]
        - components["penalties"]
    )
    return components


def grade_solution_file(
    solution_path: Path,
    hidden_tests: list[str],
    *,
    timeout_s: float = 5.0,
) -> dict[str, Any]:
    if not solution_path.exists():
        return {
            "pass_frac": 0.0,
            "passed": 0,
            "total": len(hidden_tests),
            "per_test": [],
            "error": "solution.py missing",
            "extracted_code": "",
        }
    code = solution_path.read_text(encoding="utf-8")
    grade = run_hidden_tests(code, hidden_tests, timeout_s=timeout_s)
    grade["extracted_code"] = code
    grade["error"] = None
    return grade