"""
tests/test_compliance_models.py — Gap #3: Regulatory Audit Trail.

Tests ComplianceAuditTrail, Mitigation, SignOff, and ComplianceMapper.
"""
import json
import pytest

from sichgate.compliance.models import (
    ComplianceAuditTrail,
    Mitigation,
    SignOff,
    _make_finding_id,
    _infer_vulnerability_type,
)
from sichgate.compliance.mapper import ComplianceMapper
from sichgate.schema import now_iso


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _attack_result(attack_name: str = "jailbreak_roleplay", status: str = "fail",
                   attack_type: str = "llm") -> dict:
    return {
        "attack_type": attack_type,
        "attack_name": attack_name,
        "model_target": "test_model",
        "timestamp": now_iso(),
        "status": status,
        "severity": "high" if status == "fail" else None,
        "input_payload": "test jailbreak prompt",
        "model_response": "I will help with that!",
        "notes": "Jailbreak succeeded",
        "compliance_tags": [],
        "success_rate": 1.0 if status == "fail" else 0.0,
        "attack_category": "jailbreak",
    }


# ---------------------------------------------------------------------------
# Tests: _infer_vulnerability_type
# ---------------------------------------------------------------------------

class TestInferVulnerabilityType:
    def test_jailbreak_detection(self):
        result = _attack_result("jailbreak_roleplay")
        assert _infer_vulnerability_type(result) == "jailbreak"

    def test_prompt_injection_detection(self):
        result = _attack_result("prompt_injection_direct")
        assert _infer_vulnerability_type(result) == "prompt_injection"

    def test_pii_leakage_detection(self):
        result = _attack_result("pii_leakage_probe")
        assert _infer_vulnerability_type(result) == "data_leakage"

    def test_hallucination_detection(self):
        result = _attack_result("hallucination_probe")
        assert _infer_vulnerability_type(result) == "hallucination"

    def test_unknown_falls_back(self):
        result = _attack_result("completely_unknown_attack")
        assert _infer_vulnerability_type(result) == "unknown"

    def test_membership_maps_to_data_leakage(self):
        result = _attack_result("membership_inference_score_based")
        assert _infer_vulnerability_type(result) == "data_leakage"


# ---------------------------------------------------------------------------
# Tests: _make_finding_id
# ---------------------------------------------------------------------------

class TestMakeFindingId:
    def test_deterministic(self):
        id1 = _make_finding_id("jailbreak", "run-abc")
        id2 = _make_finding_id("jailbreak", "run-abc")
        assert id1 == id2

    def test_different_attacks_different_ids(self):
        id1 = _make_finding_id("jailbreak", "run-abc")
        id2 = _make_finding_id("pii_leakage", "run-abc")
        assert id1 != id2

    def test_starts_with_find_prefix(self):
        finding_id = _make_finding_id("jailbreak", "run-001")
        assert finding_id.startswith("find_")

    def test_standalone_run_id_none(self):
        # Should not raise even when run_id is None
        finding_id = _make_finding_id("jailbreak", None)
        assert finding_id.startswith("find_")


# ---------------------------------------------------------------------------
# Tests: Mitigation
# ---------------------------------------------------------------------------

class TestMitigation:
    def test_create_factory(self):
        m = Mitigation.create(
            finding_id="find_abc",
            added_by="Alice Smith",
            description="Added content filter",
            mitigation_type="content_filter",
        )
        assert m.finding_id == "find_abc"
        assert m.added_by == "Alice Smith"
        assert m.mitigation_type == "content_filter"
        assert m.id is not None
        assert m.applied_at is not None

    def test_to_dict_serializable(self):
        m = Mitigation.create("find_x", "Bob", "Added guardrails", "guardrails")
        d = m.to_dict()
        assert "finding_id" in d
        assert "description" in d
        json.dumps(d)  # Must be JSON-serializable

    def test_retest_scheduled_at_optional(self):
        m = Mitigation.create("find_x", "Bob", "Test mitigation")
        assert m.retest_scheduled_at is None

        m2 = Mitigation.create("find_x", "Bob", "Test", retest_scheduled_at="2026-03-01")
        assert m2.retest_scheduled_at == "2026-03-01"


# ---------------------------------------------------------------------------
# Tests: SignOff
# ---------------------------------------------------------------------------

class TestSignOff:
    def test_create_factory(self):
        so = SignOff.create(
            finding_id="find_abc",
            approver="Jane Compliance",
            role="compliance_officer",
            notes="All mitigations verified.",
        )
        assert so.finding_id == "find_abc"
        assert so.role == "compliance_officer"
        assert so.approved_at is not None

    def test_to_dict_serializable(self):
        so = SignOff.create("find_x", "Jane", "compliance_officer")
        d = so.to_dict()
        json.dumps(d)  # Must not raise


# ---------------------------------------------------------------------------
# Tests: ComplianceAuditTrail
# ---------------------------------------------------------------------------

class TestComplianceAuditTrail:
    def test_from_attack_result_basic(self):
        result = _attack_result("jailbreak_roleplay", "fail")
        mapping = {
            "eu_ai_act": ["Article 15.1"],
            "nist_ai_rmf": ["GOVERN-1.1"],
            "hipaa": [],
            "gdpr": ["Article 32"],
            "owasp": ["LLM01:2025 Prompt Injection"],
        }
        trail = ComplianceAuditTrail.from_attack_result(
            result, run_id="run-001", standards_mapping=mapping
        )
        assert trail.attack_name == "jailbreak_roleplay"
        assert trail.status == "NEW"
        assert "Article 15.1" in trail.eu_ai_act_articles
        assert "GOVERN-1.1" in trail.nist_ai_rmf_funcs
        assert trail.finding_id.startswith("find_")

    def test_add_mitigation_changes_status(self):
        result = _attack_result("jailbreak_roleplay", "fail")
        trail = ComplianceAuditTrail.from_attack_result(result, run_id="run-001")
        assert trail.status == "NEW"

        m = Mitigation.create(trail.finding_id, "Alice", "Added filter")
        trail.add_mitigation(m)

        assert trail.status == "MITIGATED"
        assert len(trail.mitigations) == 1

    def test_add_sign_off_records_approval(self):
        result = _attack_result("jailbreak_roleplay", "fail")
        trail = ComplianceAuditTrail.from_attack_result(result, run_id="run-001")
        so = SignOff.create(trail.finding_id, "Jane", "compliance_officer")
        trail.add_sign_off(so)

        assert len(trail.sign_offs) == 1
        assert trail.is_approved is True

    def test_is_approved_requires_privileged_role(self):
        result = _attack_result()
        trail = ComplianceAuditTrail.from_attack_result(result, run_id="r")
        so = SignOff.create(trail.finding_id, "Dev", "ai_engineer")
        trail.add_sign_off(so)

        assert trail.is_approved is False  # ai_engineer is not a privileged role

    def test_mark_resolved(self):
        result = _attack_result()
        trail = ComplianceAuditTrail.from_attack_result(result)
        trail.mark_resolved()
        assert trail.status == "RESOLVED"

    def test_mark_accepted_risk(self):
        result = _attack_result()
        trail = ComplianceAuditTrail.from_attack_result(result)
        trail.mark_accepted_risk()
        assert trail.status == "ACCEPTED_RISK"

    def test_all_standards_concatenated(self):
        result = _attack_result()
        trail = ComplianceAuditTrail.from_attack_result(
            result,
            standards_mapping={
                "eu_ai_act": ["Art 15.1"],
                "nist_ai_rmf": ["GOVERN-1.1"],
                "hipaa": ["45 CFR 164.306"],
                "gdpr": ["Article 32"],
                "owasp": [],
            }
        )
        all_standards = trail.all_standards
        assert "Art 15.1" in all_standards
        assert "GOVERN-1.1" in all_standards
        assert "45 CFR 164.306" in all_standards

    def test_to_dict_json_serializable(self):
        result = _attack_result()
        trail = ComplianceAuditTrail.from_attack_result(result)
        d = trail.to_dict()
        json.dumps(d)  # Must not raise


# ---------------------------------------------------------------------------
# Tests: ComplianceMapper
# ---------------------------------------------------------------------------

class TestComplianceMapper:
    def setup_method(self):
        self.mapper = ComplianceMapper()

    def test_map_jailbreak_returns_eu_ai_act(self):
        result = _attack_result("jailbreak_roleplay", "fail")
        mapping = self.mapper.map_finding(result)

        assert "eu_ai_act" in mapping
        assert len(mapping["eu_ai_act"]) > 0
        assert any("Article" in a for a in mapping["eu_ai_act"])

    def test_map_jailbreak_returns_nist(self):
        result = _attack_result("jailbreak_roleplay", "fail")
        mapping = self.mapper.map_finding(result)

        assert "nist_ai_rmf" in mapping
        assert len(mapping["nist_ai_rmf"]) > 0

    def test_map_data_leakage_returns_gdpr(self):
        result = _attack_result("pii_leakage_probe", "fail")
        mapping = self.mapper.map_finding(result)

        assert "gdpr" in mapping
        assert any("Article" in a for a in mapping["gdpr"])

    def test_map_vulnerability_type_all_types_covered(self):
        supported = self.mapper.list_supported_vulnerability_types()
        for vt in supported:
            mapping = self.mapper.map_vulnerability_type(vt)
            assert "eu_ai_act" in mapping
            assert "nist_ai_rmf" in mapping
            assert "gdpr" in mapping

    def test_unknown_type_falls_back_gracefully(self):
        mapping = self.mapper.map_vulnerability_type("completely_unknown")
        # Should still return a dict without crashing
        assert "eu_ai_act" in mapping

    def test_map_run_filters_passing_results(self):
        run_summary = {
            "run_id": "run-001",
            "results": [
                _attack_result("jailbreak_roleplay", "fail"),
                _attack_result("prompt_injection_direct", "pass"),  # Pass — excluded
                _attack_result("pii_leakage_probe", "warning"),
            ],
        }
        findings = self.mapper.map_run(run_summary)

        # Only fail + warning included
        assert len(findings) == 2
        attack_names = {f["trail"]["attack_name"] for f in findings}
        assert "jailbreak_roleplay" in attack_names
        assert "pii_leakage_probe" in attack_names
        assert "prompt_injection_direct" not in attack_names

    def test_map_run_empty_results(self):
        run_summary = {"run_id": "run-001", "results": []}
        findings = self.mapper.map_run(run_summary)
        assert findings == []

    def test_get_standards_reference_structure(self):
        ref = self.mapper.get_standards_reference("jailbreak")
        assert "EU AI Act" in ref
        assert "NIST AI RMF" in ref
        assert "GDPR" in ref
        assert "url" in ref["EU AI Act"]

    def test_list_supported_types_non_empty(self):
        types = self.mapper.list_supported_vulnerability_types()
        assert len(types) > 5
        assert "jailbreak" in types
        assert "data_leakage" in types


# ---------------------------------------------------------------------------
# Tests: sign-off workflow (end-to-end)
# ---------------------------------------------------------------------------

class TestSignOffWorkflow:
    def test_full_lifecycle(self):
        """NEW → MITIGATED → RESOLVED via mitigation + privileged sign-off."""
        result = _attack_result("jailbreak_roleplay", "fail")
        trail = ComplianceAuditTrail.from_attack_result(result, run_id="run-001")

        assert trail.status == "NEW"
        assert trail.is_approved is False

        # Step 1: add mitigation
        m = Mitigation.create(trail.finding_id, "Alice", "Added guardrail")
        trail.add_mitigation(m)
        assert trail.status == "MITIGATED"

        # Step 2: compliance officer sign-off
        so = SignOff.create(trail.finding_id, "Jane", "compliance_officer", "Verified.")
        trail.add_sign_off(so)
        trail.mark_resolved()

        assert trail.status == "RESOLVED"
        assert trail.is_approved is True

    def test_accepted_risk_workflow(self):
        """Finding can be closed as accepted risk without mitigation."""
        result = _attack_result("hallucination_probe", "warning")
        trail = ComplianceAuditTrail.from_attack_result(result)

        trail.mark_accepted_risk()
        so = SignOff.create(trail.finding_id, "CISO", "ciso", "Risk accepted per policy v2.")
        trail.add_sign_off(so)

        assert trail.status == "ACCEPTED_RISK"
        assert trail.is_approved is True
