import re
from pathlib import Path
from typing import Any

from inspect_ai import Task, task
from inspect_ai.dataset import Dataset, Sample, hf_dataset
from inspect_ai.model import get_model
from inspect_ai.scorer import Score, Scorer, Target, accuracy, scorer, stderr
from inspect_ai.solver import (
    TaskState,
    generate,
    system_message,
)
from inspect_ai.util import sandbox

DATASET_HF_REPO_ID = "xlangai/DS-1000"

DEFAULT_SYSTEM_MESSAGE = (
    "Write a short code following the given format and indentation. "
    "Place the executable code between <code> and </code> tags, without any other non-executable things"
)

MODEL_SPECIFIC_SYSTEM_MESSAGES = {
    "openai/gpt-4o.*": "Only provide the code completion needed. Don't repeat the context code."
}


@task
def ds1000(library: str | None = None, perturbation: str | None = None) -> Task:
    model_name = str(get_model())
    sys_message = DEFAULT_SYSTEM_MESSAGE
    for pattern, message in MODEL_SPECIFIC_SYSTEM_MESSAGES.items():
        if re.match(pattern, model_name):
            sys_message = message
    return Task(
        name="DS-1000",
        dataset=ds1000_dataset(library, perturbation),
        solver=[
            system_message(sys_message),
            generate(),
        ],
        scorer=ds1000_scorer(),
        sandbox=("docker", (Path(__file__).parent / "compose.yaml").as_posix()),
    )


def ds1000_dataset(
    library: str | None = None, perturbation: str | None = None
) -> Dataset:
    """Loads the DS-1000 dataset from HuggingFace and converts it to the Inspect dataset.

    Optionally filters by library and perturbation type.
    """
    ds = hf_dataset(
        path=DATASET_HF_REPO_ID,
        split="test",
        sample_fields=record_to_sample,
        shuffle=True,
    )
    if library:
        ds = ds.filter(
            lambda x: x.metadata is not None and x.metadata["library"] == library
        )
    if perturbation:
        ds = ds.filter(
            lambda x: x.metadata is not None
            and x.metadata["perturbation_type"] == perturbation
        )
    return ds


def record_to_sample(record: dict[str, Any]) -> Sample:
    return Sample(
        id=record["metadata"]["problem_id"],
        input=record["prompt"],
        target=f"```python\n{record['reference_code']}\n```",
        metadata={
            "code_context": record["code_context"],
            "library": record["metadata"]["library"],
            "perturbation_type": record["metadata"]["perturbation_type"],
        },
    )


@scorer(metrics=[accuracy(), stderr()])
def ds1000_scorer() -> Scorer:
    async def score(state: TaskState, target: Target) -> Score:
        root_dir = "/ds1000"
        code_context = state.metadata["code_context"]
        solution = postprocess(state.output.completion)
        test_file = f"{root_dir}/test_{state.sample_id}.py"
        test_program = (
            f"{code_context}\n"
            + f"solution={repr(solution)}\n"
            + "test_execution(solution)\n"
            + ("test_string(solution)\n" if "test_string(" in code_context else "")
        )
        env = sandbox()
        await env.write_file(test_file, test_program)
        timed_out = False
        try:
            exec_res = await env.exec(["python", test_file], cwd=root_dir, timeout=180)
        except TimeoutError:
            timed_out = True
            exec_res = None
        return Score(
            value="C" if exec_res and exec_res.success else "I",
            answer=f"```python\n{solution}\n```",
            explanation=(
                ("## Execution timed out\n" if timed_out else "")
                + f"## Test output:\n```\n{exec_res.stdout if exec_res else ''}\n```\n"
                + f"## Test error:\n```\n{exec_res.stderr if exec_res else ''}\n```"
            ),
        )

    return score


def postprocess(code: str) -> str:
    """Cleans up generated code.

    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]
