import json
import os
import shutil
from pathlib import Path
from importlib import resources
from typing import Literal

from huggingface_hub import hf_hub_download
from inspect_ai import Task, task
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.scorer import Score, Scorer, Target, mean, scorer, stderr
from inspect_ai.solver import TaskState, use_tools
from pydantic import BaseModel, ValidationError

from astabench.constants import ASTA_BENCH_DATASET_REPO, ASTA_BENCH_DATASET_REVISION
from astabench.evals.discoverybench.eval_utils import (
    run_eval_gold_vs_gen_NL_hypo_workflow,
)
from astabench.evals.discoverybench.task_utils import format_query
from astabench.evals.utils import extract_json_from_response, not_implemented_solver
from astabench.tools import python_session


def json_to_sample(
    obj: dict[str, dict], split: Literal["test", "validation"]
) -> Sample:
    hf_repo_id = "allenai/discoverybench"

    if split == "test":
        hf_directory = "discoverybench/real/test"
    elif split == "validation":
        hf_directory = "discoverybench/real/train"
    else:
        raise ValueError(f"Invalid split: {split}. Expected 'test' or 'validation'.")
    output_base = "."

    files_dict = {}

    for dataset_metadata in obj["input"]["metadata"]["datasets"]:
        # Get the filename from the dataset sample
        filename = dataset_metadata["name"]
        data_folder = obj["input"]["data_folder"]

        hf_file_path = os.path.join(hf_directory, data_folder, filename)
        destination_path = os.path.join(
            output_base, hf_directory, data_folder, filename
        )

        # Skip if already exists
        if os.path.exists(destination_path):
            files_dict[f"{data_folder}/{filename}"] = destination_path
            continue

        # Download via Hugging Face cache
        cached_file = hf_hub_download(
            repo_id=hf_repo_id,
            filename=hf_file_path,
            repo_type="dataset",
            revision=(
                ASTA_BENCH_DATASET_REVISION
                if hf_repo_id == ASTA_BENCH_DATASET_REPO
                else None
            ),
        )

        # Ensure directory exists, then copy to destination
        os.makedirs(os.path.dirname(destination_path), exist_ok=True)
        shutil.copy(cached_file, destination_path)

        files_dict[f"{data_folder}/{filename}"] = destination_path

    return Sample(
        id=obj["id"],
        input=format_query(obj["input"])
        + "\n\nMake sure you load the dataset to analyze it (or defer to an agent that can).",
        target=[
            str(obj["gold_output"]["hypothesis"]),
            str(obj["gold_output"]["workflow"]),
        ],  # target must be str or list[str], solver can't see it
        metadata={
            "query": obj["input"]["query"],
            "metadata": obj["input"]["metadata"],
        },  # solver can see this
        files=files_dict,
    )


class DiscoveryBenchOutput(BaseModel):
    hypothesis: str
    workflow: str


def extract_output(state: TaskState) -> DiscoveryBenchOutput | str:
    text_out = state.output.completion
    dict_out = extract_json_from_response(text_out)

    if isinstance(dict_out, dict):
        return DiscoveryBenchOutput.model_validate(dict_out)
    else:
        return str(dict_out)


@scorer(metrics=[mean(), stderr()])
def score_discoverybench() -> Scorer:
    async def score(state: TaskState, target: Target) -> Score:

        try:
            output = extract_output(state)
        except Exception as e:
            if isinstance(e, ValidationError):
                msg = "Invalid output format"
            elif isinstance(e, ValueError):
                msg = "Failed to extract output"
            else:
                msg = "Error parsing agent output"

            return Score(
                value=0.0,
                answer=state.output.completion,
                explanation=f"{msg}: {str(e)}",
            )

        if isinstance(output, DiscoveryBenchOutput):
            gen_hypo = output.hypothesis
            gen_workflow = output.workflow
        else:
            gen_hypo = output
            gen_workflow = ""

        eval_result = await run_eval_gold_vs_gen_NL_hypo_workflow(
            query=state.metadata["query"],
            gold_hypo=target[0],
            gold_workflow=target[1],
            gen_hypo=gen_hypo,
            gen_workflow=gen_workflow,
            dataset_meta=state.metadata["metadata"],
            llm_used="gpt-4o-2024-08-06",
            use_column_metadata=True,
        )

        return Score(
            value=eval_result["HMS"],
            answer=state.output.completion,
            metadata=eval_result,
        )

    return score


def load_discoverybench_hf(split: Literal["test", "validation"] = "test"):
    if split == "test":
        filename = "DBTest_w_hypothesis.json"
    elif split == "validation":
        filename = "DBDev_w_hypothesis.json"
    else:
        raise ValueError(f"Invalid split: {split}. Expected 'test' or 'validation'.")

    ds_file = hf_hub_download(
        repo_id=ASTA_BENCH_DATASET_REPO,
        filename=f"tasks/discoverybench/{filename}",
        repo_type="dataset",
        revision=ASTA_BENCH_DATASET_REVISION,
    )
    with open(ds_file, "r") as f:
        return json.load(f)


@task
def discoverybench(split: Literal["test", "validation"] = "test") -> Task:
    """
    An example of input and (hidden from user) scorer_Criteria can be seen under data.json
    The expected agent output should be a structured one set in the TaskState store dictionary in the following format
    (alternatively the user can pass it via state.output.completion, but we would still expect the same format):
        "output": {
            "query_id": "nls_bmi|1|1",
            "results": [
                {
                    "hypothesis": "text",
                    "workflow": "text",
                }
            ]
        }
    """

    raw_data = load_discoverybench_hf(split=split)
    dataset = MemoryDataset(samples=[json_to_sample(x, split=split) for x in raw_data])

    return Task(
        dataset=dataset,
        solver=not_implemented_solver(),
        setup=[use_tools(python_session())],
        scorer=score_discoverybench(),
        sandbox=(
            "docker",
            (
                resources.files("astabench.util.sandbox") / "sandbox_compose.yaml"
            ).as_posix(),
        ),
    )


@task
def discoverybench_test() -> Task:
    """
    A test task for DiscoveryBench, using the test split.
    """
    return discoverybench(split="test")


@task
def discoverybench_validation() -> Task:
    """
    A validation task for DiscoveryBench, using the validation split.
    """
    return discoverybench(split="validation")


__all__ = [
    "discoverybench",
    "discoverybench_test",
    "discoverybench_validation",
]
