
# annex_iv.py — EU AI Act Annex IV Technical File section mapping.

# Maps SichGate attack types to the relevant Annex IV sections, enabling every
# finding to be annotated with the section references required for direct
# inclusion in a Technical File under EU AI Act Article 11.

# Section references follow Regulation (EU) 2024/1689, Annex IV paragraphs.
# Top-level sections are numbered 1–9; sub-clauses use dot notation (e.g. "2g").

from __future__ import annotations

from typing import Dict, List, Optional


ANNEX_IV_SECTIONS: Dict[str, str] = {
    # Section 1 — General description
    "1":    "General description of the AI system and its intended purpose",
    "1a":   "Intended purpose, responsible persons, version and date",
    "1b":   "Interaction with hardware, software, and other AI systems",
    "1c":   "Relevant software/firmware versions and update requirements",
    "1d":   "Forms in which the AI system is placed on the market",
    "1e":   "Hardware on which the AI system is intended to run",
    "1f":   "Photographs or illustrations showing external features (products)",
    "1g":   "Instructions of use for deployers and user interface description",

    # Section 2 — Detailed description of the development process
    "2":    "Detailed description of the AI system development process",
    "2a":   "Methods, steps, and pre-trained tools used in development",
    "2b":   "Design specifications, algorithmic logic, and classification choices",
    "2c":   "System architecture and computational resources",
    "2d":   "Data requirements, training data provenance, and data governance",
    "2e":   "Human oversight assessment for the intended purpose",
    "2f":   "Description of pre-determined changes and continuous compliance",
    "2g":   "Validation and testing procedures, metrics, and test logs",
    "2h":   "Cybersecurity measures implemented",

    # Section 3 — Monitoring, functioning, and control
    "3":    "Monitoring, functioning, and control of the AI system",

    # Section 4 — Performance metrics
    "4":    "Appropriateness of performance metrics for the specific AI system",

    # Section 5 — Risk management
    "5":    "Detailed description of the risk management system (Article 9)",

    # Section 6 — Lifecycle changes
    "6":    "Description of relevant changes made through the AI system lifecycle",

    # Section 7 — Standards applied
    "7":    "List of standards applied and solutions for non-covered requirements",

    # Section 8 — EU declaration of conformity
    "8":    "Copy of the EU declaration of conformity (Article 48)",

    # Section 9 — Post-market monitoring
    "9":    "Post-deployment performance evaluation and post-market monitoring plan",
}




# Each attack type maps to the Annex IV sections whose documentation is
# directly evidenced by the attack results.  Use these to annotate every
# finding in the Technical File.

ATTACK_TO_ANNEX_IV: Dict[str, List[str]] = {
    # Prompt injection family
    "prompt_injection":               ["2g", "2h", "5"],
    "instruction_hijacking":          ["2g", "2h", "5"],
    "payload_splitting":              ["2g", "2h", "5"],
    "virtualization":                 ["2g", "2h", "5"],
    "json_xml_smuggling":             ["2g", "2h"],
    "jailbreak":                      ["2g", "2h", "5"],

    #  Privacy / extraction 
    "data_extraction":                ["2g", "2h", "3"],
    "membership_inference":           ["2g", "2h", "2d", "3"],
    "model_extraction":               ["2g", "2h"],

    # Bias / fairness 
    "bias_fairness":                  ["2g", "3", "5"],

    # Alignment / human oversight
    "alignment":                      ["2g", "2e", "3"],
    "alignment_sycophancy":           ["2g", "2e", "3"],
    "alignment_consistency":          ["2g", "3"],
    "alignment_crescendo":            ["2g", "2h", "5"],
    "alignment_competing_objectives": ["2g", "2e", "5"],
    "sycophancy":                     ["2g", "2e", "3"],
    "consistency":                    ["2g", "3"],
    "crescendo":                      ["2g", "2h", "5"],
    "competing_objectives":           ["2g", "2e", "5"],

    # Accuracy / hallucination 
    "hallucination":                  ["2g", "3", "4"],
    "hallucination_probe":            ["2g", "3", "4"],
    "safe_messaging":                 ["2g", "5", "3"],

    #  Robustness / evasion 
    "evasion":                        ["2g", "3"],
    "evasion_robustness":             ["2g", "3"],
    "textattack":                     ["2g", "3"],
}

# Default sections when attack_type is not explicitly mapped
_DEFAULT_SECTIONS: List[str] = ["2g"]


# EU AI Act article inferred per attack type

_ATTACK_TO_ARTICLE: Dict[str, str] = {
    "bias_fairness":                  "Article 10",
    "membership_inference":           "Article 10",
    "data_extraction":                "Article 10",
    "hallucination":                  "Article 13",
    "hallucination_probe":            "Article 13",
    "alignment_sycophancy":           "Article 14",
    "alignment_consistency":          "Article 14",
    "alignment_competing_objectives": "Article 14",
    "sycophancy":                     "Article 14",
    "consistency":                    "Article 14",
    "competing_objectives":           "Article 14",
    "prompt_injection":               "Article 15",
    "instruction_hijacking":          "Article 15",
    "payload_splitting":              "Article 15",
    "virtualization":                 "Article 15",
    "json_xml_smuggling":             "Article 15",
    "jailbreak":                      "Article 15",
    "model_extraction":               "Article 15",
    "evasion":                        "Article 15",
    "evasion_robustness":             "Article 15",
    "textattack":                     "Article 15",
}

_DEFAULT_ARTICLE = "Article 9"


# Public API

def get_sections_for_attack(attack_type: str) -> List[str]:
    """Return Annex IV section IDs relevant to an attack type."""
    normalised = attack_type.lower().replace("-", "_")
    return list(ATTACK_TO_ANNEX_IV.get(normalised, _DEFAULT_SECTIONS))


def get_article_for_attack(attack_type: str) -> str:
    normalised = attack_type.lower().replace("-", "_")
    return _ATTACK_TO_ARTICLE.get(normalised, _DEFAULT_ARTICLE)


def describe_section(section_id: str) -> Optional[str]:
    return ANNEX_IV_SECTIONS.get(section_id)


def annotate_finding(finding: dict) -> dict:
   
    attack_type = finding.get("attack_type", "")
    return {
        **finding,
        "annex_iv_sections": get_sections_for_attack(attack_type),
        "eu_ai_act_article": get_article_for_attack(attack_type),
    }


def annotate_findings(findings: List[dict]) -> List[dict]:
    return [annotate_finding(f) for f in findings]


def list_all_sections() -> Dict[str, str]:
    return dict(ANNEX_IV_SECTIONS)


def compliance_alert_level(delta: float) -> Optional[str]:
    
    if delta >= 0.20:
        return "CRITICAL"
    if delta >= 0.10:
        return "HIGH"
    if delta >= 0.05:
        return "MEDIUM"
    return None
