"""Subprocess-based code testing utility for RL reward computation.

Supports both stdin-based (Codeforces/AtCoder style) and call-based
(LeetCode functional style) test execution.

Uses subprocess.run instead of multiprocessing.Process to work reliably
inside serverless / daemon worker environments.
"""
from __future__ import annotations

import json
import subprocess
import sys
import textwrap
import zlib
import pickle
import base64
from typing import Optional, Tuple

_DEFAULT_TIMEOUT = 10
_MAX_TEST_CASES = 20

# --- v2: display bounds for format_test_results --------------------------- #
# Hard upper-bound on per-case tool output footprint, regardless of how big
# the underlying actual/expected strings are. This decouples test-case
# complexity from rollout context size (avoids "samples too long → fallback
# multimodal sample" path that broke training).
_INPUT_DISPLAY_MAX = 300
_OUTPUT_HEAD = 300
_OUTPUT_TAIL = 100
_PASS_OUTPUT_HEAD = 200
_PASS_OUTPUT_TAIL = 50


def extract_code(response: str) -> Optional[str]:
    """Extract Python code from the last ```python``` fenced block."""
    if "```python" in response:
        code = response.split("```python")[-1].split("```")[0].strip()
        return code or None
    parts = response.split("```")
    if len(parts) >= 3:
        code = parts[-2].strip()
        return code or None
    return None


def _compare_output(actual: str, expected: str) -> bool:
    """Compare two output strings with tolerance for whitespace differences."""
    actual_lines = [l.rstrip() for l in actual.strip().splitlines()]
    expected_lines = [l.rstrip() for l in expected.strip().splitlines()]
    actual_lines = [l for l in actual_lines if l]
    expected_lines = [l for l in expected_lines if l]

    if actual_lines == expected_lines:
        return True

    if len(actual_lines) != len(expected_lines):
        return False

    for a, e in zip(actual_lines, expected_lines):
        if a == e:
            continue
        try:
            if abs(float(a) - float(e)) < 1e-6:
                continue
        except (ValueError, OverflowError):
            pass
        return False
    return True


class _TestResult:
    """Detailed result of a single test case execution."""
    __slots__ = ("passed", "actual", "expected", "error", "timed_out")

    def __init__(self, *, passed: bool, actual: str = "", expected: str = "",
                 error: str = "", timed_out: bool = False):
        self.passed = passed
        self.actual = actual
        self.expected = expected
        self.error = error
        self.timed_out = timed_out


def _run_stdin_test(
    code: str, test_input: str, expected_output: str, timeout: int
) -> bool:
    """Run code with stdin, compare stdout with expected output."""
    r = _run_stdin_test_detailed(code, test_input, expected_output, timeout)
    return r.passed


def _run_stdin_test_detailed(
    code: str, test_input: str, expected_output: str, timeout: int
) -> _TestResult:
    try:
        result = subprocess.run(
            [sys.executable, "-c", code],
            input=test_input,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        if result.returncode != 0:
            stderr = result.stderr.strip()[-500:] if result.stderr else ""
            return _TestResult(passed=False, expected=expected_output.strip(),
                               actual="", error=f"Runtime Error:\n{stderr}")
        ok = _compare_output(result.stdout, expected_output)
        return _TestResult(passed=ok, expected=expected_output.strip(),
                           actual=result.stdout.strip())
    except subprocess.TimeoutExpired:
        return _TestResult(passed=False, expected=expected_output.strip(),
                           error="Time Limit Exceeded", timed_out=True)
    except Exception as e:
        return _TestResult(passed=False, expected=expected_output.strip(),
                           error=f"Execution Error: {e}")


def _run_functional_test(
    code: str, fn_name: str, test_input: str, expected_output: str, timeout: int
) -> bool:
    r = _run_functional_test_detailed(code, fn_name, test_input,
                                       expected_output, timeout)
    return r.passed


def _run_functional_test_detailed(
    code: str, fn_name: str, test_input: str, expected_output: str, timeout: int
) -> _TestResult:
    """Run call-based test with detailed result."""
    input_lines = [l for l in test_input.strip().split("\n") if l.strip()]
    args_json = ", ".join(input_lines)

    wrapper = (
        "import json, sys\n"
        "from typing import *\n"
        f"{code}\n"
        f"_args = [{args_json}]\n"
        "try:\n"
        f'    _sol = Solution()\n'
        f'    _result = getattr(_sol, "{fn_name}")(*_args)\n'
        "except Exception as _e:\n"
        '    print("__ERROR__:" + repr(_e), file=sys.stderr)\n'
        "    sys.exit(1)\n"
        "print(json.dumps(_result))\n"
    )
    try:
        result = subprocess.run(
            [sys.executable, "-c", wrapper],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        if result.returncode != 0:
            stderr = result.stderr.strip()[-500:] if result.stderr else ""
            return _TestResult(passed=False, expected=expected_output.strip(),
                               actual="", error=f"Runtime Error:\n{stderr}")
        actual = json.loads(result.stdout.strip())
        expected = json.loads(expected_output.strip())
        ok = actual == expected
        return _TestResult(passed=ok, expected=expected_output.strip(),
                           actual=result.stdout.strip())
    except subprocess.TimeoutExpired:
        return _TestResult(passed=False, expected=expected_output.strip(),
                           error="Time Limit Exceeded", timed_out=True)
    except json.JSONDecodeError:
        return _TestResult(passed=False, expected=expected_output.strip(),
                           actual=result.stdout.strip() if result else "",
                           error="Output is not valid JSON")
    except Exception as e:
        return _TestResult(passed=False, expected=expected_output.strip(),
                           error=f"Execution Error: {e}")


def run_tests(
    code: str,
    test_cases: dict,
    fn_name: Optional[str] = None,
    timeout: int = _DEFAULT_TIMEOUT,
    max_cases: int = _MAX_TEST_CASES,
) -> Tuple[float, dict]:
    """Run code against test cases and return (score, metadata).

    score = passed / total (proportional).
    Caller can derive pass@1 from metadata["passed"] == metadata["total"].
    """
    inputs = test_cases.get("inputs", [])
    outputs = test_cases.get("outputs", [])
    total = min(len(inputs), len(outputs), max_cases)

    if total == 0:
        return 0.0, {"error": "no_test_cases"}

    passed = 0
    first_failure = None

    for i in range(total):
        if fn_name:
            ok = _run_functional_test(code, fn_name, inputs[i], outputs[i], timeout)
        else:
            ok = _run_stdin_test(code, inputs[i], outputs[i], timeout)

        if ok:
            passed += 1
        elif first_failure is None:
            first_failure = i

    score = passed / total
    meta = {"total": total, "passed": passed}
    if first_failure is not None:
        meta["first_failure_index"] = first_failure
    return score, meta


# ---------------------------------------------------------------------------
#  Helpers for normalising test-case formats from different data sources
# ---------------------------------------------------------------------------

def normalise_test_cases(rollout_extra: dict) -> Tuple[dict, Optional[str]]:
    """Extract and normalise test cases from rollout_extra.

    Supports:
      - deepcoder format:  rollout_extra["test_cases"] = JSON string of
            {"inputs": [...], "outputs": [...]}
      - livecodebench format: rollout_extra["public_test_cases"] (JSON string)
            + rollout_extra["private_test_cases"] (base64+zlib+pickle encoded)

    Returns (test_cases_dict, fn_name_or_None).
    """
    if "test_cases" in rollout_extra:
        tc = rollout_extra["test_cases"]
        if isinstance(tc, str):
            tc = json.loads(tc)
        fn_name = tc.get("fn_name")
        return tc, fn_name

    all_tests = []

    pub = rollout_extra.get("public_test_cases", "[]")
    if isinstance(pub, str):
        pub = json.loads(pub)
    all_tests.extend(pub)

    priv_raw = rollout_extra.get("private_test_cases", "")
    if priv_raw:
        try:
            decoded = base64.b64decode(priv_raw)
            decompressed = zlib.decompress(decoded)
            priv_str = pickle.loads(decompressed)
            priv = json.loads(priv_str)
            all_tests.extend(priv)
        except Exception:
            pass

    fn_name = None
    metadata = rollout_extra.get("metadata", "{}")
    if isinstance(metadata, str):
        metadata = json.loads(metadata)
    fn_name = metadata.get("func_name")

    inputs = [t["input"] for t in all_tests]
    outputs = [t["output"] for t in all_tests]
    return {"inputs": inputs, "outputs": outputs}, fn_name


# ---------------------------------------------------------------------------
#  Public test-case execution (for rollout tool)
# ---------------------------------------------------------------------------

def run_public_tests(
    code: str,
    public_test_cases: list[dict],
    fn_name: Optional[str] = None,
    timeout: int = _DEFAULT_TIMEOUT,
) -> str:
    """Execute code against public test cases and return formatted results.

    Each entry in *public_test_cases* is ``{"input": ..., "output": ...,
    "testtype": "stdin"|"functional"}``.

    Returns a human-readable string describing pass/fail per case with
    actual vs expected output and any errors.
    """
    if not public_test_cases:
        return "No sample test cases available."

    results: list[_TestResult] = []
    for tc in public_test_cases:
        inp = tc["input"]
        exp = tc["output"]
        testtype = tc.get("testtype", "stdin")
        if testtype == "functional" and fn_name:
            r = _run_functional_test_detailed(code, fn_name, inp, exp, timeout)
        else:
            r = _run_stdin_test_detailed(code, inp, exp, timeout)
        results.append(r)

    return format_test_results(results, public_test_cases)



# 测试结果输出采用有界截断策略：即使 expected output 很大，模型只看到
# 首段 + 末段 + 长度统计 + 首处 diff 位置，保证 rollout context 有硬上限。

def _truncate_diag(s: str, head: int = _OUTPUT_HEAD, tail: int = _OUTPUT_TAIL) -> str:
    """Smart truncation: short → as-is; long → head + length annotation + tail."""
    if not isinstance(s, str):
        s = str(s)
    n = len(s)
    if n <= head + tail + 50:
        return s
    return f"{s[:head]}\n...(truncated, total {n} chars)...\n{s[-tail:]}"


def _first_diff_index(a: str, b: str) -> int:
    """First index where a[i] != b[i], or min(len) if one is a prefix of the other."""
    for i, (ca, cb) in enumerate(zip(a, b)):
        if ca != cb:
            return i
    return min(len(a), len(b))


def format_test_results(
    results: list[_TestResult],
    test_cases: list[dict],
) -> str:
    """Bounded per-case footprint:
        input ≤ 300 chars, each of actual/expected ≤ ~430 chars (head+tail+anno).
    With ~6 cases, total tool output stays under ~3KB ≈ 750 tokens —
    safe for max_length=16384 even after multiple rollout turns.
    """
    lines: list[str] = []
    total = len(results)
    passed = sum(1 for r in results if r.passed)

    for i, (r, tc) in enumerate(zip(results, test_cases), 1):
        status = "PASS" if r.passed else "FAIL"

        if r.passed:
            # PASS 的 case 只标状态，不展示 input/output，节省 context
            lines.append(f"Test {i}/{total}: PASS")
        else:
            lines.append(f"--- Test {i}/{total}: FAIL ---")
            inp = tc["input"].strip()
            if len(inp) > _INPUT_DISPLAY_MAX:
                lines.append(
                    f"Input ({len(inp)} chars):\n"
                    f"{inp[:_INPUT_DISPLAY_MAX]}\n...(truncated)"
                )
            else:
                lines.append(f"Input:\n{inp}")

            if r.error:
                lines.append(f"Error: {r.error}")
            actual = r.actual or ""
            expected = r.expected or ""

            if actual and expected:
                la, le = len(actual), len(expected)
                if la != le:
                    lines.append(
                        f"Diff summary: length mismatch (actual={la} chars, expected={le} chars)"
                    )
                else:
                    di = _first_diff_index(actual, expected)
                    if di < la:
                        lines.append(
                            f"Diff summary: both lengths={la}, first differing char at index {di}"
                        )
            if actual:
                lines.append(f"Your output:\n{_truncate_diag(actual)}")
            lines.append(f"Expected output:\n{_truncate_diag(expected)}")
            lines.append("")

    lines.append(f"\nResult: {passed}/{total} tests passed.")
    if passed == total:
        lines.append(
            "All tests passed. If you are confident in your solution, "
            "output the final code in a ```python``` code block."
        )
    return "\n".join(lines)
