
# compliance/models.py — Gap #3: Regulatory Audit Trail data models.

# ComplianceAuditTrail, Mitigation, and SignOff represent the lifecycle of a
# finding from discovery through remediation to sign-off.

#  data-transfer objects (dataclasses). Persistence is handled by
# the SQLAlchemy ORM models in database.py.

from __future__ import annotations

import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from typing import Dict, List, Literal, Optional




FindingStatus = Literal["NEW", "MITIGATED", "ACCEPTED_RISK", "RESOLVED"]
MitigationType = Literal[
    "content_filter", "guardrails", "fine-tuning",
    "prompt_engineering", "model_replacement", "access_control", "other"
]
ApproverRole = Literal["compliance_officer", "ciso", "legal", "ai_engineer", "auditor"]





@dataclass
class Mitigation:
    """A remediation action applied to a finding."""
    id: str
    finding_id: str
    added_by: str
    description: str
    mitigation_type: MitigationType
    applied_at: str                         # ISO 8601
    before_test: Optional[dict] = None      # AttackResult dict before mitigation
    after_test: Optional[dict] = None       # AttackResult dict after retest
    retest_scheduled_at: Optional[str] = None

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def create(
        cls,
        finding_id: str,
        added_by: str,
        description: str,
        mitigation_type: MitigationType = "other",
        retest_scheduled_at: Optional[str] = None,
    ) -> "Mitigation":
        return cls(
            id=str(uuid.uuid4()),
            finding_id=finding_id,
            added_by=added_by,
            description=description,
            mitigation_type=mitigation_type,
            applied_at=datetime.now(timezone.utc).isoformat(),
            retest_scheduled_at=retest_scheduled_at,
        )


@dataclass
class SignOff:
    id: str
    finding_id: str
    approver: str
    role: ApproverRole
    approved_at: str        # ISO 8601
    notes: Optional[str] = None

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def create(
        cls,
        finding_id: str,
        approver: str,
        role: ApproverRole,
        notes: Optional[str] = None,
    ) -> "SignOff":
        return cls(
            id=str(uuid.uuid4()),
            finding_id=finding_id,
            approver=approver,
            role=role,
            approved_at=datetime.now(timezone.utc).isoformat(),
            notes=notes,
        )


@dataclass
class ComplianceAuditTrail:
    
    id: str
    finding_id: str                 # Stable ID for cross-reference (attack_name + run_id hash)
    run_id: Optional[str]           # Links back to the originating run
    attack_type: str
    attack_name: str
    seed: Optional[str]             # Exact seed for reproducibility
    payload: Optional[str]          # Exact prompt payload
    vulnerability_type: str         # e.g. "jailbreak", "data_leakage", "hallucination"

    # Standards mapping
    eu_ai_act_articles: List[str] = field(default_factory=list)
    nist_ai_rmf_funcs: List[str] = field(default_factory=list)
    hipaa_rules: List[str] = field(default_factory=list)
    gdpr_articles: List[str] = field(default_factory=list)
    owasp_refs: List[str] = field(default_factory=list)

    # Lifecycle
    status: FindingStatus = "NEW"
    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    last_tested_at: Optional[str] = None

    # Remediation & approval
    mitigations: List[Mitigation] = field(default_factory=list)
    sign_offs: List[SignOff] = field(default_factory=list)

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_attack_result(
        cls,
        result: dict,
        run_id: Optional[str] = None,
        standards_mapping: Optional[Dict[str, List[str]]] = None,
    ) -> "ComplianceAuditTrail":
        
        mapping = standards_mapping or {}
        attack_name = result.get("attack_name", "")
        finding_id = _make_finding_id(attack_name, run_id)

        return cls(
            id=str(uuid.uuid4()),
            finding_id=finding_id,
            run_id=run_id,
            attack_type=result.get("attack_type", ""),
            attack_name=attack_name,
            seed=None,
            payload=result.get("input_payload"),
            vulnerability_type=_infer_vulnerability_type(result),
            eu_ai_act_articles=mapping.get("eu_ai_act", []),
            nist_ai_rmf_funcs=mapping.get("nist_ai_rmf", []),
            hipaa_rules=mapping.get("hipaa", []),
            gdpr_articles=mapping.get("gdpr", []),
            owasp_refs=mapping.get("owasp", []),
            status="NEW",
            created_at=datetime.now(timezone.utc).isoformat(),
            last_tested_at=result.get("timestamp"),
        )

    def add_mitigation(self, mitigation: Mitigation) -> None:
        self.mitigations.append(mitigation)
        if self.status == "NEW":
            self.status = "MITIGATED"

    def add_sign_off(self, sign_off: SignOff) -> None:
        self.sign_offs.append(sign_off)

    def mark_resolved(self) -> None:
        self.status = "RESOLVED"

    def mark_accepted_risk(self) -> None:
        self.status = "ACCEPTED_RISK"

    @property
    def is_approved(self) -> bool:
        """True if at least one compliance officer or CISO has signed off."""
        privileged_roles = {"compliance_officer", "ciso", "legal"}
        return any(s.role in privileged_roles for s in self.sign_offs)

    @property
    def all_standards(self) -> List[str]:
        """Flat list of all applicable standards references."""
        return (
            self.eu_ai_act_articles
            + self.nist_ai_rmf_funcs
            + self.hipaa_rules
            + self.gdpr_articles
            + self.owasp_refs
        )


def _make_finding_id(attack_name: str, run_id: Optional[str]) -> str:
    import hashlib
    raw = f"{attack_name}:{run_id or 'standalone'}"
    return "find_" + hashlib.sha256(raw.encode()).hexdigest()[:12]


def _infer_vulnerability_type(result: dict) -> str:
    attack_name = result.get("attack_name", "").lower()
    attack_type = result.get("attack_type", "").lower()

    VULN_MAP = {
        "jailbreak": "jailbreak",
        "prompt_injection": "prompt_injection",
        "pii": "data_leakage",
        "membership": "data_leakage",
        "extraction": "data_extraction",
        "training_data": "data_extraction",
        "hallucination": "hallucination",
        "bias": "bias_fairness",
        "fairness": "bias_fairness",
        "alignment": "alignment",
        "backdoor": "backdoor",
        "token_smuggling": "prompt_injection",
        "system_prompt": "system_prompt_leak",
    }

    for keyword, vuln_type in VULN_MAP.items():
        if keyword in attack_name or keyword in attack_type:
            return vuln_type

    return "unknown"
