"""
CyberMetric: A Benchmark Dataset based on Retrieval-Augmented Generation for Evaluating LLMs in Cybersecurity Knowledge

Norbert Tihanyi, Mohamed Amine Ferrag, Ridhi Jain, Tamas Bisztray, Merouane Debbah
https://arxiv.org/abs/2402.07688

# Run all subsets
inspect eval cybermetric.py

# Run specific subsets
inspect eval cybermetric.py@cybermetric_80
inspect eval cybermetric.py@cybermetric_500
inspect eval cybermetric.py@cybermetric_2000
inspect eval cybermetric.py@cybermetric_10000

"""

import json
from pathlib import Path
from typing import Any

import requests
from inspect_ai import Task, task
from inspect_ai.dataset import Dataset, Sample, json_dataset
from inspect_ai.scorer import choice
from inspect_ai.solver import multiple_choice, system_message
from platformdirs import user_cache_dir

GITHUB_DATASET_URL_TEMPLATE = "https://raw.githubusercontent.com/cybermetric/CyberMetric/205262cdf5022ba890e792efd176fb19d42913fa/{name}-v1.json"
SYSTEM_MESSAGE = "You are a security expert who answers questions."
CACHE_DIR = Path(user_cache_dir("inspect_evals")) / "cybermetric" / "data"


@task
def cybermetric_80() -> Task:
    """Evaluates model on CyberMetric-80"""
    return create_task("CyberMetric-80")


@task
def cybermetric_500() -> Task:
    """Evaluates model on CyberMetric-500"""
    return create_task("CyberMetric-500")


@task
def cybermetric_2000() -> Task:
    """Evaluates model on CyberMetric-2000"""
    return create_task("CyberMetric-2000")


@task
def cybermetric_10000() -> Task:
    """Evaluates model on CyberMetric-10000"""
    return create_task("CyberMetric-10000")


def create_task(dataset_name: str) -> Task:
    return Task(
        dataset=load_dataset(dataset_name),
        solver=[system_message(SYSTEM_MESSAGE), multiple_choice()],
        scorer=choice(),
    )


def load_dataset(name: str) -> Dataset:
    """
    Load a dataset from the CyberMetric repository on GitHub.

    The json files on GitHub have the questions nested under a "questions" key, which makes it incompatible with the `json_dataset` loader. This function reads the json file, extracts the questions, and writes them to a jsonl file, which can then be read by `json_dataset`.
    """
    dataset_url = GITHUB_DATASET_URL_TEMPLATE.format(name=name)
    json_file = CACHE_DIR / f"{name}.json"
    jsonl_file = CACHE_DIR / f"{name}.jsonl"

    if not jsonl_file.exists():
        _ensure_data(json_file, dataset_url)

        # Read data from json file and un-nest
        with open(json_file, "r") as f:
            data = json.load(f)["questions"]

        # Write data to back to jsonl file
        with jsonl_file.open("w") as f:
            f.writelines(json.dumps(record) + "\n" for record in data)

    dataset = json_dataset(
        json_file=str(jsonl_file),
        name=name,
        sample_fields=record_to_sample,
        auto_id=True,
    )
    return dataset


def _ensure_data(json_file: Path, dataset_url: str) -> None:
    if not json_file.exists():
        json_file.parent.mkdir(parents=True, exist_ok=True)
        response = requests.get(dataset_url)
        response.raise_for_status()
        with open(json_file, "wb") as f:
            f.write(response.content)


def record_to_sample(record: dict[str, Any]) -> Sample:
    options = ", ".join([f"{key}) {value}" for key, value in record["answers"].items()])
    input = f"Question: {record['question']}\nOptions: {options}\n\n"
    return Sample(
        input=input, choices=list(record["answers"].keys()), target=record["solution"]
    )
