import json
from pathlib import Path

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, json_dataset
from inspect_ai.log import ToolEvent, transcript
from inspect_ai.scorer import INCORRECT, Score, Scorer, Target, accuracy, scorer, stderr
from inspect_ai.solver import Solver, TaskState, solver, use_tools
from inspect_ai.tool import ToolDef

import astabench.evals.demo.arithmetic.task
from astabench.evals.demo.arithmetic.task import INPUT_TEMPLATE, TASK_INSTRUCTIONS
from astabench.evals.utils import not_implemented_solver
from astabench.tools import calculator


def json_to_sample_tools(obj: dict[str, str]) -> Sample:
    return Sample(
        id=obj["id"],
        input=INPUT_TEMPLATE.format(
            task_instructions=TASK_INSTRUCTIONS
            + "  You MUST use the calculator tool to help you with this task.",
            question=obj["question"],
        ),
        target=json.dumps(
            {
                "answer": obj["answer"],
                "expected_calculator_calls": obj["expected_calculator_calls"],
            }
        ),
    )


@solver
def set_up_arithmetic_tools() -> Solver:
    """Set up the calculator tool for the arithmetic task."""
    return use_tools(calculator())


@scorer(metrics=[accuracy(), stderr()])
def check_arithmetic_tool_calls() -> Scorer:
    """Scorer to check if the solution is correct and used the calculator tool."""

    async def score(state: TaskState, target: Target) -> Score:
        # Check if calculator tool is present; if not, is suggests that the
        # solver overwrote the tool list rather than appending
        calc_name = ToolDef(calculator()).name
        if not any(ToolDef(tool).name == calc_name for tool in state.tools):
            return Score(
                value=INCORRECT,
                answer="",
                explanation="No calculator tool found in solver state",
            )

        # Get calls to the calculator tool
        actual_calls = []
        for event in transcript().events:
            if not (isinstance(event, ToolEvent) and event.function == calc_name):
                continue
            actual_calls.append(
                [
                    event.arguments["op"],
                    event.arguments["num1"],
                    event.arguments["num2"],
                ]
            )

        expected_calls = json.loads(target.text)["expected_calculator_calls"]

        # Expected calls are like [["*", 2, 3], ["+", 1, 2]].  We need to
        # compare these against the actual calls made by the agent
        # (order-independent)

        n_recall = 0
        for expected_call in expected_calls:
            found = False
            for arg_list in actual_calls:
                if arg_list == expected_call:
                    found = True
                    break
            if found:
                n_recall += 1

        # Avoid division by zero if no calls were made
        recall = n_recall / max(len(actual_calls), 1)

        return Score(
            value=recall,
            answer=f"Answer: {state.output.completion}\n\nCalculator calls: {actual_calls}",
            explanation=f"Expected calls: {expected_calls}",
        )

    return score


@scorer(metrics=[accuracy(), stderr()])
def check_arithmetic() -> Scorer:
    """Use the arithmetic task's scorer to check the answer value."""

    async def score(state: TaskState, target: Target) -> Score:
        # The simple arithmetic scorer just expects a float target, so we need
        # to extract that from our complex json target
        target_val = float(json.loads(target.text)["answer"])

        return await astabench.evals.demo.arithmetic.task.check_arithmetic()(
            state,
            Target(target=str(target_val)),
        )

    return score


@task
def arithmetic_with_tools():
    return Task(
        dataset=json_dataset(
            str(Path(__file__).parent / "data.json"),
            sample_fields=json_to_sample_tools,
        ),
        setup=set_up_arithmetic_tools(),
        # Solver that just tells the user to specify a real solver via cli
        solver=not_implemented_solver(),
        scorer=[check_arithmetic_tool_calls(), check_arithmetic()],
    )
