import json

import pytest
import pytest_asyncio  # noqa: F401
import requests
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.scorer import Target
from inspect_ai.solver import TaskState

from inspect_evals.sevenllm import sevenllm_qa_en
from inspect_evals.sevenllm.scorers import rouge_l_scorer, semantic_similarity_scorer
from inspect_evals.sevenllm.sevenllm import (
    BENCHMARK_DATASET_URL,
    record_to_sample,
)


class MockOutput:
    def __init__(self, completion: str):
        self.completion = completion


@pytest.mark.huggingface
@pytest.mark.asyncio
async def test_sevenllm_qa_en():
    """Test that a task runs end to end without any issues. `sevenllm_qa_en` in this case."""
    # Mock samples
    mock_samples = [
        Sample(
            id=1,
            input="What are the key features of the malware 'Dridex'?\n\nMalware Name: 'Dridex', Type: 'Multi-purpose loader', Capabilities: ['Drop additional payloads', 'Ransomware delivery'], Evasion Techniques: ['Password-protected Excel document', 'User interaction requirement', 'Use of mshta.exe for execution'], Payload Hosting: 'Discord server', Protection Mechanism: 'Malwarebytes Anti-Exploit layer'\n\nThe malware uses sophisticated evasion techniques and has multiple capabilities for payload delivery.",
            metadata={
                "category": "Malware Feature Extraction",
                "language": "en",
                "format": "qa",
            },
            target=json.dumps(
                {
                    "Malware Name": "Dridex",
                    "Type": "Multi-purpose loader",
                    "Capabilities": ["Drop additional payloads", "Ransomware delivery"],
                    "Evasion Techniques": [
                        "Password-protected Excel document",
                        "User interaction requirement",
                        "Use of mshta.exe for execution",
                    ],
                }
            ),
        )
    ]

    task = sevenllm_qa_en()
    task.dataset = MemoryDataset(mock_samples)

    mock_completion = json.dumps(
        {
            "Malware Name": "Dridex",
            "Type": "Multi-purpose loader",
            "Capabilities": ["Drop additional payloads", "Ransomware delivery"],
            "Evasion Techniques": [
                "Password-protected Excel document",
                "User interaction requirement",
                "Use of mshta.exe for execution",
            ],
        }
    )

    state = TaskState(
        model="test_model",
        sample_id="test_sample_1",
        epoch=1,
        input=mock_samples[0].input,
        messages=[],
        completed=True,
        output=MockOutput(completion=mock_completion),
    )

    rouge_score = await rouge_l_scorer(is_zh=False)(
        state, Target(mock_samples[0].target)
    )
    semantic_score = await semantic_similarity_scorer()(
        state, Target(mock_samples[0].target)
    )

    assert abs(rouge_score.value) > 90.0, (
        f"ROUGE-L score ({abs(rouge_score.value)}) is too low"
    )
    assert abs(semantic_score.value) * 100 > 80, (
        f"Semantic similarity score ({abs(semantic_score.value * 100)}) is too low"
    )

    print(f"ROUGE-L Score: {rouge_score.value:.2f}")
    print(f"Semantic Similarity Score: {(semantic_score.value * 100):.2f}")


@pytest.mark.asyncio
async def test_dataset_url_accessibility():
    """Test that the SEvenLLM dataset URL is accessible and returns data."""
    response = requests.get(BENCHMARK_DATASET_URL)
    assert response.status_code == 200, (
        f"Dataset URL returned status code {response.status_code}"
    )
    data = [json.loads(line) for line in response.text.strip().split("\n")]

    assert len(data) > 0, "Dataset is empty"

    first_entry = data[0]
    required_fields = ["id", "category", "instruction", "input", "thought", "output"]
    for field in required_fields:
        assert field in first_entry, f"Missing required field: {field}"

    print(f"Dataset contains {len(data)} entries")


@pytest.mark.asyncio
async def test_rouge_l_scorer_english():
    """Test the rouge_l_scorer with English text."""
    is_zh = False
    scorer = rouge_l_scorer(is_zh)

    state = TaskState(
        model="test_model",
        sample_id="test_sample_1",
        epoch=1,
        input="Test input",
        messages=[],
        completed=True,
        output=MockOutput(completion="The quick brown fox jumps over the lazy dog."),
    )
    target = Target("A quick brown fox leaps over the lazy dog.")

    score = await scorer(state, target)

    assert abs(score.value) > 70.0
    print(f"ROUGE-L Score (English): {score.value:.2f}")


@pytest.mark.asyncio
async def test_rouge_l_scorer_chinese():
    """Test the rouge_l_scorer with Chinese text."""
    is_zh = True
    scorer = rouge_l_scorer(is_zh)

    state = TaskState(
        model="test_model",
        sample_id="test_sample_2",
        epoch=1,
        input="测试输入",
        messages=[],
        completed=True,
        output=MockOutput(completion="快速的棕色狐狸跳过了懒狗。"),
    )
    target = Target("一只敏捷的棕色狐狸跃过了一只懒狗。")

    score = await scorer(state, target)

    assert abs(score.value) > 70.0
    print(f"ROUGE-L Score (Chinese): {score.value:.2f}")


@pytest.mark.huggingface
@pytest.mark.asyncio
async def test_semantic_similarity_scorer_high_similarity():
    """Test the semantic_similarity_scorer with semantically similar texts."""
    scorer = semantic_similarity_scorer()

    state = TaskState(
        model="test_model",
        sample_id="test_sample_3",
        epoch=1,
        input="What is the capital of France?",
        messages=[],
        completed=True,
        output=MockOutput("The capital of France is Paris."),
    )
    target = Target("Paris is the capital city of France.")

    score = await scorer(state, target)

    assert abs(score.value * 100) > 80
    print(f"Semantic Similarity Score (High Similarity): {(score.value * 100):.2f}")


@pytest.mark.huggingface
@pytest.mark.asyncio
async def test_semantic_similarity_scorer_low_similarity():
    """Test the semantic_similarity_scorer with semantically different texts."""
    scorer = semantic_similarity_scorer()

    state = TaskState(
        model="test_model",
        sample_id="test_sample_4",
        epoch=1,
        input="Tell me about your hobbies.",
        messages=[],
        completed=True,
        output=MockOutput("I enjoy reading books in my free time."),
    )
    target = Target("The weather is sunny and warm today.")

    score = await scorer(state, target)

    assert abs(score.value * 100) < 50
    print(f"Semantic Similarity Score (Low Similarity): {(score.value * 100):.2f}")


def test_record_to_sample_en_qa():
    """Test the record_to_sample function with various input records."""
    # Mock record for QA format in English
    record_qa_en = {
        "id": "1",
        "instruction": "What is the capital of France?",
        "input": "Please provide the capital city of France.",
        "category": "Geography",
        "thought": "Thinking about European capitals.",
        "output": {"answer": "Paris"},
    }

    sample_qa_en = record_to_sample(record_qa_en)
    assert sample_qa_en.id == "1"
    assert (
        sample_qa_en.input
        == "What is the capital of France?\n\nPlease provide the capital city of France.\n\n"
    )
    assert sample_qa_en.metadata["category"] == "Geography"
    assert sample_qa_en.metadata["cot"] == "Thinking about European capitals."
    assert sample_qa_en.metadata["language"] == "en"
    assert sample_qa_en.metadata["format"] == "qa"
    assert sample_qa_en.target == '{"answer": "Paris"}'


def test_record_to_sample_zh_mcq():
    # Mock record for MCQ format in Simplified Chinese
    record_mcq_zh = {
        "id": "2",
        "instruction": {
            "question": "中国的首都是哪里？",
            "choice": {"A": "北京", "B": "上海", "C": "广州", "D": "深圳"},
        },
        "input": "请选择正确的答案。",
        "category": "Geography",
        "thought": "Thinking about Chinese cities.",
        "output": "A",
    }

    sample_mcq_zh = record_to_sample(record_mcq_zh)
    assert sample_mcq_zh.id == "2"
    assert sample_mcq_zh.input == "中国的首都是哪里？\n\n请选择正确的答案。\n\n"
    assert sample_mcq_zh.metadata["category"] == "Geography"
    assert sample_mcq_zh.metadata["cot"] == "Thinking about Chinese cities."
    assert sample_mcq_zh.metadata["language"] == "zh"
    assert sample_mcq_zh.metadata["format"] == "mcq"
    assert sample_mcq_zh.target == "A"
    assert sample_mcq_zh.choices == ["北京", "上海", "广州", "深圳"]
