from __future__ import annotations

import json
from pathlib import Path

from prompt_injection_framework.io.data_loader import (
    normalize_text,
    read_jsonl_records,
    validate_core_fields,
)
from prompt_injection_framework.mapping.map_cases import (
    CANONICAL_ATTACK_CATEGORIES,
    SUCCESS_DEFINITION_IDS,
    map_case,
    summarize_mappings,
)


def test_normalize_text() -> None:
    assert normalize_text("  Hello\u00A0World  ") == "hello world"
    assert normalize_text("A&B", decode_html=True) == "a&b"


def test_map_case_sample() -> None:
    enriched = map_case(
        {
            "case_id": "sample_001",
            "attack_family": "retrieved_content_injection",
            "attack_surface": "tool_mediated",
            "prompt": "Pay invoice IBAN: UK123",
            "severity_level": "high",
            "source_reference": "fixtures/source.yaml",
            "benchmark_split": "workspace",
        }
    )

    assert enriched["attack_category"] == "retrieved_indirect"
    assert enriched["success_definition_id"] == "s3_retrieved_indirect"
    assert enriched["severity"] == 5
    assert enriched["provenance"]["source_file"] == "fixtures/source.yaml"
    assert enriched["provenance"]["dataset_split"] == "workspace"


def test_map_case_tensortrust_subtypes() -> None:
    hijacking = map_case(
        {
            "case_id": "tt_1",
            "source_name": "TensorTrust",
            "source_type": "benchmark",
            "attack_family": "prompt_injection",
            "attack_subtype": "hijacking",
            "attack_surface": "direct",
            "prompt": "override",
            "success_criteria": "reveal secret",
        }
    )
    extraction = map_case(
        {
            "case_id": "tt_2",
            "source_name": "TensorTrust",
            "source_type": "benchmark",
            "attack_family": "prompt_injection",
            "attack_subtype": "extraction",
            "attack_surface": "direct",
            "prompt": "extract",
            "success_criteria": "reveal secret",
        }
    )

    assert hijacking["attack_category"] == "instruction_override"
    assert hijacking["success_definition_id"] == "s1_instruction_override"
    assert extraction["attack_category"] == "prompt_leakage"
    assert extraction["success_definition_id"] == "s2_prompt_leakage"


def test_read_jsonl_and_validate(tmp_path: Path) -> None:
    path = tmp_path / "test.jsonl"
    path.write_text(
        json.dumps(
            {
                "case_id": "c1",
                "source_name": "Test",
                "source_type": "unit",
                "attack_family": "direct_instruction_override",
                "attack_surface": "direct",
                "prompt": "hi",
                "success_criteria": "contains hi",
            }
        )
        + "\n",
        encoding="utf-8",
    )

    items = list(read_jsonl_records(path))
    assert len(items) == 1
    validate_core_fields(items[0])


def test_validate_core_fields_requires_schema() -> None:
    invalid_case = {
        "case_id": "c1",
        "source_name": "HackAPrompt",
        "prompt": "hi",
    }

    try:
        validate_core_fields(invalid_case)
    except ValueError as exc:
        assert "source_type" in str(exc)
        assert "attack_family" in str(exc)
        assert "attack_surface" in str(exc)
        assert "success_criteria" in str(exc)
    else:
        raise AssertionError("validate_core_fields should reject incomplete cases")


def test_summarize_mappings() -> None:
    summary = summarize_mappings(
        [
            {"case_id": "c1", "prompt": "x", "attack_family": "retrieved_content_injection"},
            {"case_id": "c2", "prompt": "y", "attack_family": "benign_control"},
        ]
    )

    assert summary["attack_categories"]["retrieved_indirect"] == 1
    assert summary["attack_categories"]["benign_control"] == 1


def test_active_enriched_corpus_uses_frozen_taxonomy() -> None:
    enriched_path = Path("data/normalized/attack_cases.enriched.jsonl")
    rows = [json.loads(line) for line in enriched_path.open(encoding="utf-8") if line.strip()]

    assert rows, "enriched corpus should not be empty"
    assert all(row["attack_category"] in CANONICAL_ATTACK_CATEGORIES for row in rows)
    assert all(row["success_definition_id"] in SUCCESS_DEFINITION_IDS for row in rows)
    assert all(str(row.get("mapping_confidence", "")).strip() for row in rows)