from typing import Any

import pytest
from inspect_ai.dataset import Sample

from inspect_evals.bigcodebench.bigcodebench import (
    INSTRUCTION_PROMPT,
    get_record_to_sample,
)

SIMPLE_RECORD: dict[str, Any] = {
    "task_id": "BigCodeBench/9999999",
    "complete_prompt": (
        '''def add(a: int, b: int) -> int:
    """
    Add two integers and return their sum.

    Args:
        a: First integer
        b: Second integer

    Returns:
        The sum of a and b

    Example:
        >>> add(1, 2)
        3
    """
'''
    ),
    "instruct_prompt": "Write a function that adds two integers and returns their sum.",
    "canonical_solution": "    return a + b",
    "code_prompt": "def add(a: int, b: int) -> int:\n",
    "test": (
        """def test_add():
    assert add(1, 2) == 3
    assert add(-1, 1) == 0
    assert add(0, 0) == 0"""
    ),
    "entry_point": "add",
    "doc_struct": (
        """{"description": ["Add two integers and return their sum."], "params": ["a: First integer", "b: Second integer"], "returns": ["The sum of a and b"], "examples": [">>> add(1, 2)", "3"]}"""
    ),
    "libs": "[]",
}

# Actual record 1116 from BigCodeBench
COMPLEX_RECORD: dict[str, Any] = {
    "task_id": "BigCodeBench/1116",
    "complete_prompt": (
        '''import random
import statistics

# Constants
AGE_RANGE = (22, 60)

def task_func(dict1):
    """
    Calculate the mean, the median, and the mode(s) of the age of the employees in the department "EMP$$."
    Generate random ages for each employee within the range [22, 60].

    Parameters:
    dict1 (dict): A dictionary with department codes as keys and number of employees
                  as values.

    Returns:
    tuple: A tuple of mean, median, and a list of mode(s) of employee ages.

    Requirements:
    - random
    - statistics

    Example:
    >>> random.seed(0)
    >>> d = {'EMP$$': 10, 'MAN$$': 5, 'DEV$$': 8, 'HR$$': 7}
    >>> stats = task_func(d)
    >>> print(stats)
    (44.7, 46.5, [46, 48, 24, 38, 54, 53, 47, 41, 52, 44])
    """
'''
    ),
    "instruct_prompt": (
        """Calculate the mean, the median, and the mode(s) of the age of the employees in the department "EMP$$." Generate random ages for each employee within the range [22, 60].
The function should output with:
    tuple: A tuple of mean, median, and a list of mode(s) of employee ages.
You should write self-contained code starting with:
import random
import statistics
# Constants
AGE_RANGE = (22, 60)
def task_func(dict1):
"""
    ),
    "canonical_solution": (
        """    emp_ages = []

    for prefix, num_employees in dict1.items():
        if not prefix.startswith('EMP$$'):
            continue

        for _ in range(num_employees):
            age = random.randint(*AGE_RANGE)
            emp_ages.append(age)

    # If no employees in EMP$$ department
    if not emp_ages:
        return 0, 0, []

    mean_age = statistics.mean(emp_ages)
    median_age = statistics.median(emp_ages)
    mode_age = statistics.multimode(emp_ages)

    return mean_age, median_age, mode_age"""
    ),
    "code_prompt": (
        """import random
import statistics
# Constants
AGE_RANGE = (22, 60)
def task_func(dict1):
"""
    ),
    "test": (
        """import unittest
class TestCases(unittest.TestCase):

    def test_case_1(self):
        random.seed(0)
        # Input: 10 employees in "EMP$$" department
        d = {'EMP$$': 10}
        mean_age, median_age, mode_age = task_func(d)

        # Checks
        self.assertTrue(22 <= mean_age <= 60)
        self.assertTrue(22 <= median_age <= 60)
        self.assertTrue(all(22 <= age <= 60 for age in mode_age))

    def test_case_2(self):
        random.seed(0)
        # Input: Different number of employees in multiple departments
        d = {'EMP$$': 10, 'MAN$$': 5, 'DEV$$': 8, 'HR$$': 7}
        mean_age, median_age, mode_age = task_func(d)

        # Checks
        self.assertTrue(22 <= mean_age <= 60)
        self.assertTrue(22 <= median_age <= 60)
        self.assertTrue(all(22 <= age <= 60 for age in mode_age))

    def test_case_3(self):
        random.seed(0)
        # Input: No employees in "EMP$$" department
        d = {'MAN$$': 5, 'DEV$$': 8, 'HR$$': 7}
        mean_age, median_age, mode_age = task_func(d)

        # Checks
        self.assertEqual(mean_age, 0)
        self.assertEqual(median_age, 0)
        self.assertEqual(mode_age, [])

    def test_case_4(self):
        random.seed(0)
        # Input: Large number of employees in "EMP$$" department to increase likelihood of multiple modes
        d = {'EMP$$': 1000}
        mean_age, median_age, mode_age = task_func(d)

        # Checks
        self.assertTrue(22 <= mean_age <= 60)
        self.assertTrue(22 <= median_age <= 60)
        self.assertTrue(all(22 <= age <= 60 for age in mode_age))

    def test_case_5(self):
        random.seed(0)
        # Input: Only one employee in "EMP$$" department
        d = {'EMP$$': 1}
        mean_age, median_age, mode_age = task_func(d)

        # Checks
        self.assertTrue(22 <= mean_age <= 60)
        self.assertEqual(mean_age, median_age)
        self.assertEqual([mean_age], mode_age)"""
    ),
    "entry_point": "task_func",
    "doc_struct": (
        """{"description": ["Calculate the mean, the median, and the mode(s) of the age of the employees in the department "EMP$$.", "Generate random ages for each employee within the range [22, 60]."], "notes": [], "params": ["dict1 (dict): A dictionary with department codes as keys and number of employees", "as values."], "returns": ["tuple: A tuple of mean, median, and a list of mode(s) of employee ages."], "reqs": ["random", "statistics"], "raises": [], "examples": [">>> random.seed(0)", ">>> d = {'EMP$$': 10, 'MAN$$': 5, 'DEV$$': 8, 'HR$$': 7}", ">>> stats = task_func(d)", ">>> print(stats)", "(44.7, 46.5, [46, 48, 24, 38, 54, 53, 47, 41, 52, 44])"]}"""
    ),
    "libs": "['statistics', 'random']",
}

TEST_RECORDS = [
    (SIMPLE_RECORD, 9999999),
    (COMPLEX_RECORD, 1116),
]


@pytest.mark.parametrize("record,expected_id", TEST_RECORDS)
def test_get_record_to_sample(record: dict[str, Any], expected_id: int):
    record_to_sample = get_record_to_sample()
    sample = record_to_sample(record)

    # Verify the sample structure and content
    assert isinstance(sample, Sample)
    assert sample.id == expected_id  # ID from task_id
    assert sample.input == INSTRUCTION_PROMPT + record["complete_prompt"]
    assert sample.target == record["complete_prompt"] + record["canonical_solution"]
    assert sample.metadata == {
        "test": record["test"],
        "libs": record["libs"],
        "task_id_metadata": record["task_id"],
    }


@pytest.mark.parametrize("record,expected_id", TEST_RECORDS)
def test_get_record_to_sample_custom_prompt(record: dict[str, Any], expected_id: int):
    # Test with custom instruction prompt
    custom_prompt = "Write a Python function:\n"
    record_to_sample = get_record_to_sample(instruction_prompt=custom_prompt)
    sample = record_to_sample(record)

    # Verify custom prompt is used correctly
    assert isinstance(sample, Sample)
    assert sample.id == expected_id  # ID should still be the same
    assert sample.input == custom_prompt + record["complete_prompt"]
    assert sample.target == record["complete_prompt"] + record["canonical_solution"]
    assert sample.metadata == {
        "test": record["test"],
        "libs": record["libs"],
        "task_id_metadata": record["task_id"],
    }
