import re
from pathlib import Path

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, json_dataset
from inspect_ai.scorer import (
    CORRECT,
    INCORRECT,
    Score,
    Scorer,
    Target,
    accuracy,
    scorer,
    stderr,
)
from inspect_ai.solver import TaskState

from astabench.evals.utils import not_implemented_solver

TASK_INSTRUCTIONS = 'Answer the given arithmetic question.  Output your answer as JSON like {{"answer": 1234}}.'
INPUT_TEMPLATE = """{task_instructions}\n\n{question}"""


def json_to_sample(obj: dict[str, str]) -> Sample:
    # NOTE: inspect_ai includes built-in data loaders; if the fields in
    # data.json match up with the field names in Sample then json_dataset can
    # be applied as-is; this is an example of a custom loader
    return Sample(
        id=obj["id"],
        # Populated as the first user message in task state `messages`
        input=INPUT_TEMPLATE.format(
            task_instructions=TASK_INSTRUCTIONS,
            question=obj["question"],
        ),
        # Target isn't included in `TaskState`` so solver can't see it at all;
        # unfortunately, it must be str or list[str].
        target=str(obj["answer"]),
        # Metadata is included in `TaskState`` so solver can technically see
        # it; some built-in solvers use it to fill prompt template values, but
        # otherwise it can be used for any arbitrary data.
        metadata={
            "operations": obj.get("operations"),
        },
    )


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

    # NOTE: This is just an example custom scorer; inspect_ai also includes
    # built-in scorers for things like exact match or even llm grading:
    # https://inspect.ai-safety-institute.org.uk/scorers.html
    async def score(state: TaskState, target: Target) -> Score:
        raw_output = state.output.completion
        explanation = "<Failed to parse output>"
        is_correct = False

        m = re.search(r"\"answer\":\s*(?P<answer>-?\d+(\.\d+)?)", raw_output)
        if m is not None:
            answer = float(m.group("answer"))
            explanation = f"Parsed answer: {answer} (correct answer = {target.text})"
            is_correct = abs(answer - float(target.text)) < 1e-6
        return Score(
            value=CORRECT if is_correct else INCORRECT,
            answer=raw_output,
            explanation=explanation,
        )

    return score


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