"""
C/C++ code execution reward functions for evaluating C/C++ code correctness.

This module provides functions to evaluate the correctness of C/C++ code by:
1. Extracting code blocks from messages
2. Executing the code using the Piston execution engine
3. Comparing the output with expected results or running against test cases
"""

import asyncio
import json
import os
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union

import aiohttp

from ..models import EvaluateResult, Message, MetricResult
from ..reward_function import reward_function


@dataclass
class TestResult:
    """
    Represents the result of a single test case execution.
    """

    test_name: str
    score: float = 0.0
    status: str = "SKIPPED"
    feedback: str = ""
    actual_output: str = ""
    expected_output: str = ""


class PistonError(Exception):
    """Exception raised for errors from the Piston API."""

    pass


class PistonClient:
    """
    A client that communicates with Piston API endpoints for code execution.

    Piston is a general purpose code execution engine:
    https://github.com/engineer-man/piston
    """

    def __init__(
        self,
        base_endpoint: str = "https://emkc.org/api/v2/piston",
        session: Optional[aiohttp.ClientSession] = None,
        timeout: int = 30,
    ):
        self.base_endpoint = base_endpoint
        self._session = session
        self.timeout = timeout

    @property
    def session(self):
        if self._session is None:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(sock_read=self.timeout),
                connector=aiohttp.TCPConnector(
                    limit=10,
                    ttl_dns_cache=300,
                    keepalive_timeout=30,
                ),
            )
        return self._session

    async def close(self):
        """Close the session."""
        if self._session:
            await self._session.close()
            self._session = None

    async def get_runtimes(self) -> List[Dict[str, Any]]:
        """Get list of supported runtimes."""
        async with self.session.get(f"{self.base_endpoint}/runtimes") as response:
            if response.status != 200:
                raise PistonError(f"Error getting runtimes: {response.status}")
            return await response.json()

    async def execute(
        self,
        language: str,
        version: str,
        files: List[Dict[str, str]],
        stdin: str = "",
        args: List[str] = [],
        compile_timeout: int = 10000,
        run_timeout: int = 3000,
        compile_memory_limit: int = -1,
        run_memory_limit: int = -1,
    ) -> Dict[str, Any]:
        """
        Execute code using the Piston API.

        Args:
            language: Programming language (e.g., "c", "cpp")
            version: Version of the language (e.g., "10.2.0")
            files: List of files to include in execution (each with "name" and "content")
            stdin: Standard input to provide to the program
            args: Command-line arguments to pass to the program
            compile_timeout: Maximum compilation time in milliseconds
            run_timeout: Maximum execution time in milliseconds
            compile_memory_limit: Maximum memory for compilation in bytes (-1 for unlimited)
            run_memory_limit: Maximum memory for execution in bytes (-1 for unlimited)

        Returns:
            Dictionary containing the execution results
        """
        payload = {
            "language": language,
            "version": version,
            "files": files,
            "stdin": stdin,
            "args": args,
            "compile_timeout": compile_timeout,
            "run_timeout": run_timeout,
            "compile_memory_limit": compile_memory_limit,
            "run_memory_limit": run_memory_limit,
        }

        async with self.session.post(
            f"{self.base_endpoint}/execute",
            json=payload,
            headers={"Content-Type": "application/json"},
        ) as response:
            if response.status != 200:
                error_text = await response.text()
                raise PistonError(
                    f"Error executing code: {response.status} - {error_text}"
                )

            result = await response.json()

            if "message" in result:
                raise PistonError(result["message"])

            return result


def get_piston_client(endpoint: Optional[str] = None) -> PistonClient:
    """
    Get a Piston client instance.

    Args:
        endpoint: Optional custom Piston API endpoint

    Returns:
        PistonClient instance
    """
    piston_endpoint = endpoint or os.environ.get(
        "PISTON_ENDPOINT", "https://emkc.org/api/v2/piston"
    )
    assert isinstance(piston_endpoint, str)
    return PistonClient(base_endpoint=piston_endpoint)


def extract_code_blocks(text: str, language: str = "cpp") -> List[Dict[str, str]]:
    """
    Extract code blocks from text.

    Args:
        text: The text to extract code blocks from
        language: Language to filter by (e.g., "cpp", "c")

    Returns:
        List of dictionaries with "code" and "language" keys
    """
    pattern = r"```(\w*)\n([\s\S]*?)\n```"
    matches = re.findall(pattern, text)

    code_blocks = []
    for lang, code in matches:
        lang = lang.lower()

        if language and lang:
            if language == "cpp" and lang not in ["cpp", "c++"]:
                continue
            elif language == "c" and lang != "c":
                continue
            elif language not in ["c", "cpp"] and language != lang:
                continue

        detected_lang = lang if lang else "unknown"
        code_blocks.append({"language": detected_lang, "code": code.strip()})

    return code_blocks


def add_cpp_includes(code: str) -> str:
    """
    Add common C++ includes if they're missing.

    Args:
        code: C++ code

    Returns:
        Code with added includes if necessary
    """
    if not code:
        return code

    includes = []

    if "#include <iostream>" not in code:
        includes.append("#include <iostream>")
    if "#include <vector>" not in code:
        includes.append("#include <vector>")
    if "#include <string>" not in code:
        includes.append("#include <string>")
    if "#include <bits/stdc++.h>" not in code:
        includes.append("#include <bits/stdc++.h>")
    if "using namespace std;" not in code and "std::" not in code:
        includes.append("using namespace std;")

    if includes:
        return "\n".join(includes) + "\n\n" + code

    return code


def add_c_includes(code: str) -> str:
    """
    Add common C includes if they're missing.

    Args:
        code: C code

    Returns:
        Code with added includes if necessary
    """
    if not code:
        return code

    includes = []

    if "#include <stdio.h>" not in code:
        includes.append("#include <stdio.h>")
    if "#include <stdlib.h>" not in code:
        includes.append("#include <stdlib.h>")
    if "#include <string.h>" not in code:
        includes.append("#include <string.h>")

    if includes:
        return "\n".join(includes) + "\n\n" + code

    return code


async def execute_cpp_code(
    code: str,
    stdin: str = "",
    language: str = "cpp",
    version: str = "11.4.0",
    timeout: int = 5000,
    memory_limit: int = 512000000,
    piston_endpoint: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Execute C/C++ code using the Piston API.

    Args:
        code: C/C++ code to execute
        stdin: Standard input to provide to the program
        language: "c" or "cpp"
        version: Version of the compiler to use
        timeout: Maximum execution time in milliseconds
        memory_limit: Maximum memory in bytes
        piston_endpoint: Optional custom Piston API endpoint

    Returns:
        Dictionary with execution results
    """
    if language == "cpp":
        code = add_cpp_includes(code)
    else:
        code = add_c_includes(code)

    client = get_piston_client(piston_endpoint)

    try:
        main_file = {
            "name": "main.cpp" if language == "cpp" else "main.c",
            "content": code,
        }

        result = await client.execute(
            language=language,
            version=version,
            files=[main_file],
            stdin=stdin,
            compile_timeout=timeout,
            run_timeout=timeout,
            run_memory_limit=memory_limit,
        )

        if "compile" in result and result["compile"]["code"] != 0:
            return {
                "success": False,
                "output": None,
                "error": f"Compilation error: {result['compile']['stderr']}",
            }

        if "run" in result:
            if result["run"]["code"] == 0:
                return {
                    "success": True,
                    "output": result["run"]["stdout"],
                    "error": None,
                }
            else:
                return {
                    "success": False,
                    "output": (
                        result["run"]["stdout"] if result["run"]["stdout"] else None
                    ),
                    "error": f"Runtime error (exit code {result['run']['code']}): {result['run']['stderr']}",
                }

        return {
            "success": False,
            "output": None,
            "error": "Unknown error during execution",
        }

    except PistonError as e:
        return {
            "success": False,
            "output": None,
            "error": f"Piston error: {str(e)}",
        }
    except Exception as e:
        return {"success": False, "output": None, "error": f"Error: {str(e)}"}
    finally:
        loop = asyncio.get_event_loop()
        loop.create_task(client.close())


def compare_outputs(actual: str, expected: str) -> float:
    """
    Compare actual and expected outputs to calculate a similarity score.

    Args:
        actual: Actual output from code execution
        expected: Expected output

    Returns:
        Similarity score between 0.0 and 1.0
    """
    if actual is None:
        actual = ""
    if expected is None:
        expected = ""

    actual_norm = re.sub(r"\s+", " ", actual.strip())
    expected_norm = re.sub(r"\s+", " ", expected.strip())

    if actual_norm == expected_norm:
        return 1.0

    try:
        actual_num = float(actual_norm)
        expected_num = float(expected_norm)

        if expected_num == 0:
            return 1.0 if actual_num == 0 else 0.0

        rel_diff = abs(actual_num - expected_num) / abs(expected_num)

        if rel_diff <= 0.001:
            return 1.0
        elif rel_diff <= 0.01:
            return 0.95
        elif rel_diff <= 0.1:
            return 0.7
        else:
            return max(0.0, 1.0 - min(1.0, rel_diff))
    except (ValueError, TypeError):
        pass

    if "\n" in actual_norm or "\n" in expected_norm:
        actual_lines = actual_norm.split("\n")
        expected_lines = expected_norm.split("\n")

        common_len = min(len(actual_lines), len(expected_lines))
        if common_len == 0:
            return 0.0

        line_similarities = []
        for i in range(common_len):
            if actual_lines[i] == expected_lines[i]:
                line_similarities.append(1.0)
            else:
                line_similarities.append(
                    string_similarity(actual_lines[i], expected_lines[i])
                )

        total_weight = sum(1 / (i + 1) for i in range(common_len))
        weighted_sum = sum(
            (1 / (i + 1)) * sim for i, sim in enumerate(line_similarities)
        )
        similarity = weighted_sum / total_weight if total_weight > 0 else 0.0

        length_penalty = min(len(actual_lines), len(expected_lines)) / max(
            len(actual_lines), len(expected_lines)
        )

        return similarity * length_penalty

    return string_similarity(actual_norm, expected_norm)


def string_similarity(s1: str, s2: str) -> float:
    """
    Calculate string similarity.

    Args:
        s1: First string
        s2: Second string

    Returns:
        Similarity score between 0.0 and 1.0
    """
    if not s1 and not s2:
        return 1.0
    if not s1 or not s2:
        return 0.0

    distance = levenshtein_distance(s1, s2)
    max_len = max(len(s1), len(s2))
    return 1.0 - (distance / max_len if max_len > 0 else 0.0)


def levenshtein_distance(s1: str, s2: str) -> int:
    """
    Calculate the Levenshtein distance between two strings.

    Args:
        s1: First string
        s2: Second string

    Returns:
        Edit distance between strings
    """
    if len(s1) < len(s2):
        return levenshtein_distance(s2, s1)

    if not s2:
        return len(s1)

    previous_row: List[int] = list(range(len(s2) + 1))
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            substitutions = previous_row[j] + (c1 != c2)
            current_row.append(min(insertions, deletions, substitutions))
        previous_row = current_row

    return previous_row[-1]


async def run_cpp_test_cases(
    code: str,
    test_cases: List[Dict[str, Any]],
    language: str = "cpp",
    version: str = "11.4.0",
    timeout: int = 5000,
    memory_limit: int = 512000000,
    piston_endpoint: Optional[str] = None,
) -> List[TestResult]:
    """
    Run C/C++ code against multiple test cases.

    Args:
        code: C/C++ code to execute
        test_cases: List of test cases with "input" and "expected_output" keys
        language: "c" or "cpp"
        version: Version of the compiler to use
        timeout: Maximum execution time in milliseconds
        memory_limit: Maximum memory in bytes
        piston_endpoint: Optional custom Piston API endpoint

    Returns:
        List of TestResult objects
    """
    results = []

    for i, test_case in enumerate(test_cases):
        test_input = test_case.get("input", "")
        expected_output = test_case.get("expected_output", "")
        test_name = test_case.get("name", f"Test {i+1}")

        execution_result = await execute_cpp_code(
            code=code,
            stdin=test_input,
            language=language,
            version=version,
            timeout=timeout,
            memory_limit=memory_limit,
            piston_endpoint=piston_endpoint,
        )

        test_result = TestResult(test_name=test_name, expected_output=expected_output)

        if execution_result["success"]:
            actual_output = execution_result["output"]
            test_result.actual_output = actual_output
            similarity = compare_outputs(actual_output, expected_output)
            test_result.score = similarity

            if similarity >= 0.99:
                test_result.status = "AC"
            elif similarity > 0:
                test_result.status = "PA"
            else:
                test_result.status = "WA"
            test_result.feedback = f"Similarity: {similarity:.2f}"
        else:
            test_result.status = (
                "CE" if "Compilation error" in execution_result["error"] else "RE"
            )
            test_result.feedback = execution_result["error"]
            test_result.score = 0.0

        results.append(test_result)
        if test_result.score == 0.0:
            break

    return results


@reward_function
def ioi_cpp_code_reward(
    messages: List[Message],
    ground_truth: Union[Optional[str], Optional[List[Dict[str, Any]]]],
    language: str = "cpp",
    version: str = "11.4.0",
    timeout: int = 5000,
    memory_limit: int = 512000000,
    piston_endpoint: Optional[str] = None,
    pass_threshold: float = 0.99,
    **kwargs: Any,
) -> EvaluateResult:
    """
    Wrapper function for the asynchronous implementation to make it compatible with the reward_function decorator.
    """
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        return _ioi_cpp_code_reward_impl(
            messages=messages,
            ground_truth=ground_truth,
            language=language,
            version=version,
            timeout=timeout,
            memory_limit=memory_limit,
            piston_endpoint=piston_endpoint,
            pass_threshold=pass_threshold,
            **kwargs,
        )
    finally:
        loop.close()


def _ioi_cpp_code_reward_impl(
    messages: List[Message],
    ground_truth: Union[Optional[str], Optional[List[Dict[str, Any]]]],
    language: str = "cpp",
    version: str = "11.4.0",
    timeout: int = 5000,
    memory_limit: int = 512000000,
    piston_endpoint: Optional[str] = None,
    pass_threshold: float = 0.99,
    **kwargs: Any,
) -> EvaluateResult:
    """
    Evaluate C/C++ code correctness using the Piston execution engine.

    This function evaluates code for competitive programming problems (like IOI)
    by compiling and executing C/C++ code against test cases.

    Args:
        messages: Generated conversation messages
        ground_truth: Expected output string or list of test case dictionaries.
        language: Programming language ("c" or "cpp")
        version: Version of the compiler to use
        timeout: Maximum execution time in milliseconds
        memory_limit: Maximum memory in bytes
        piston_endpoint: Optional custom Piston API endpoint
        pass_threshold: Similarity threshold for considering a test passed
        **kwargs: Additional keyword arguments

    Returns:
        EvaluateResult with score and metrics
    """
    metrics: Dict[str, MetricResult] = {}

    if (
        not messages
        or not isinstance(messages[-1], Message)
        or messages[-1].role != "assistant"
        or messages[-1].content is None
    ):
        return EvaluateResult(
            score=0.0,
            reason="Invalid or missing assistant response in messages.",
            metrics={
                "error": MetricResult(
                    score=0.0,
                    is_score_valid=False,
                    reason="Last message not a valid assistant response.",
                )
            },
        )

    response_content = messages[-1].content

    expected_output_str_from_gt: Optional[str] = None
    test_cases_from_gt: Optional[List[Dict[str, Any]]] = None

    if isinstance(ground_truth, str):
        expected_output_str_from_gt = ground_truth
    elif isinstance(ground_truth, list):
        if all(isinstance(item, dict) for item in ground_truth):
            test_cases_from_gt = ground_truth
        else:
            return EvaluateResult(
                score=0.0,
                reason="Invalid ground_truth format: if list, must be list of test case dicts.",
                metrics={
                    "error": MetricResult(
                        score=0.0,
                        is_score_valid=False,
                        reason="Invalid ground_truth list format.",
                    )
                },
            )
    elif ground_truth is not None:
        return EvaluateResult(
            score=0.0,
            reason="Invalid ground_truth format: expected string, list of test case dicts, or None.",
            metrics={
                "error": MetricResult(
                    score=0.0,
                    is_score_valid=False,
                    reason="Invalid ground_truth format.",
                )
            },
        )

    code_blocks = extract_code_blocks(response_content, language)

    if not code_blocks:
        return EvaluateResult(
            score=0.0,
            reason=f"No {language} code blocks found in model's response.",
            metrics={
                "error": MetricResult(
                    score=0.0,
                    is_score_valid=False,
                    reason=f"No {language} code blocks found in model's response.",
                )
            },
        )

    code = code_blocks[0]["code"]

    metrics["extracted_code"] = MetricResult(
        score=0.0,
        is_score_valid=True,
        reason=f"Extracted code:\n```{language}\n{code}\n```",
    )

    if expected_output_str_from_gt and not test_cases_from_gt:
        metrics["expected_output"] = MetricResult(
            score=0.0,
            is_score_valid=True,
            reason=f"Expected output:\n{expected_output_str_from_gt}",
        )

    if test_cases_from_gt:
        results = asyncio.get_event_loop().run_until_complete(
            run_cpp_test_cases(
                code=code,
                test_cases=test_cases_from_gt,
                language=language,
                version=version,
                timeout=timeout,
                memory_limit=memory_limit,
                piston_endpoint=piston_endpoint,
            )
        )

        passed = sum(1 for result in results if result.score >= pass_threshold)
        total = len(results)
        overall_score = passed / total if total > 0 else 0.0
        final_reason = f"{passed}/{total} tests passed ({overall_score:.2%})."

        metrics["test_results"] = MetricResult(
            score=overall_score,
            is_score_valid=overall_score >= pass_threshold,
            reason=json.dumps(
                [
                    {
                        "test_name": result.test_name,
                        "status": result.status,
                        "score": result.score,
                        "feedback": result.feedback,
                    }
                    for result in results
                ],
                indent=2,
            ),
        )

        metrics["pass_rate"] = MetricResult(
            score=overall_score,
            is_score_valid=overall_score == 1.0,
            reason=f"{passed}/{total} tests passed ({overall_score:.2%})",
        )

        return EvaluateResult(score=overall_score, reason=final_reason, metrics=metrics)

    elif expected_output_str_from_gt:
        execution_result = asyncio.get_event_loop().run_until_complete(
            execute_cpp_code(
                code=code,
                language=language,
                version=version,
                timeout=timeout,
                memory_limit=memory_limit,
                piston_endpoint=piston_endpoint,
            )
        )

        if execution_result["success"]:
            output = execution_result["output"]
            final_reason = "Code executed successfully."

            metrics["execution_result"] = MetricResult(
                score=1.0,
                is_score_valid=True,
                reason=f"Code executed successfully with output:\n{output}",
            )

            similarity = compare_outputs(output, expected_output_str_from_gt)
            match_reason = f"Output similarity: {similarity:.2f}\n\nExpected:\n{expected_output_str_from_gt}\n\nActual:\n{output}"
            final_reason += f" Output similarity: {similarity:.2f}."

            metrics["output_match"] = MetricResult(
                score=similarity,
                is_score_valid=similarity >= pass_threshold,
                reason=match_reason,
            )

            return EvaluateResult(
                score=similarity, reason=final_reason, metrics=metrics
            )
        else:
            error = execution_result["error"]
            final_reason = f"Code execution failed: {error}"

            metrics["execution_result"] = MetricResult(
                score=0.0,
                is_score_valid=False,
                reason=f"Code execution failed with error:\n{error}",
            )

            return EvaluateResult(score=0.0, reason=final_reason, metrics=metrics)

    else:
        execution_result = asyncio.get_event_loop().run_until_complete(
            execute_cpp_code(
                code=code,
                language=language,
                version=version,
                timeout=timeout,
                memory_limit=memory_limit,
                piston_endpoint=piston_endpoint,
            )
        )

        if execution_result["success"]:
            output = execution_result["output"]
            final_reason = (
                "Code executed successfully (no expected output for comparison)."
            )

            metrics["execution_result"] = MetricResult(
                score=1.0,
                is_score_valid=True,
                reason=f"Code executed successfully with output:\n{output}",
            )

            return EvaluateResult(score=1.0, reason=final_reason, metrics=metrics)
        else:
            error = execution_result["error"]
            final_reason = f"Code execution failed: {error}"
            metrics["execution_result"] = MetricResult(
                score=0.0,
                is_score_valid=False,
                reason=f"Code execution failed with error:\n{error}",
            )

            return EvaluateResult(score=0.0, reason=final_reason, metrics=metrics)


@reward_function
def binary_cpp_code_reward(
    messages: List[Message],
    ground_truth: Union[Optional[str], Optional[List[Dict[str, Any]]]],
    language: str = "cpp",
    version: str = "11.4.0",
    timeout: int = 5000,
    memory_limit: int = 512000000,
    piston_endpoint: Optional[str] = None,
    pass_threshold: float = 0.99,
    **kwargs: Any,
) -> EvaluateResult:
    """
    Evaluate C/C++ code correctness and return a binary result (passed/failed).

    This function is a wrapper around ioi_cpp_code_reward that returns 1.0 if the
    score is at or above the pass_threshold, and 0.0 otherwise.

    Args:
        messages: Generated conversation messages
        ground_truth: Expected output string or list of test case dictionaries.
        language: Programming language ("c" or "cpp")
        version: Version of the compiler to use
        timeout: Maximum execution time in milliseconds
        memory_limit: Maximum memory in bytes
        piston_endpoint: Optional custom Piston API endpoint
        pass_threshold: Similarity threshold for considering a test passed
        **kwargs: Additional keyword arguments

    Returns:
        EvaluateResult with binary score (0.0 or 1.0) and metrics
    """
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        reward_output = _ioi_cpp_code_reward_impl(
            messages=messages,
            ground_truth=ground_truth,
            language=language,
            version=version,
            timeout=timeout,
            memory_limit=memory_limit,
            piston_endpoint=piston_endpoint,
            pass_threshold=pass_threshold,
            **kwargs,
        )

        score = reward_output.score
        binary_score = 1.0 if score >= pass_threshold else 0.0
        metrics = dict(reward_output.metrics)
        final_reason = f"Binary score based on threshold {pass_threshold:.2f}. Original score: {score:.2f}."
        metrics["binary_result"] = MetricResult(
            score=binary_score,
            is_score_valid=binary_score == 1.0,
            reason=f"{'Passed' if binary_score > 0 else 'Failed'} (threshold: {pass_threshold:.2f}, actual: {score:.2f})",
        )

        return EvaluateResult(score=binary_score, reason=final_reason, metrics=metrics)
    finally:
        loop.close()
