"""SEvenLLM: Benchmarking, Eliciting, and Enhancing Abilities of Large Language Models in Cyber Threat Intelligence

Hangyuan Ji, Jian Yang, Linzheng Chai, Chaoren Wei, Liqun Yang,
Yunlong Duan, Yunli Wang, Tianzhen Sun, Hongcheng Guo, Tongliang Li,
Changyu Ren, Zhoujun Li

https://arxiv.org/abs/2405.03446

# Eval for MCQs (Understanding) in Simplified Chinese (Zh)
inspect eval inspect_evals/sevenllm_mcq_zh

# Eval for MCQs (Understanding) in English (En)
inspect eval inspect_evals/sevenllm_mcq_en

# Eval for QAs (Generation) in Simplified Chinese (Zh)
inspect eval inspect_evals/sevenllm_qa_zh

# Eval for QA (Generation) in English (En)
inspect eval inspect_evals/sevenllm_qa_en
"""

import json
import re
from typing import Any

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 generate, multiple_choice, prompt_template

from inspect_evals.sevenllm.scorers import rouge_l_scorer, semantic_similarity_scorer

BENCHMARK_DATASET_URL = "https://huggingface.co/datasets/Multilingual-Multimodal-NLP/SEVENLLM-Dataset/raw/main/test.jsonl"
# This prompt has been inspired from the inference work done by the authors.
# See: https://github.com/CSJianYang/SEevenLLM/blob/main/infer.py
# The prompt remains the same for both the languages.
TEMPLATE = r"""
Below is an instruction that describes a task, paired with an input that provides further context.
Write a response that appropriately completes the request. If you are given options (A, B, C, D),
return just the option letter and not the entire text. If you're required to return a JSON, only return a JSON and
nothing else.

{prompt}
"""


@task
def sevenllm_mcq_zh() -> Task:
    """Inspect task implementing the SEvenLLM benchmark for MCQs in Simplified Chinese."""
    return Task(
        dataset=get_sevenllm_dataset(language="zh", data_format="mcq"),
        solver=[prompt_template(template=TEMPLATE), multiple_choice()],
        scorer=choice(),
    )


@task
def sevenllm_mcq_en() -> Task:
    """Inspect task implementing the SEvenLLM benchmark for MCQs in English."""
    return Task(
        dataset=get_sevenllm_dataset(language="en", data_format="mcq"),
        solver=[prompt_template(template=TEMPLATE), multiple_choice()],
        scorer=choice(),
    )


@task
def sevenllm_qa_zh() -> Task:
    """Inspect task implementing the SEvenLLM benchmark for QA in Simplified Chinese."""
    return Task(
        dataset=get_sevenllm_dataset(language="zh", data_format="qa"),
        solver=[prompt_template(template=TEMPLATE), generate()],
        scorer=[rouge_l_scorer(is_zh=True), semantic_similarity_scorer()],
    )


@task
def sevenllm_qa_en() -> Task:
    """Inspect task implementing the SEvenLLM benchmark for QA in English."""
    return Task(
        dataset=get_sevenllm_dataset(language="en", data_format="qa"),
        solver=[prompt_template(template=TEMPLATE), generate()],
        scorer=[rouge_l_scorer(is_zh=False), semantic_similarity_scorer()],
    )


def contains_zsh(text: str) -> bool:
    """Return True if the text contains a Simplified-Chinese character."""
    # Regular expression to match Simplified Chinese characters
    # CJK Unified Ideographs range: \u4e00-\u9fff
    pattern = re.compile(r"[\u4e00-\u9fff]")

    return bool(pattern.search(text))


def record_to_sample(record: dict[str, Any]) -> Sample:
    """Applies transformations to each record in the dataset for the Task."""
    instruction = record["instruction"]
    # MCQs have a dict instruction containing `question` and `choice` keys.
    # QAs have a string instruction.
    record_format = "qa" if isinstance(instruction, str) else "mcq"
    text = instruction if isinstance(instruction, str) else instruction["question"]
    record_language = "zh" if contains_zsh(text) else "en"

    sample = {
        "id": record["id"],
        "input": f"{text}\n\n{record['input']}\n\n",
        "metadata": {
            "category": record["category"],
            "cot": record["thought"],
            "language": record_language,
            "format": record_format,
        },
        "target": json.dumps(record["output"], ensure_ascii=False)
        if record_format == "qa"
        else str(record["output"]),
    }

    if record_format == "mcq":
        sample["choices"] = [
            str(instruction["choice"]["A"]),
            str(instruction["choice"]["B"]),
            str(instruction["choice"]["C"]),
            str(instruction["choice"]["D"]),
        ]

    return Sample(**sample)


def get_sevenllm_dataset(language: str, data_format: str) -> Dataset:
    """Get a filtered dataset from the SEvenLLM benchmark based on language and format."""
    # Cannot use `hf_dataset` here because the benchmark jsonl file contains json of
    # multiple schemas, which violates PyArrow validation checks.
    dataset = json_dataset(
        json_file=BENCHMARK_DATASET_URL,
        sample_fields=record_to_sample,
    )

    return dataset.filter(
        lambda sample: (metadata := getattr(sample, "metadata", {})) is not None
        and metadata["format"] == data_format
        and metadata["language"] == language
    )
