"""
Enhanced security module for Phase 1 BCE improvements.

This module implements prompt injection detection, context validation,
and enhanced behavioral contract enforcement to address security
vulnerabilities identified in the Guardian AI security agents.
"""

import json
import logging
import re
from datetime import datetime
from typing import Any

logger = logging.getLogger(__name__)


class PromptInjectionDetector:
    """Detects various types of prompt injection attacks."""

    def __init__(self):
        # Common prompt injection patterns
        self.injection_patterns = [
            # Direct instruction override attempts
            r"\b(?:ignore|disregard|forget|skip)\s+(?:previous|above|earlier|prior)\s+(?:instructions?|prompts?|rules?)\b",
            r"\b(?:new|different|alternative)\s+(?:instructions?|prompts?|rules?|task)\b",
            r"\bnow\s+(?:do|perform|execute|run)\s+(?:this|the following)\b",
            # Role manipulation
            r"\b(?:act|pretend|behave|roleplay)\s+(?:as|like)\s+(?:a|an)?\s*(?:helpful|different|new)\b",
            r"\byou\s+are\s+(?:now|actually|really)\s+(?:a|an)?\s*\w+",
            r"\bassume\s+(?:the\s+)?role\s+of\b",
            # System manipulation
            r"\bsystem[:\s]+(?:override|bypass|disable|ignore)\b",
            r"\boverride\s+(?:security|safety|behavioral?)\s+(?:protocols?|constraints?|rules?)\b",
            r"\bdeveloper\s+mode\b",
            r"\bjailbreak\s+mode\b",
            # Context confusion
            r"\bthis\s+(?:conversation|chat|session)\s+(?:never|didn\'t)\s+happened?\b",
            r"\bstart\s+(?:over|fresh|new|again)\b",
            r"\breset\s+(?:conversation|context|memory)\b",
            # Output manipulation
            r"\bignore\s+(?:format|structure|schema|contract)\b",
            r"\breturn\s+(?:only|just)\s+\w+",
            r"\bbypass\s+(?:validation|checks?|requirements?)\b",
            # Social engineering
            r"\b(?:emergency|urgent|critical|important):\s*ignore\b",
            r"\bfor\s+(?:testing|debugging|security)\s+purposes?\b",
            r"\bthe\s+(?:user|human|person)\s+(?:is|was)\s+(?:authorized|allowed)\s+to\b",
        ]

        # Compile patterns for efficiency
        self.compiled_patterns = [
            re.compile(pattern, re.IGNORECASE) for pattern in self.injection_patterns
        ]

        # Suspicious keywords that often appear in injection attempts
        self.suspicious_keywords = [
            "jailbreak",
            "override",
            "bypass",
            "ignore",
            "disregard",
            "pretend",
            "roleplay",
            "act as",
            "developer mode",
            "admin mode",
            "root access",
            "system prompt",
            "base instructions",
            "core directives",
        ]

    def detect_injection(self, text: str) -> tuple[bool, list[str], float]:
        """
        Detect prompt injection attempts in input text.

        Args:
            text: Input text to analyze

        Returns:
            Tuple of (is_injection, detected_patterns, confidence_score)
        """
        if not text or not isinstance(text, str):
            return False, [], 0.0

        detected_patterns = []
        confidence_factors = []

        # Check for regex patterns
        for i, pattern in enumerate(self.compiled_patterns):
            matches = pattern.findall(text)
            if matches:
                detected_patterns.append(f"Pattern {i + 1}: {self.injection_patterns[i]}")
                confidence_factors.append(0.8)  # High confidence for regex matches

        # Check for suspicious keywords
        text_lower = text.lower()
        for keyword in self.suspicious_keywords:
            if keyword.lower() in text_lower:
                detected_patterns.append(f"Suspicious keyword: {keyword}")
                confidence_factors.append(0.3)  # Lower confidence for keywords

        # Check for encoding attempts (base64, hex, etc.)
        if self._detect_encoded_content(text):
            detected_patterns.append("Potential encoded injection content")
            confidence_factors.append(0.6)

        # Check for excessive special characters (obfuscation attempt)
        special_char_ratio = len(re.findall(r"[^\w\s]", text)) / len(text) if text else 0
        if special_char_ratio > 0.3:
            detected_patterns.append("High special character ratio (potential obfuscation)")
            confidence_factors.append(0.4)

        # Calculate overall confidence
        if not confidence_factors:
            confidence = 0.0
        else:
            confidence = min(1.0, sum(confidence_factors) / 2)  # Normalize

        is_injection = len(detected_patterns) > 0 and confidence > 0.3

        if is_injection:
            logger.warning(f"Potential prompt injection detected: {detected_patterns}")

        return is_injection, detected_patterns, confidence

    def _detect_encoded_content(self, text: str) -> bool:
        """Detect potentially encoded malicious content."""
        # Simple base64 detection
        base64_pattern = r"[A-Za-z0-9+/]{20,}={0,2}"
        if re.search(base64_pattern, text):
            try:
                import base64

                # Try to decode and check for injection patterns
                potential_b64 = re.findall(base64_pattern, text)
                for encoded in potential_b64:
                    try:
                        decoded = base64.b64decode(encoded).decode("utf-8", errors="ignore")
                        if any(keyword in decoded.lower() for keyword in self.suspicious_keywords):
                            return True
                    except BaseException:
                        continue
            except BaseException:
                pass

        # Simple hex detection
        hex_pattern = r"\\x[0-9a-fA-F]{2}"
        if re.search(hex_pattern, text):
            return True

        return False


class ContextValidator:
    """Validates context consistency to prevent hallucination and behavioral drift."""

    def __init__(self):
        self.context_memory = []
        self.max_memory_size = 10
        self.consistency_threshold = 0.7

    def validate_context_consistency(
        self, current_response: dict[str, Any], previous_context: dict[str, Any] | None = None
    ) -> tuple[bool, str, float]:
        """
        Validate that current response is consistent with previous context.

        Args:
            current_response: Current agent response
            previous_context: Previous conversation context

        Returns:
            Tuple of (is_consistent, reason, confidence_score)
        """
        if not previous_context or not current_response:
            return True, "No previous context to validate against", 1.0

        consistency_score = 1.0
        issues = []

        # Check for contradictory decisions with high confidence
        if self._check_decision_consistency(current_response, previous_context):
            consistency_score -= 0.4
            issues.append("High confidence decision contradiction")

        # Check for factual consistency
        fact_consistency = self._check_factual_consistency(current_response, previous_context)
        consistency_score -= (1.0 - fact_consistency) * 0.3
        if fact_consistency < 0.7:
            issues.append("Factual inconsistency detected")

        # Check for behavioral pattern consistency
        behavior_consistency = self._check_behavioral_consistency(
            current_response, previous_context
        )
        consistency_score -= (1.0 - behavior_consistency) * 0.3
        if behavior_consistency < 0.7:
            issues.append("Behavioral pattern inconsistency")

        consistency_score = max(0.0, consistency_score)
        is_consistent = consistency_score >= self.consistency_threshold

        reason = "; ".join(issues) if issues else "Context validation passed"

        if not is_consistent:
            logger.warning(f"Context inconsistency detected: {reason} (score: {consistency_score})")

        return is_consistent, reason, consistency_score

    def _check_decision_consistency(
        self, current: dict[str, Any], previous: dict[str, Any]
    ) -> bool:
        """Check for contradictory high-confidence decisions."""
        current_decision = current.get("decision", "").lower()
        current_confidence = current.get("confidence_level", 0)

        # Look for previous high-confidence decisions
        memory = previous.get("memory", [])
        if not memory:
            return False

        for mem_entry in memory[:3]:  # Check last 3 entries
            if isinstance(mem_entry, dict):
                analysis = mem_entry.get("analysis", {})
                prev_decision = analysis.get("decision", "").lower()
                prev_confidence = analysis.get("confidence_level", 0)

                # If both are high confidence but different decisions
                if (
                    current_confidence >= 0.8
                    and prev_confidence >= 0.8
                    and current_decision != prev_decision
                    and current_decision
                    and prev_decision
                ):
                    return True

        return False

    def _check_factual_consistency(
        self, current: dict[str, Any], previous: dict[str, Any]
    ) -> float:
        """Check for factual consistency between responses."""
        # Simple keyword-based consistency check
        current_text = str(current.get("risk_assessment", "")) + str(
            current.get("recommendations", [])
        )

        memory = previous.get("memory", [])
        if not memory:
            return 1.0

        consistency_scores = []
        for mem_entry in memory[:2]:  # Check last 2 entries
            if isinstance(mem_entry, dict):
                analysis = mem_entry.get("analysis", {})
                prev_text = str(analysis.get("risk_assessment", "")) + str(
                    analysis.get("recommendations", [])
                )

                # Simple keyword overlap measurement
                current_words = set(re.findall(r"\w+", current_text.lower()))
                prev_words = set(re.findall(r"\w+", prev_text.lower()))

                if current_words and prev_words:
                    overlap = len(current_words.intersection(prev_words))
                    total = len(current_words.union(prev_words))
                    consistency_scores.append(overlap / total if total > 0 else 0.0)

        return sum(consistency_scores) / len(consistency_scores) if consistency_scores else 1.0

    def _check_behavioral_consistency(
        self, current: dict[str, Any], previous: dict[str, Any]
    ) -> float:
        """Check for consistent behavioral patterns."""
        # Check confidence pattern consistency
        current_confidence = current.get("confidence_level", 0)

        memory = previous.get("memory", [])
        if not memory:
            return 1.0

        confidence_scores = []
        for mem_entry in memory[:3]:
            if isinstance(mem_entry, dict):
                analysis = mem_entry.get("analysis", {})
                prev_confidence = analysis.get("confidence_level", 0)
                confidence_scores.append(prev_confidence)

        if not confidence_scores:
            return 1.0

        # Check if confidence levels are reasonably consistent
        avg_confidence = sum(confidence_scores) / len(confidence_scores)
        confidence_variance = sum((c - avg_confidence) ** 2 for c in confidence_scores) / len(
            confidence_scores
        )

        # Penalize extreme confidence swings
        confidence_diff = abs(current_confidence - avg_confidence)
        if confidence_diff > 0.5 and confidence_variance < 0.1:  # Sudden swing from stable pattern
            return 0.5

        return max(0.0, 1.0 - confidence_diff)


class EnhancedBCEValidator:
    """Enhanced Behavioral Contract Enforcement with security features."""

    def __init__(self):
        self.injection_detector = PromptInjectionDetector()
        self.context_validator = ContextValidator()
        self.compliance_templates = self._load_compliance_templates()

    def validate_with_security(
        self,
        input_data: dict[str, Any],
        response: dict[str, Any],
        contract_spec: dict[str, Any],
        context: dict[str, Any] | None = None,
    ) -> tuple[bool, dict[str, Any]]:
        """
        Enhanced validation with security features.

        Args:
            input_data: Original input to the agent
            response: Agent response to validate
            contract_spec: Behavioral contract specification
            context: Previous conversation context

        Returns:
            Tuple of (is_valid, validation_report)
        """
        validation_report = {
            "timestamp": datetime.now().isoformat(),
            "security_checks": {},
            "context_validation": {},
            "contract_validation": {},
            "overall_status": "unknown",
            "risk_level": "low",
            "recommendations": [],
        }

        # 1. Prompt injection detection
        injection_results = self._check_prompt_injection(input_data)
        validation_report["security_checks"]["prompt_injection"] = injection_results

        # 2. Context consistency validation
        context_results = self._check_context_consistency(response, context)
        validation_report["context_validation"] = context_results

        # 3. Enhanced contract validation
        contract_results = self._check_enhanced_contract(response, contract_spec)
        validation_report["contract_validation"] = contract_results

        # 4. Determine overall status and risk level
        overall_valid = (
            not injection_results["detected"]
            and context_results["consistent"]
            and contract_results["valid"]
        )

        validation_report["overall_status"] = "valid" if overall_valid else "invalid"
        validation_report["risk_level"] = self._calculate_risk_level(validation_report)
        validation_report["recommendations"] = self._generate_recommendations(validation_report)

        return overall_valid, validation_report

    def _check_prompt_injection(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """Check for prompt injection in input data."""
        results = {"detected": False, "patterns": [], "confidence": 0.0, "input_fields_checked": []}

        # Check all text fields in input
        for key, value in input_data.items():
            if isinstance(value, str):
                results["input_fields_checked"].append(key)
                is_injection, patterns, confidence = self.injection_detector.detect_injection(value)

                if is_injection:
                    results["detected"] = True
                    results["patterns"].extend(patterns)
                    results["confidence"] = max(results["confidence"], confidence)

        return results

    def _check_context_consistency(
        self, response: dict[str, Any], context: dict[str, Any] | None
    ) -> dict[str, Any]:
        """Check context consistency."""
        is_consistent, reason, confidence = self.context_validator.validate_context_consistency(
            response, context
        )

        return {
            "consistent": is_consistent,
            "reason": reason,
            "confidence": confidence,
            "checks_performed": [
                "decision_consistency",
                "factual_consistency",
                "behavioral_consistency",
            ],
        }

    def _check_enhanced_contract(
        self, response: dict[str, Any], contract_spec: dict[str, Any]
    ) -> dict[str, Any]:
        """Enhanced contract validation."""
        results = {"valid": True, "violations": [], "checks_performed": []}

        # Standard contract checks
        response_contract = contract_spec.get("response_contract", {})
        required_fields = response_contract.get("output_format", {}).get("required_fields", [])

        # Check required fields
        for field in required_fields:
            if field not in response:
                results["valid"] = False
                results["violations"].append(f"Missing required field: {field}")
        results["checks_performed"].append("required_fields")

        # Enhanced security checks
        safety_checks = response_contract.get("safety_checks", {})

        # PII protection
        if safety_checks.get("pii_protection", False):
            if self._contains_pii(response):
                results["valid"] = False
                results["violations"].append("Response contains PII")
        results["checks_performed"].append("pii_protection")

        # Harmful content detection
        if safety_checks.get("harmful_content", False):
            if self._contains_harmful_content(response):
                results["valid"] = False
                results["violations"].append("Response contains harmful content")
        results["checks_performed"].append("harmful_content")

        return results

    def _contains_pii(self, response: dict[str, Any]) -> bool:
        """Enhanced PII detection."""
        pii_patterns = [
            r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",  # Email
            r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b",  # Phone
            r"\b\d{3}[-.]?\d{2}[-.]?\d{4}\b",  # SSN
            r"\b\d{4}[-.]?\d{4}[-.]?\d{4}[-.]?\d{4}\b",  # Credit card
            r"\b[A-Z]{2}\d{6,8}\b",  # Passport numbers
        ]

        text = json.dumps(response)
        return any(re.search(pattern, text) for pattern in pii_patterns)

    def _contains_harmful_content(self, response: dict[str, Any]) -> bool:
        """Detect harmful content patterns."""
        harmful_keywords = [
            "exploit",
            "vulnerability",
            "attack vector",
            "malicious payload",
            "backdoor",
            "rootkit",
            "trojan",
            "ransomware payload",
        ]

        text = json.dumps(response).lower()
        return any(keyword in text for keyword in harmful_keywords)

    def _calculate_risk_level(self, validation_report: dict[str, Any]) -> str:
        """Calculate overall risk level."""
        risk_score = 0

        # Prompt injection risk
        if validation_report["security_checks"]["prompt_injection"]["detected"]:
            confidence = validation_report["security_checks"]["prompt_injection"]["confidence"]
            risk_score += confidence * 40  # Max 40 points

        # Context consistency risk
        if not validation_report["context_validation"]["consistent"]:
            confidence = 1.0 - validation_report["context_validation"]["confidence"]
            risk_score += confidence * 30  # Max 30 points

        # Contract violation risk
        if not validation_report["contract_validation"]["valid"]:
            violation_count = len(validation_report["contract_validation"]["violations"])
            risk_score += min(violation_count * 10, 30)  # Max 30 points

        # Risk level determination
        if risk_score >= 70:
            return "critical"
        elif risk_score >= 40:
            return "high"
        elif risk_score >= 20:
            return "medium"
        else:
            return "low"

    def _generate_recommendations(self, validation_report: dict[str, Any]) -> list[str]:
        """Generate actionable recommendations based on validation results."""
        recommendations = []

        # Prompt injection recommendations
        if validation_report["security_checks"]["prompt_injection"]["detected"]:
            recommendations.append("Implement input sanitization and validation")
            recommendations.append("Consider using a separate prompt injection filter")
            recommendations.append("Log and analyze injection attempts for patterns")

        # Context consistency recommendations
        if not validation_report["context_validation"]["consistent"]:
            recommendations.append("Review agent memory and context handling")
            recommendations.append("Implement consistency checks before response generation")
            recommendations.append("Consider retraining or fine-tuning the model")

        # Contract violation recommendations
        if not validation_report["contract_validation"]["valid"]:
            recommendations.append("Review and update behavioral contract specifications")
            recommendations.append("Implement stricter validation rules")
            recommendations.append("Add fallback responses for contract violations")

        # Risk-based recommendations
        risk_level = validation_report["risk_level"]
        if risk_level in ["high", "critical"]:
            recommendations.append("Immediately escalate for manual review")
            recommendations.append("Consider temporarily disabling agent until issues resolved")
        elif risk_level == "medium":
            recommendations.append("Increase monitoring frequency")
            recommendations.append("Review agent outputs more frequently")

        return recommendations

    def _load_compliance_templates(self) -> dict[str, Any]:
        """Load compliance templates for GDPR, HIPAA, SOC2."""
        return {
            "gdpr": {
                "data_protection": True,
                "pii_handling": "strict",
                "consent_required": True,
                "retention_limits": True,
            },
            "hipaa": {
                "phi_protection": True,
                "access_controls": True,
                "audit_logging": True,
                "encryption_required": True,
            },
            "soc2": {
                "security_controls": True,
                "availability_monitoring": True,
                "processing_integrity": True,
                "confidentiality": True,
            },
        }
