import json
import logging
import os
import re
import numpy as np
from dataclasses import dataclass
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from importlib import resources
from typing import Dict, List, Any, Union, Tuple, Type
from datasets import Dataset, DatasetDict, IterableDatasetDict, load_dataset
from inspect_ai import Task, task
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.solver import TaskState
from inspect_ai.util import sandbox
from inspect_ai.scorer import (
    CORRECT,
    INCORRECT,
    Score,
    Scorer,
    Target,
    accuracy,
    scorer,
    stderr,
    Metric,
    metric,
    SampleScore,
)
from astabench.evals.utils import not_implemented_solver
from astabench.constants import ASTA_BENCH_DATASET_REPO, ASTA_BENCH_DATASET_REVISION

logger = logging.getLogger(__name__)

EVAL_TEMPLATE_PATH = (Path(__file__).parent / "_sandbox_eval_code.py").as_posix()

with open(EVAL_TEMPLATE_PATH) as f:
    EVAL_TEMPLATE = f.read()


def postprocess(code: str) -> str:
    """Cleans up generated code. COPIED FROM INSPECT_EVAL DS1000 (see citation below),
    :see: https://github.com/UKGovernmentBEIS/inspect_evals/blob/main/src/inspect_evals/ds1000/ds1000.py

    This method based on the paper's repository with additions marked below with "new pattern".
    https://github.com/xlang-ai/DS-1000/blob/478e065b9e2f3b23cbc4ff0659cf26be278b091d/test_ds1000.py#L58-L67
    """
    code = extract_from_tags(code, "```python\n", "\n```")
    code = extract_from_tags(code, "```\n", "\n```")  # new pattern.
    code = extract_from_tags(code, "<code>", "</code>")
    code = extract_from_tags(code, "", "</code>")
    code = extract_from_tags(code, "", "\nEND SOLUTION")
    code = extract_from_tags(code, "", "\n### END SOLUTION")  # new pattern.
    return code


def extract_from_tags(text: str, start_tag: str, end_tag: str) -> str:
    """Extracts code from between first occurence of optional tags."""
    start_index = len(start_tag) if text.startswith(start_tag) else 0
    end_index = text.find(end_tag, len(start_tag))
    end_index = end_index if end_index != -1 else len(text)
    return text[start_index:end_index]


TASK_INSTRUCTIONS = """You goal is to answer questions about different utilities in well-known Python libraries. Please directly write the Python code associated with the following question in text form and nothing more (e.g., 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 will then get directly executed as-is (in an environment with all the needed dependencies included) and checked against a target reference implementation (e.g., by checking the output of your executed code against the reference output.)

All questions will involve printing certain variables from the code; e.g., when the question says `return/output/display X` this requires printing `X`, so be sure that your code prints something (these printed values are what get evaluated. Be sure not to add additional text to these print statements unless asked to do so explicitly in the question). 

question: {question}
"""


def json_to_sample(obj: dict[str, str], split: str) -> Sample:
    """Loads the diagnostic `question` (task to be executed) and `solution`
    (i.e., a piece ofcode executing question).

    Each code example takes the following form:

        ```json
        {
           "id"       : <ID>, ## unique question/instance identifier
           "doc_id"   L <ID>, ## identifier for original documentation where question/code came from
           "question" : <QUERY>, ## code question
           "solution" : <SOLUTION> ## solution code
           "variables": <VAR_LIST> ## optional list of varibles that code should contian
           "executable_variables" : <VAR_LIST> ## optional variables that can be compared for equality.
        }
        ```

    """
    return Sample(
        id=obj["id"],
        input=TASK_INSTRUCTIONS.format(question=obj["question"]),
        target=obj["solution"],
        metadata={
            "target_variables": obj.get("variables", []),
            "executable_variables": obj.get("variable_eval", []),
            "doc_id": obj["doc_id"],
        },
    )


@dataclass
class ExecutionOutput:
    """Special helper class for storing the output of code executation, as well
    as computing additional metrics on top of what the executor provides.


    Parameters:
        solution_exec: A Boolean indicating in solution code is executable.
        model_exec: A Bollean indicating whether the model produced code is executable.
        solution_out: The output produced by the solution code (used for computing EM).
        model_out: The output produced by the model output (used for computing EM).
        solution_err: Any execution output for solution code printed to stderr.
        model_err: Any execution output for model code printed to stderr.
        variable_eq: List of variable equivalences (optional, can be empty).
        has_variable: List of variables checked for equality (optional, can be empty)

    """

    solution_exec: bool
    model_exec: bool
    solution_out: str
    model_out: str
    solution_err: str
    model_err: str
    variable_eq: List[Tuple[str, bool]]
    has_variable: List[Tuple[str, bool]]

    def __post_init__(self):
        ### compute extra metrics needed for getting instance-level scores
        self.em = False
        self.solution_out = self.solution_out.strip()
        self.model_out = self.model_out.strip()
        self.em = 0.0

        ## compute exact match metric
        if (self.solution_exec and self.model_exec) and (
            self.solution_out == self.model_out
        ):
            self.em = 1.0

        ## first variable eval
        if not self.has_variable:
            self.var_eval = -1.0
        else:
            self.var_eval = sum([1.0 for k, v in self.has_variable if v is True]) / len(
                self.has_variable
            )

        ## second variable eval
        if not self.variable_eq:
            self.var_eq_eval = -1.0
        else:
            self.var_eq_eval = sum(
                [1.0 for k, v in self.variable_eq if v is True]
            ) / len(self.variable_eq)

    @classmethod
    def empty_output(cls):
        """Gives an empty output (used to represent output for e.g., timeouts).

        Returns:
            An `empty` ExecutionOutput instance with output set to `ERROR`.

        """
        return cls(
            solution_exec=False,
            model_exec=False,
            solution_out="ERROR",
            model_out="ERROR",
            model_err=None,
            solution_err=None,
            variable_eq=[],
            has_variable=[],
        )

    @property
    def output_string(self) -> str:
        """Produces a string representation of the model output that can be used for debugging"""
        return f"em={self.em}\nmodel_exec={self.model_exec}, solution_exec={self.solution_exec}\nsolution_out:\n=====\n{self.solution_out}\nmodel output:\n\n====\n\n{self.model_out}"


@metric
def output_em() -> Metric:
    """Checks the exact match accuracy between solution code output and model code output."""

    def metric_fn(scores: list[SampleScore]) -> float:
        return sum([s.score.value["em"] for s in scores]) / len(scores)

    return metric_fn


@metric
def cluster_em() -> Metric:
    """Computer the mean exact match accuracy across instance clusters that are clustered by `doc_id`."""

    def metric_fn(scores: list[SampleScore]) -> float:
        cluster_lists = defaultdict(list)
        for s in scores:
            doc_id = s.score.metadata["doc_id"]
            em = s.score.value["em"]
            cluster_lists[doc_id].append(em)

        cluster_eval = []
        for cluster, result_list in cluster_lists.items():
            cacc = sum(result_list) / len(result_list)
            cluster_eval.append(cacc)

        return np.mean(cluster_eval)

    return metric_fn


@metric
def model_executable() -> Metric:
    """Checks the ratio of model produced code that is executable."""

    def metric_fn(scores: list[SampleScore]) -> float:
        return sum([s.score.value["model_exec"] for s in scores]) / len(scores)

    return metric_fn


@metric
def solution_executable() -> Metric:
    """Checks the ratio of solution code that is executable (should be 100%, used for checking)."""

    def metric_fn(scores: list[SampleScore]) -> float:
        return sum([s.score.value["solution_exec"] for s in scores]) / len(scores)

    return metric_fn


@metric
def variable_acc() -> Metric:
    """Checks the average accuracy of solution code containing the target variables specified in question"""

    def metric_fn(scores: list[SampleScore]) -> float:
        found = [
            s.score.value["var_eval"] for s in scores if s.score.value["var_eval"] != -1
        ]
        return sum(found) / len(found) if found else 0.0

    return metric_fn


@metric
def variable_equality_acc() -> Metric:
    """Checks the average accuracy of solution code containing equal variables specified in question"""

    def metric_fn(scores: list[SampleScore]) -> float:
        found = [
            s.score.value["var_eq_eval"]
            for s in scores
            if s.score.value["var_eq_eval"] != -1
        ]
        return sum(found) / len(found) if found else 0.0

    return metric_fn


@scorer(
    metrics=[
        output_em(),
        model_executable(),
        solution_executable(),
        variable_acc(),
        variable_equality_acc(),
        cluster_em(),
    ]
)
def diagnostic_score(evaluation_timeout: int = 180) -> Scorer:
    """Scorer to check diagnostic coding abilities. The input to the scorer is
    a code solution (`target_code` below), optional details about the
    variables that should be contained in the code (`target_variables`) and
    variables that can be checked for equivalence (`executable_variables`),
    along with a code solution produced by the model being tested. The ultimate
    goal is to check that the model produced code that  is (loosely) the same
    as the solution code in functionality and the output it produces.

    Arguments:
        evaluation_timeout: The timeout value for running the code execution.

    NOTE: Given the inherent challenge of checking two pieces of code for
    equality, and the limitations of executation-based evaluation, the provided
    scores are not perfect and subject to error. We try to mitigate this by
    reporting several different metrics described below.

    The following instance-level metrics are computed to establish this:

         -- exact match (EM): Reports whether or not the model code produces
            exactly the same output of the solution code when executed.

         -- executability: Indicates whether the model produced is executable.

         -- variable containment: The ratio of target variables contained in
            the model produced code (only computed when variable annotations
            are provided. If code is not executable, variables are still recorded
            before the code crashes, so partial credit is given in some cases).

         -- variable equivalence: The ratio of variables that are equal to the
            solution code variables (only computed when annotations are provided
            and the objects in the variables implement equality. As above, in
            the cases of code crashing variables prior to the crash are stored).

    The score is computed by executing both the solution code and the model produced
    code together. This is done via `exec` as shown above, which gives us a way to compare
    variables.

    """

    async def score(state: TaskState, target: Target) -> Score:
        """Instance-level scoring.

        Arguments:
            state: The model output containing the code to evaluation.
            target: The target annotations containing the solution code,
                and the (optional) variable annotations.

        """
        target_code = target.text.strip()
        inst_id = state.sample_id
        doc_id = state.metadata["doc_id"]
        root_dir = "/code_diagnostics"
        test_file = f"{root_dir}/test_{state.sample_id}.py"
        test_code_file = f"{root_dir}/test_{state.sample_id}.json"

        metrics_path = f"{root_dir}/{state.sample_id}_metrics.json"
        required_variables = state.metadata["target_variables"]
        eq_variables = state.metadata["executable_variables"]
        model_code = postprocess(state.output.completion)

        exec_details = {
            "solution": target_code,
            "model_code": model_code,
            "metrics_out": metrics_path,
            "required_variables": required_variables,
            "comparable_variables": eq_variables,
        }

        env = sandbox()
        await env.write_file(test_file, EVAL_TEMPLATE)
        await env.write_file(test_code_file, json.dumps(exec_details))

        try:
            exec_out = await env.exec(
                ["python", test_file, "--code_file", test_code_file],
                cwd=root_dir,
                timeout=evaluation_timeout,
            )

            if exec_out.stderr.strip():
                logger.warning(
                    f"Code executation evaluation return error: {exec_out.stderr}"
                )

            metrics = await env.read_file(metrics_path)
            out_metrics = ExecutionOutput(**json.loads(metrics))
            if not out_metrics.solution_exec:
                logger.error(
                    f"Solution code not executable, skipping instance for evaluation, please check. id={inst_id}"
                )

        except TimeoutError:
            out_metrics = ExecutionOutput.empty_output()
            logger.error(
                f"Timed out on executation, skipping instance for evaluation. Consider updating timeout. id={inst_id}"
            )

        except Exception as e:
            out_metrics = ExecutionOutput.empty_output()
            logger.error(
                f"Exception encountered on instance, e={e}, skipping from evaluation. id={inst_id}"
            )

        return Score(
            value={
                "em": out_metrics.em,
                "model_exec": 1.0 if out_metrics.model_exec else 0.0,
                "solution_exec": 1.0 if out_metrics.solution_exec else 0.0,
                "var_eval": out_metrics.var_eval,
                "var_eq_eval": out_metrics.var_eq_eval,
            },
            metadata={"doc_id": doc_id},
            explanation=out_metrics.output_string,
            answer=model_code,
        )

    return score


def load_diagnostic(split: str = "dev_small") -> list[dict]:
    """Loads code diagnostic dataset from huggingface.

    Arguments:
        split: The particular diagnostic split to use from original dataset.
             Defaults to `dev_small` (i.e., the small dev set that removes expensive
             to execute examples from the ordinary `dev` split).

    Raises:
        AssertionError when the provided split is not in `dev_small`, `dev` and
        and `test`, or when dataset is not of the right type.

    """
    assert split in set(["dev_small", "dev", "test"]), f"Unknown split: {split}"

    dataset = load_dataset(
        ASTA_BENCH_DATASET_REPO,
        data_files=f"tasks/library_diagnostic/{split}.jsonl",
        token=os.getenv("HF_ACCESS_TOKEN"),
        revision=ASTA_BENCH_DATASET_REVISION,
    )
    assert isinstance(dataset, Dataset | DatasetDict | IterableDatasetDict)
    return MemoryDataset([json_to_sample(d, split) for d in list(dataset["train"])])


@task
def code_diagnostics(split: str = "dev_small") -> Task:
    """Runs the code diagnostics task.

    Arguments:
        split: The particular diagnostics split to use, defaulting to
            `dev_small` (i.e., smaller dev with expensive code execution
            problems removed).

    """
    return Task(
        dataset=load_diagnostic(split),
        scorer=diagnostic_score(),
        solver=not_implemented_solver(),
        sandbox=(
            "docker",
            (
                resources.files("astabench.util.sandbox") / "sandbox_compose.yaml"
            ).as_posix(),
        ),
    )


@task
def code_diagnostics_dev() -> Task:
    """Runs the code diagnostics task on the dev set.

    Arguments:
        split: The particular diagnostics split to use, defaulting to
            `small_dev` (i.e., smaller dev with expensive code execution
            problems removed).

    """
    return code_diagnostics(split="dev")


@task
def code_diagnostics_small_dev() -> Task:
    """Runs the code diagnostics task on the small dev set.

    Arguments:
        split: The particular diagnostics split to use, defaulting to
            `small_dev` (i.e., smaller dev with expensive code execution
            problems removed).

    """
    return code_diagnostics(split="dev_small")


@task
def code_diagnostics_test() -> Task:
    """Runs the code diagnostics task on the test set.

    Arguments:
        split: The particular diagnostics split to use, defaulting to
            `small_dev` (i.e., smaller dev with expensive code execution
            problems removed).

    """
    return code_diagnostics(split="test")
