"""Sandboxed code execution and test case validation."""

from __future__ import annotations

import logging
import re
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path

from cohort.benchmark.config import CodingTask, TestCase

logger = logging.getLogger(__name__)


@dataclass
class TestResult:
    """Result of running one test case."""

    passed: bool
    input: str
    expected: str
    actual: str
    error: str = ""
    timed_out: bool = False


def extract_code(text: str) -> str:
    """Extract Python code from LLM response.

    Looks for ```python ... ``` blocks first, then plain ``` blocks.
    If no fenced code blocks, returns the entire text (agent may have
    produced raw code).
    """
    # Try python-fenced blocks first
    python_blocks = re.findall(
        r"```python\s*\n(.*?)```", text, re.DOTALL
    )
    if python_blocks:
        return python_blocks[-1].strip()

    # Try any fenced block
    any_blocks = re.findall(r"```\s*\n(.*?)```", text, re.DOTALL)
    if any_blocks:
        return any_blocks[-1].strip()

    # Fallback: look for lines that look like code (import, def, class, print)
    lines = text.split("\n")
    code_lines = []
    in_code = False
    for line in lines:
        stripped = line.strip()
        if stripped.startswith((
            "import ", "from ", "def ", "class ",
            "print(", "if ", "for ", "while ",
        )):
            in_code = True
        if in_code:
            code_lines.append(line)

    return "\n".join(code_lines).strip() if code_lines else ""


def run_code_against_test(
    code: str,
    test_case: TestCase,
    time_limit: float = 5.0,
) -> TestResult:
    """Run code with test case input and check output."""
    if not code.strip():
        return TestResult(
            passed=False,
            input=test_case.input,
            expected=test_case.expected_output,
            actual="",
            error="No code to execute",
        )

    # Write code to temp file
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False, encoding="utf-8"
    ) as f:
        f.write(code)
        code_path = f.name

    try:
        result = subprocess.run(
            ["python", code_path],
            input=test_case.input,
            capture_output=True,
            text=True,
            timeout=time_limit,
            cwd=tempfile.gettempdir(),
        )

        actual = result.stdout.strip()
        expected = test_case.expected_output.strip()

        # Normalize whitespace for comparison
        actual_norm = "\n".join(line.rstrip() for line in actual.split("\n"))
        expected_norm = "\n".join(line.rstrip() for line in expected.split("\n"))

        passed = actual_norm == expected_norm

        return TestResult(
            passed=passed,
            input=test_case.input,
            expected=expected,
            actual=actual,
            error=result.stderr.strip() if not passed else "",
        )

    except subprocess.TimeoutExpired:
        return TestResult(
            passed=False,
            input=test_case.input,
            expected=test_case.expected_output.strip(),
            actual="",
            error=f"Timed out after {time_limit}s",
            timed_out=True,
        )

    except Exception as e:
        return TestResult(
            passed=False,
            input=test_case.input,
            expected=test_case.expected_output.strip(),
            actual="",
            error=str(e),
        )

    finally:
        Path(code_path).unlink(missing_ok=True)


def evaluate_task(
    code: str, task: CodingTask
) -> tuple[list[TestResult], list[TestResult]]:
    """Run code against all test cases (sample + hidden).

    Returns (sample_results, hidden_results).
    """
    sample_results = [
        run_code_against_test(code, tc, task.time_limit_seconds)
        for tc in task.sample_cases
    ]
    hidden_results = [
        run_code_against_test(code, tc, task.time_limit_seconds)
        for tc in task.hidden_cases
    ]
    return sample_results, hidden_results
