"""HumanEval benchmark loader wrapping ember.api.data.

HumanEval is a code generation benchmark with execution-based evaluation.
Fresh for EB1 evaluation (not in router training data).
"""

import re
from typing import Any, Iterator

from ember.api.data import HuggingFaceSource
from embergraphrouter.evals.benchmarks.base import (
    Benchmark,
    BenchmarkDomain,
    BenchmarkExample,
    EvaluationScore,
    MetricType,
)
from embergraphrouter.evals.benchmarks.registry import register_benchmark


def _extract_code(prediction: str, entry_point: str | None = None) -> str:
    """Extract Python code from model response.

    Handles common patterns:
    - Raw code
    - ```python ... ``` blocks
    - ``` ... ``` blocks
    - Code mixed with explanations
    """
    text = prediction.strip()

    # Extract from markdown code blocks
    code_block = re.search(r"```(?:python)?\s*\n?(.*?)```", text, re.DOTALL)
    if code_block:
        return code_block.group(1).strip()

    # If starts with def or looks like code, use as-is
    if text.startswith("def ") or text.startswith("class ") or text.startswith("import "):
        # Find the function if entry_point specified
        if entry_point:
            # Find the specific function
            func_pattern = rf"(def {re.escape(entry_point)}\s*\([^)]*\).*?)(?=\ndef |\Z)"
            match = re.search(func_pattern, text, re.DOTALL)
            if match:
                return match.group(1).strip()
        return text

    # Try to find any function definition
    func_match = re.search(r"(def \w+\s*\([^)]*\).*?)(?=\ndef |\Z)", text, re.DOTALL)
    if func_match:
        return func_match.group(1).strip()

    return text


def _safe_execute(code: str, test_code: str, timeout: float = 5.0) -> tuple[bool, str]:
    """Execute code with test cases in isolated environment.

    Args:
        code: Function implementation.
        test_code: Test assertions to run.
        timeout: Execution timeout in seconds.

    Returns:
        Tuple of (passed, error_message).
    """
    import multiprocessing
    import queue

    def worker(code: str, test_code: str, result_queue: Any) -> None:
        try:
            # Create isolated namespace
            namespace: dict[str, Any] = {}
            exec(code, namespace)
            exec(test_code, namespace)
            result_queue.put((True, ""))
        except AssertionError as e:
            result_queue.put((False, f"AssertionError: {e}"))
        except Exception as e:
            result_queue.put((False, f"{type(e).__name__}: {e}"))

    result_queue: queue.Queue[tuple[bool, str]] = multiprocessing.Queue()
    process = multiprocessing.Process(
        target=worker,
        args=(code, test_code, result_queue),
    )
    process.start()
    process.join(timeout=timeout)

    if process.is_alive():
        process.terminate()
        process.join()
        return False, "Timeout exceeded"

    try:
        return result_queue.get_nowait()
    except queue.Empty:
        return False, "No result from execution"


@register_benchmark("humaneval")
class HumanEvalBenchmark(Benchmark):
    """HumanEval code generation benchmark.

    Wraps HuggingFace openai_humaneval dataset. Evaluation runs generated
    code against test cases.

    This benchmark is NOT in EB1 router training data, making it suitable
    for uncontaminated evaluation.

    Note: Code execution requires appropriate sandboxing in production.
    """

    name = "humaneval"
    version = "1.0"
    domain = BenchmarkDomain.CODING
    metric_type = MetricType.PASS_AT_K
    description = "Python code generation with execution-based evaluation"

    def __init__(
        self,
        max_examples: int | None = None,
        execute_code: bool = False,
    ):
        """Initialize HumanEval benchmark.

        Args:
            max_examples: Maximum examples to load (None = all 164).
            execute_code: Whether to actually execute code for evaluation.
                If False, uses simple heuristics (useful for testing).
        """
        self._max_examples = max_examples
        self._execute_code = execute_code

    def load(self, subset: str | None = None) -> Iterator[BenchmarkExample]:
        """Load HumanEval examples via HuggingFace.

        Args:
            subset: Ignored for HumanEval (single split).

        Yields:
            BenchmarkExample instances with code completion prompts.
        """
        source = HuggingFaceSource(
            "openai_humaneval",
            split="test",
        )

        count = 0
        for batch in source.read_batches():
            for row in batch:
                if self._max_examples and count >= self._max_examples:
                    return

                task_id = str(row.get("task_id", f"HumanEval/{count}"))
                prompt = str(row.get("prompt", ""))
                canonical_solution = str(row.get("canonical_solution", ""))
                test_code = str(row.get("test", ""))
                entry_point = str(row.get("entry_point", ""))

                yield BenchmarkExample(
                    id=task_id,
                    query=prompt,
                    reference=canonical_solution,
                    metadata={
                        "domain": "coding",
                        "source": "humaneval",
                        "entry_point": entry_point,
                        "language": "python",
                    },
                    test_cases=[{
                        "test_code": test_code,
                        "entry_point": entry_point,
                    }],
                )
                count += 1

    def evaluate(
        self,
        prediction: str,
        example: BenchmarkExample,
    ) -> EvaluationScore:
        """Evaluate prediction by running test cases.

        Args:
            prediction: Model's code output.
            example: Ground truth example with test cases.

        Returns:
            EvaluationScore with pass/fail result.
        """
        entry_point = example.metadata.get("entry_point", "")
        code = _extract_code(prediction, entry_point)

        # Combine prompt (function signature) with generated body
        prompt = example.query
        if code.startswith("def "):
            # Full function provided, use as-is
            full_code = code
        else:
            # Assume it's just the body, combine with prompt
            full_code = prompt + code

        if not self._execute_code:
            # Heuristic evaluation: check if it looks like valid Python
            has_return = "return " in code or "yield " in code
            has_function = "def " in full_code
            looks_valid = has_function and (has_return or "pass" in code)

            return EvaluationScore(
                correct=looks_valid,
                score=1.0 if looks_valid else 0.0,
                details={
                    "metric": "heuristic",
                    "has_return": has_return,
                    "has_function": has_function,
                    "note": "Code execution disabled",
                },
            )

        # Execute code with test cases
        test_cases = example.test_cases or []
        if not test_cases:
            return EvaluationScore(
                correct=False,
                score=0.0,
                details={"error": "No test cases available"},
            )

        test_code = test_cases[0].get("test_code", "")
        passed, error = _safe_execute(full_code, test_code)

        return EvaluationScore(
            correct=passed,
            score=1.0 if passed else 0.0,
            details={
                "metric": "execution",
                "passed": passed,
                "error": error if not passed else None,
            },
        )
