"""Memory Immune System — quarantines suspicious memories.

Protects the knowledge base from:
- Prompt injections masquerading as memories
- Hallucinated facts with low reliability
- Rapid-fire writes (potential spam/injection attack)
- Excessively long or malformed content
"""

from __future__ import annotations

import re
import time
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timezone

from mnemosynth.core.types import MemoryNode, MemoryStatus


@dataclass
class ThreatReport:
    """Report from immune system analysis."""
    clean: bool = True
    threats: list[str] = field(default_factory=list)
    risk_score: float = 0.0  # 0.0 = clean, 1.0 = definitely malicious
    quarantined: bool = False


# Prompt injection patterns
INJECTION_PATTERNS = [
    r"ignore\s+(previous|all|above)\s+(instructions?|prompts?|rules?)",
    r"you\s+are\s+now\s+[A-Z]",
    r"system\s*:\s*",
    r"<\|?system\|?>",
    r"act\s+as\s+(if|though)",
    r"pretend\s+(you|to|that)",
    r"override\s+(your|the|all)",
    r"disregard\s+(all|previous|your)",
    r"forget\s+(all|everything|your)\s+(rules|instructions|guidelines)",
    r"new\s+instructions?\s*:",
    r"jailbreak",
    r"DAN\s+mode",
]

# Compiled patterns for efficiency
_COMPILED_PATTERNS = [re.compile(p, re.IGNORECASE) for p in INJECTION_PATTERNS]


class ImmuneSystem:
    """Memory immune system — screens, quarantines, and rate-limits.

    Think of it as the blood-brain barrier for AI memory.
    """

    def __init__(self, config=None):
        self.config = config
        self.max_content_length = 5000
        self.min_content_length = 3
        self.rate_limit_per_minute = 60
        self.quarantine_threshold = 0.6

        # Rate limiting state
        self._write_timestamps: dict[str, list[float]] = defaultdict(list)

    def screen(self, memory: MemoryNode, session_id: str = "default") -> ThreatReport:
        """Screen a memory before it enters the knowledge base.

        Returns a ThreatReport with risk assessment.
        """
        report = ThreatReport()

        # Check 1: Content length
        if len(memory.content) > self.max_content_length:
            report.threats.append(
                f"Content too long ({len(memory.content)} chars, max {self.max_content_length})"
            )
            report.risk_score += 0.3

        if len(memory.content.strip()) < self.min_content_length:
            report.threats.append(
                f"Content too short ({len(memory.content.strip())} chars)"
            )
            report.risk_score += 0.2

        # Check 2: Prompt injection patterns
        injection_score = self._check_injection(memory.content)
        if injection_score > 0:
            report.threats.append(
                f"Prompt injection pattern detected (score: {injection_score:.2f})"
            )
            report.risk_score += injection_score

        # Check 3: Rate limiting
        if self._is_rate_limited(session_id):
            report.threats.append(
                f"Rate limit exceeded ({self.rate_limit_per_minute} writes/min)"
            )
            report.risk_score += 0.4

        # Check 4: Suspicious characters
        suspicious = self._check_suspicious_chars(memory.content)
        if suspicious:
            report.threats.append(f"Suspicious characters: {suspicious}")
            report.risk_score += 0.15

        # Check 5: Low source reliability
        if memory.source.reliability < 0.3:
            report.threats.append(
                f"Low source reliability ({memory.source.reliability:.2f})"
            )
            report.risk_score += 0.2

        # Final assessment
        report.risk_score = min(1.0, report.risk_score)
        report.clean = report.risk_score < self.quarantine_threshold

        if not report.clean:
            report.quarantined = True
            memory.quarantine()

        # Record write timestamp for rate limiting
        self._record_write(session_id)

        return report

    def _check_injection(self, text: str) -> float:
        """Check for prompt injection patterns. Returns 0.0 to 0.7."""
        matches = 0
        for pattern in _COMPILED_PATTERNS:
            if pattern.search(text):
                matches += 1

        if matches == 0:
            return 0.0
        return min(0.7, 0.3 * matches)

    def _is_rate_limited(self, session_id: str) -> bool:
        """Check if the session is being rate limited."""
        now = time.time()
        # Clean old timestamps (older than 60 seconds)
        self._write_timestamps[session_id] = [
            ts for ts in self._write_timestamps[session_id]
            if now - ts < 60
        ]

        return len(self._write_timestamps[session_id]) >= self.rate_limit_per_minute

    def _record_write(self, session_id: str) -> None:
        """Record a write timestamp for rate limiting."""
        self._write_timestamps[session_id].append(time.time())

    def _check_suspicious_chars(self, text: str) -> str:
        """Check for suspicious Unicode or control characters."""
        suspicious = []

        # Zero-width characters (hidden text)
        if re.search(r'[\u200b\u200c\u200d\ufeff\u2060]', text):
            suspicious.append("zero-width chars")

        # Excessive newlines (formatting attack)
        if text.count('\n') > 50:
            suspicious.append("excessive newlines")

        # Excessive special characters
        special_ratio = sum(1 for c in text if not c.isalnum() and not c.isspace()) / max(len(text), 1)
        if special_ratio > 0.4:
            suspicious.append("high special-char ratio")

        return ", ".join(suspicious)

    def get_quarantined_count(self) -> int:
        """Get the number of messages quarantined in this session."""
        return sum(
            1 for timestamps in self._write_timestamps.values()
            for _ in timestamps
        )
