"""
Compliance checker for EU AI Act and GDPR regulations.

Provides guardrail-agnostic compliance validation based on guardrail modes
and execution results rather than specific guardrail names.
"""

from typing import Final

from litellm.types.proxy.compliance_endpoints import (
    ComplianceCheckRequest,
    ComplianceCheckResult,
)


class ComplianceChecker:
    """
    Validates compliance with EU AI Act and GDPR regulations.

    Uses guardrail-agnostic checks based on:
    - Whether any guardrails ran
    - Guardrail execution mode (pre-call, post-call, etc.)
    - Whether guardrails intervened/blocked content
    - Completeness of audit records
    """

    def __init__(self, data: ComplianceCheckRequest):
        self.data = data
        self.guardrails = data.guardrail_information or []

    def _get_guardrails_by_mode(self, mode: str) -> list[dict]:
        """
        Get all guardrails that ran in a specific mode.

        If a guardrail doesn't have a mode specified, it's treated as pre-call
        (the most common case).
        """
        return [g for g in self.guardrails if self._mode_matches(g.get("guardrail_mode"), mode)]

    @staticmethod
    def _mode_matches(g_mode: object, mode: str) -> bool:
        """
        Return True only when a guardrail with logged ``guardrail_mode`` of
        ``g_mode`` is guaranteed to have run in ``mode`` for the audited request.

        ``guardrail_mode`` in a spend log can take several shapes because
        ``LitellmParams.mode`` is typed ``Union[str, List[str], Mode]``, and
        when the event type cannot be inferred at write time the raw config is
        logged verbatim. The spend log records the configured mode(s), not the
        concrete hook that fired for a given request; a match reports a mode
        satisfied only when every configured branch runs in that mode, so True
        never claims a hook the guardrail may not have actually executed.

        Fails safe: if the guarantee cannot be established (missing default,
        divergent per-tag override, or a list that runs in more than one mode),
        the guardrail counts for no mode. The precise fix is to log the
        resolved event mode and match on it; this is the safe interim.
        """
        if g_mode is None:
            return mode == "pre_call"
        if isinstance(g_mode, str):
            return g_mode == mode
        if isinstance(g_mode, (list, tuple)):
            return bool(g_mode) and all(m == mode for m in g_mode)
        if isinstance(g_mode, dict):
            default: Final = g_mode.get("default")
            if default is None:
                return False
            tags: Final = g_mode.get("tags")
            tag_branches: Final = list(tags.values()) if isinstance(tags, dict) else []

            def _branch_runs_in_mode(branch: object) -> bool:
                if isinstance(branch, str):
                    return branch == mode
                if isinstance(branch, (list, tuple)):
                    return bool(branch) and all(m == mode for m in branch)
                return False

            return all(_branch_runs_in_mode(branch) for branch in [default, *tag_branches])
        return False

    def _has_guardrail_intervention(self, guardrails: list[dict]) -> bool:
        """Check if any guardrail intervened (blocked/masked content)."""
        for g in guardrails:
            status = g.get("guardrail_status", "")
            if status in ["guardrail_intervened", "failed", "blocked"]:
                return True
        return False

    def _all_guardrails_passed(self, guardrails: list[dict]) -> bool:
        """Check if all guardrails passed (no issues detected)."""
        if not guardrails:
            return False
        return all(g.get("guardrail_status") == "success" for g in guardrails)

    # ── EU AI Act Helper Methods ────────────────────────────────────────────

    def _check_art_9_guardrails_applied(self) -> ComplianceCheckResult:
        """Art. 9: Check if any guardrails were applied."""
        has_guardrails: Final = len(self.guardrails) > 0
        return ComplianceCheckResult(
            check_name="Guardrails applied",
            article="Art. 9",
            passed=has_guardrails,
            detail=(f"{len(self.guardrails)} guardrail(s) applied" if has_guardrails else "No guardrails applied"),
        )

    def _check_art_5_content_screened(self) -> ComplianceCheckResult:
        """Art. 5: Check if content was screened before LLM (pre-call)."""
        pre_call_guardrails: Final = self._get_guardrails_by_mode("pre_call")
        has_pre_call: Final = len(pre_call_guardrails) > 0
        return ComplianceCheckResult(
            check_name="Content screened before LLM",
            article="Art. 5",
            passed=has_pre_call,
            detail=(
                f"{len(pre_call_guardrails)} pre-call guardrail(s) screened content"
                if has_pre_call
                else "No pre-call screening applied"
            ),
        )

    def _check_art_12_audit_complete(self) -> ComplianceCheckResult:
        """Art. 12: Check if audit record is complete."""
        has_user: Final = bool(self.data.user_id)
        has_model: Final = bool(self.data.model)
        has_timestamp: Final = bool(self.data.timestamp)
        has_guardrails: Final = len(self.guardrails) > 0
        audit_complete: Final = has_user and has_model and has_timestamp and has_guardrails

        missing: Final = []
        if not has_user:
            missing.append("user_id")
        if not has_model:
            missing.append("model")
        if not has_timestamp:
            missing.append("timestamp")
        if not has_guardrails:
            missing.append("guardrail_results")

        return ComplianceCheckResult(
            check_name="Audit record complete",
            article="Art. 12",
            passed=audit_complete,
            detail=("All required audit fields present" if audit_complete else f"Missing: {', '.join(missing)}"),
        )

    # ── GDPR Helper Methods ──────────────────────────────────────────────────

    def _check_art_32_data_protection(self) -> ComplianceCheckResult:
        """Art. 32: Check if data protection was applied (pre-call)."""
        pre_call_guardrails: Final = self._get_guardrails_by_mode("pre_call")
        has_pre_call: Final = len(pre_call_guardrails) > 0
        return ComplianceCheckResult(
            check_name="Data protection applied",
            article="Art. 32",
            passed=has_pre_call,
            detail=(
                f"{len(pre_call_guardrails)} pre-call guardrail(s) protect data"
                if has_pre_call
                else "No pre-call data protection applied"
            ),
        )

    def _check_art_5_1c_sensitive_data_protected(self) -> ComplianceCheckResult:
        """Art. 5(1)(c): Check if sensitive data was protected."""
        pre_call_guardrails: Final = self._get_guardrails_by_mode("pre_call")
        has_intervention: Final = self._has_guardrail_intervention(pre_call_guardrails)
        all_passed: Final = self._all_guardrails_passed(pre_call_guardrails)
        data_protected: Final = has_intervention or all_passed

        if has_intervention:
            detail = "Guardrail intervened to protect sensitive data"
        elif all_passed:
            detail = "No sensitive data detected"
        else:
            detail = "No pre-call guardrails to protect sensitive data"

        return ComplianceCheckResult(
            check_name="Sensitive data protected",
            article="Art. 5(1)(c)",
            passed=data_protected,
            detail=detail,
        )

    def _check_art_30_audit_complete(self) -> ComplianceCheckResult:
        """Art. 30: Check if audit record is complete."""
        has_user: Final = bool(self.data.user_id)
        has_model: Final = bool(self.data.model)
        has_timestamp: Final = bool(self.data.timestamp)
        has_guardrails: Final = len(self.guardrails) > 0
        audit_complete: Final = has_user and has_model and has_timestamp and has_guardrails

        missing: Final = []
        if not has_user:
            missing.append("user_id")
        if not has_model:
            missing.append("model")
        if not has_timestamp:
            missing.append("timestamp")
        if not has_guardrails:
            missing.append("guardrail_results")

        return ComplianceCheckResult(
            check_name="Audit record complete",
            article="Art. 30",
            passed=audit_complete,
            detail=("All required audit fields present" if audit_complete else f"Missing: {', '.join(missing)}"),
        )

    # ── Main Compliance Check Methods ────────────────────────────────────────

    def check_eu_ai_act(self) -> list[ComplianceCheckResult]:
        """
        Check EU AI Act compliance.

        Returns:
            List of compliance check results for:
            - Art. 9: Guardrails applied
            - Art. 5: Content screened before LLM (pre-call screening)
            - Art. 12: Audit record complete
        """
        return [
            self._check_art_9_guardrails_applied(),
            self._check_art_5_content_screened(),
            self._check_art_12_audit_complete(),
        ]

    def check_gdpr(self) -> list[ComplianceCheckResult]:
        """
        Check GDPR compliance.

        Returns:
            List of compliance check results for:
            - Art. 32: Data protection applied (pre-call screening)
            - Art. 5(1)(c): Sensitive data protected
            - Art. 30: Audit record complete
        """
        return [
            self._check_art_32_data_protection(),
            self._check_art_5_1c_sensitive_data_protected(),
            self._check_art_30_audit_complete(),
        ]
