from typing import Any, cast

from inspect_ai.scorer import (
    Score,
    Scorer,
    ScoreReducer,
    Target,
    multi_scorer,
    score_reducer,
    scorer,
)
from inspect_ai.solver import TaskState
from inspect_ai.util import ExecResult, sandbox

from .metrics import (
    percentage_main_problems_solved,
    percentage_subproblems_solved,
    total_main_problems_solved,
    total_subproblems_solved,
)
from .util import get_generated_code, subproblem_str_to_int

"""
Nota bene:

Due to a current (0.3.47) Inspect issue where scores are being converted to floats before being processed by metrics, correct answers are scored as 1.0, and incorrect answers as 0.0, rather than the usual CORRECT or INCORRECT.

"""


@score_reducer
def reduce_subproblems() -> ScoreReducer:
    """
    Combine subproblem scores into a main problem score.

    The value of a subproblem score is a dict whose key is its subproblem number, and whose value is either 1.0, if it was answered correctly, or 0.0 if it was ans wered incorrectly.

    The value of a main problem is a dict, whose keys are subproblem numbers and whose values are 1.0 if the corresponding subproblem was answered correctly, and 0.0 otherwise.

    Returns:
    --------
    ScoreReducer
       A function to convert subproblem scores to main problem scores.

    """

    def reduce(scores: list[Score]) -> Score:
        result = {
            k: v
            for score in scores
            for k, v in cast(dict[str, Any], score.value).items()
        }

        explanations = "\n".join(
            [score.explanation for score in scores if score.explanation is not None]
        )

        return Score(
            value=result,
            explanation=explanations,
        )

    return reduce


def verify_subproblem(subproblem: dict[str, Any], timeout: int) -> Scorer:
    """
    Verify if a subproblem is correct.

    A subproblem is considered correct if all its test cases pass i.e. execute without erroring.

    """

    async def score(state: TaskState, target: Target) -> Score:
        def compose_code() -> str:
            """
            Compose code to be executed.

            This consists of:
              * dependencies, as defined for the main problem in state.metadata["dependencies"]
              * functions definitions produced by each previous subproblems
              * the code produced in response to **this** subproblem
              * the tests associated with **this** subproblem

            """
            subproblem_no = subproblem["step_number"]
            test_cases = subproblem["test_cases"]

            # Collect dependencies and the model's subproblem solutions.
            code_sections: list[str] = []
            code_sections.append(state.metadata["required_dependencies"])
            code_sections.append(
                "from test_util import are_dicts_close, cmp_tuple_or_list"
            )
            code_sections += get_generated_code(state)[
                : subproblem_str_to_int(subproblem_no)
            ]

            # Collect test cases and data.
            code_sections.append("# Test cases:")
            code_sections.append("from process_data import process_hdf5_to_tuple")
            code_sections.append(
                f"targets = process_hdf5_to_tuple('{subproblem_no}', {len(test_cases)})"
            )

            for i, test_case in enumerate(test_cases):
                code_sections.append(f"target = targets[{i}]")
                code_sections.append(test_case)

            # Compose code sections to a single string.
            code = "\n".join(code_sections)
            return code

        def format_explanation(code: str, result: ExecResult[str]) -> str:
            """Format the score explanation, based on the result of code execution."""
            subproblem_no = subproblem["step_number"]
            explanation = (
                f"The following code was executed:\n\n```python\n{code}\n```\n"
            )
            if result.success:
                explanation += (
                    f"All test cases passed for subproblem {subproblem_no}.\n"
                )
            else:
                explanation += "Code did not pass all test cases.\n"
                if result.stderr:
                    explanation += f"Error details:\n```python\n{result.stderr}\n```\n"
            return explanation

        code = compose_code()
        try:
            result = await sandbox().exec(
                cmd=["python", "-c", code],
                timeout=timeout,
            )
            explanation = format_explanation(code, result)
        except TimeoutError:
            result = ExecResult(False, 1, "", "Verification timed out.")
            explanation = "execution timed out"

        return Score(
            value={subproblem["step_number"]: 1.0 if result.success else 0.0},
            answer=code,
            explanation=explanation,
        )

    return score


@scorer(
    metrics=[
        percentage_main_problems_solved(),
        percentage_subproblems_solved(),
        total_main_problems_solved(),
        total_subproblems_solved(),
    ]
)
def verify(timeout: int) -> Scorer:
    """Score a problem, by individually scoring its subproblems and composing the resulting scores."""

    async def score(state: TaskState, target: Target) -> Score:
        subproblem_scorers = [
            verify_subproblem(subproblem=subproblem, timeout=timeout)
            for subproblem in state.metadata["sub_steps"]
        ]

        outer_scorer = multi_scorer(
            scorers=subproblem_scorers,
            reducer=reduce_subproblems(),
        )
        res = await outer_scorer(state, target)
        return res

    return score
