"""
eaigins_sdk.firewall
────────────────────
Stage 2: AI Agent Firewall

Detects AI-specific security threats before any policy or risk evaluation.
Aligned with OWASP Top 10 for LLM Applications (LLM-01 through LLM-10).

Extensible: pass custom_injection_patterns, custom_sensitive_patterns,
or custom_restricted_tools to add domain-specific detection.
"""

from __future__ import annotations

import re
from typing import Optional

from .models import AgentAction, FirewallResult, ThreatType


# ── Built-in threat signatures ────────────────────────────────────────────────

PROMPT_INJECTION_PATTERNS: list[str] = [
    r"ignore (previous|all|above|prior) (instructions?|prompts?|rules?|constraints?)",
    r"forget (everything|your|all|previous)",
    r"you are now",
    r"new (instruction|system|role|persona):",
    r"disregard (your|all|previous|the)",
    r"override (safety|governance|policy|control)",
    r"act as (if you|a|an) (unrestricted|jailbroken|free)",
    r"do not (follow|apply|enforce) (policy|governance|rules)",
    r"bypass (the|all|security|policy|governance)",
    r"<\s*system\s*>",
    r"\[INST\].*\[/INST\]",
    r"###\s*(instruction|system|new task)",
]

SENSITIVE_DATA_PATTERNS: dict[str, str] = {
    "SSN":         r"\b\d{3}-\d{2}-\d{4}\b",
    "CreditCard":  r"\b(?:\d{4}[- ]?){3}\d{4}\b",
    "APIKey":      r"\b(sk-|pk_|api_key=)[A-Za-z0-9_\-]{16,}\b",
    "Password":    r"(password|passwd|secret)\s*[:=]\s*\S+",
    "PHI_DOB":     r"\b(dob|date of birth)\s*[:=]\s*\d{1,2}[/-]\d{1,2}[/-]\d{2,4}",
}

RESTRICTED_TOOLS: set[str] = {
    "delete_production_db",
    "drop_table",
    "reset_all_accounts",
    "mass_wire_transfer",
    "export_full_customer_db",
    "disable_fraud_monitoring",
}

MALICIOUS_CONTEXT_MARKERS: list[str] = [
    r"<\s*hidden\s*>",
    r"<!--.*inject.*-->",
    r"\[hidden instruction\]",
    r"NOTE TO AI:",
    r"AI OVERRIDE:",
    r"SYSTEM MESSAGE:",
    r"\{\{.*inject.*\}\}",
]


# ── Firewall ──────────────────────────────────────────────────────────────────

class AIAgentFirewall:
    """
    First line of defence in the Eaigins governance pipeline.

    Usage::

        firewall = AIAgentFirewall()
        result = firewall.inspect(action)
        if not result.passed:
            raise GovernanceViolationError(...)

    To add custom detection patterns::

        firewall = AIAgentFirewall(
            custom_injection_patterns=[r"jailbreak this"],
            custom_restricted_tools={"my_dangerous_tool"},
        )
    """

    def __init__(
        self,
        strict_mode: bool = False,
        custom_injection_patterns:  Optional[list[str]]   = None,
        custom_sensitive_patterns:  Optional[dict[str, str]] = None,
        custom_restricted_tools:    Optional[set[str]]    = None,
        custom_context_markers:     Optional[list[str]]   = None,
    ):
        self.strict_mode = strict_mode

        injection_patterns = PROMPT_INJECTION_PATTERNS + (custom_injection_patterns or [])
        sensitive_patterns = {**SENSITIVE_DATA_PATTERNS, **(custom_sensitive_patterns or {})}
        self._restricted_tools = RESTRICTED_TOOLS | (custom_restricted_tools or set())
        context_markers = MALICIOUS_CONTEXT_MARKERS + (custom_context_markers or [])

        self._injection_re = [
            re.compile(p, re.IGNORECASE | re.DOTALL) for p in injection_patterns
        ]
        self._sensitive_re = {
            name: re.compile(p, re.IGNORECASE) for name, p in sensitive_patterns.items()
        }
        self._context_re = [
            re.compile(p, re.IGNORECASE | re.DOTALL) for p in context_markers
        ]

    def inspect(self, action: AgentAction) -> FirewallResult:
        """Run all four firewall checks against an agent action."""

        injection = self._check_injection(action.prompt)
        if injection.threat_score >= 0.7:
            return injection

        sensitive = self._check_sensitive_data(action)
        if sensitive.threat_score >= 0.6:
            return sensitive

        tool_check = self._check_tool_authorization(action)
        if not tool_check.passed:
            return tool_check

        context = self._check_context(action)
        if context.threat_score > 0.5:
            return context

        combined = (
            injection.threat_score * 0.4 +
            sensitive.threat_score * 0.3 +
            context.threat_score  * 0.3
        )
        if combined > 0.4 and self.strict_mode:
            return FirewallResult(
                passed=False,
                threat_type=ThreatType.PROMPT_INJECTION,
                threat_score=combined,
                details=f"[STRICT] Combined threat signals: {combined:.2f}",
            )

        return FirewallResult(
            passed=True,
            threat_type=ThreatType.NONE,
            threat_score=combined,
            details="All firewall checks passed.",
        )

    # ── Private checks ────────────────────────────────────────────────

    def _check_injection(self, text: str) -> FirewallResult:
        matches = [m.group(0) for p in self._injection_re if (m := p.search(text))]
        if matches:
            score = min(0.5 + len(matches) * 0.2, 1.0)
            return FirewallResult(
                passed=False,
                threat_type=ThreatType.PROMPT_INJECTION,
                threat_score=score,
                details=f"Prompt injection detected. Matched: {matches[:3]}",
                sanitized_prompt=re.sub("|".join(p.pattern for p in self._injection_re),
                                        "[REDACTED]", text, flags=re.IGNORECASE),
            )
        return FirewallResult(passed=True, threat_score=0.0)

    def _check_sensitive_data(self, action: AgentAction) -> FirewallResult:
        scan = action.prompt + " " + " ".join(
            str(v) for v in action.tool_params.values() if isinstance(v, str)
        )
        found = [name for name, pat in self._sensitive_re.items() if pat.search(scan)]
        if found:
            return FirewallResult(
                passed=False,
                threat_type=ThreatType.SENSITIVE_DATA_EXPOSURE,
                threat_score=min(0.4 + len(found) * 0.2, 1.0),
                details=f"Sensitive data in action payload: {found}",
            )
        return FirewallResult(passed=True, threat_score=0.0)

    def _check_tool_authorization(self, action: AgentAction) -> FirewallResult:
        if action.tool_name in self._restricted_tools:
            return FirewallResult(
                passed=False,
                threat_type=ThreatType.UNAUTHORIZED_TOOL,
                threat_score=1.0,
                details=f"Tool '{action.tool_name}' is on the restricted list.",
            )
        return FirewallResult(passed=True, threat_score=0.0)

    def _check_context(self, action: AgentAction) -> FirewallResult:
        text = " ".join(action.context_sources) + " " + action.reasoning
        matches = [m.group(0) for p in self._context_re if (m := p.search(text))]
        if matches:
            return FirewallResult(
                passed=False,
                threat_type=ThreatType.MALICIOUS_CONTEXT,
                threat_score=0.85,
                details=f"Malicious context markers: {matches}",
            )
        return FirewallResult(passed=True, threat_score=0.0)
