import json
import re
import warnings
from dataclasses import dataclass
from json import JSONDecodeError
from pathlib import Path
from importlib import resources
from typing import Dict, List, Optional

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 inspect_ai.util import sandbox

from astabench.evals.utils import not_implemented_solver
from astabench.types.code_execution import CodeAgentOutput

TASK_INSTRUCTIONS = "Please write the code associated with each query and nothing more; do not wrap it in any special quotations. If you are reasoning about the code, the final `submit` should just be the raw code. Your code with then get directly executed."
INPUT_TEMPLATE = """{task_instructions}\n\n{question}"""


def parse_code_output(output: str) -> CodeAgentOutput:
    try:
        json_out = json.loads(output)
        exec_res = json_out.get("exec_result", None)
        output_code = json_out.get("code", None)

        if exec_res is None and output_code is None:
            warnings.warn(
                f"OUTPUT ERROR: your agent did not output one or more of required outputs needed for evaluation (`exec_res` or `output_code`), skipping: {output}"
            )

    except JSONDecodeError:
        exec_res = None
        output_code = output

    return CodeAgentOutput(exec_res=exec_res, output_code=output_code)


def json_to_sample(obj: dict[str, str]) -> Sample:
    return Sample(
        id=obj["id"],
        input=INPUT_TEMPLATE.format(
            task_instructions=TASK_INSTRUCTIONS,
            question=obj["question"],
        ),
        target=str(obj["answer"]),
    )


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

    async def score(state: TaskState, target: Target) -> Score:
        is_correct = False
        timed_out = False

        agent_output = parse_code_output(state.output.completion)
        exec_res = agent_output.exec_res

        ## perform execution of not already done
        if exec_res is None:
            output_code = agent_output.output_code
            in_quotes = re.search(r"```python(.*)```", output_code, re.DOTALL)
            if in_quotes:
                output_code = in_quotes.groups()[0]

            root_dir = "/code_execution"
            test_file = f"{root_dir}/test_{state.sample_id}.py"
            test_program = output_code
            env = sandbox()
            await env.write_file(test_file, test_program)

            try:
                exec_res = await env.exec(
                    ["python", test_file], cwd=root_dir, timeout=180
                )
                exec_res = exec_res.stdout.strip()

            except TimeoutError:
                timed_out = True
                exec_res = None

        if str(exec_res) == str(target.text.strip()):
            is_correct = True

        return Score(
            value=CORRECT if is_correct else INCORRECT,
            answer=exec_res,
            explanation=(
                "## Execution timed out\n"
                if timed_out
                else "" + f"\n## Test output\n {exec_res}"
            ),
        )

    return score


@task
def code_demo():
    return Task(
        dataset=json_dataset(
            str(Path(__file__).parent / "data.json"),
            sample_fields=json_to_sample,
        ),
        solver=not_implemented_solver(),
        scorer=check_code_execution(),
        sandbox=(
            "docker",
            (
                resources.files("astabench.util.sandbox") / "sandbox_compose.yaml"
            ).as_posix(),
        ),
    )
