# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Spam Filter for Familiar Email Server.

Provides basic spam detection using:
- Keyword/pattern matching
- Header analysis
- Sender reputation (basic)
- SPF/DKIM verification status
"""

import logging
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class SpamRule:
    """A spam detection rule."""

    name: str
    pattern: str
    score: float
    description: str
    field: str = "body"  # body, subject, sender, headers

    def __post_init__(self):
        self._regex = re.compile(self.pattern, re.IGNORECASE)

    def match(self, text: str) -> bool:
        return bool(self._regex.search(text))


# Default spam rules
DEFAULT_RULES = [
    # High-confidence spam indicators
    SpamRule("nigerian_prince", r"nigerian?\s+prince", 5.0, "Nigerian prince scam", "body"),
    SpamRule(
        "lottery_winner",
        r"you\s+(have\s+)?won\s+.*?(lottery|million|prize)",
        4.0,
        "Lottery scam",
        "body",
    ),
    SpamRule("urgent_action", r"urgent\s+action\s+required", 3.0, "Urgent action spam", "subject"),
    SpamRule("wire_transfer", r"wire\s+transfer", 2.5, "Wire transfer request", "body"),
    SpamRule(
        "click_here", r"click\s+here\s+(now|immediately|urgently)", 2.0, "Urgent click spam", "body"
    ),
    # Pharmaceutical spam
    SpamRule("viagra", r"\bviagra\b", 3.0, "Pharmaceutical spam", "body"),
    SpamRule("cialis", r"\bcialis\b", 3.0, "Pharmaceutical spam", "body"),
    SpamRule(
        "pharmacy_spam", r"(online|canadian|discount)\s+pharmacy", 2.5, "Pharmacy spam", "body"
    ),
    # Financial spam
    SpamRule(
        "crypto_scam",
        r"(bitcoin|crypto)\s+(opportunity|investment|profit)",
        2.5,
        "Crypto scam",
        "body",
    ),
    SpamRule("get_rich", r"get\s+rich\s+(quick|fast)", 3.0, "Get rich quick scam", "body"),
    SpamRule(
        "make_money",
        r"make\s+\$?\d+[,\d]*\s*(per|a)\s+(day|week|month)",
        2.0,
        "Money making spam",
        "body",
    ),
    # Suspicious patterns
    SpamRule("excessive_caps", r"[A-Z\s]{20,}", 1.5, "Excessive capitals", "subject"),
    SpamRule("excessive_exclaim", r"!{3,}", 1.0, "Excessive exclamation marks", "subject"),
    SpamRule("free_offer", r"free\s+(gift|offer|trial|money)", 1.5, "Free offer spam", "body"),
    SpamRule("act_now", r"act\s+now", 1.0, "Urgency spam", "body"),
    SpamRule("limited_time", r"limited\s+time\s+(offer|only)", 1.0, "Limited time spam", "body"),
    # Phishing indicators
    SpamRule(
        "verify_account",
        r"verify\s+your\s+(account|identity|information)",
        2.5,
        "Phishing attempt",
        "body",
    ),
    SpamRule(
        "suspended_account",
        r"(account|access)\s+(has\s+been\s+)?(suspended|locked|disabled)",
        2.5,
        "Phishing attempt",
        "body",
    ),
    SpamRule(
        "password_reset",
        r"password\s+(reset|expired|will\s+expire)",
        2.0,
        "Potential phishing",
        "body",
    ),
    SpamRule("confirm_identity", r"confirm\s+your\s+identity", 2.0, "Phishing attempt", "body"),
    # Header-based rules
    SpamRule("no_subject", r"^$", 1.0, "Empty subject", "subject"),
    SpamRule("re_no_thread", r"^re:\s*$", 1.5, "Fake reply", "subject"),
    # Sender patterns
    SpamRule("random_sender", r"^[a-z0-9]{15,}@", 1.5, "Random sender address", "sender"),
    SpamRule("numeric_sender", r"^\d+@", 1.0, "Numeric sender", "sender"),
]


class SpamFilter:
    """
    Spam detection for incoming email.

    Uses a scoring system where each rule adds to the total score.
    Messages above the threshold are marked as spam.
    """

    def __init__(
        self,
        threshold: float = 5.0,
        rules: Optional[List[SpamRule]] = None,
    ):
        self.threshold = threshold
        self.rules = rules or DEFAULT_RULES

        # Whitelist/blacklist
        self.whitelist: set = set()
        self.blacklist: set = set()

        # Stats
        self._total_checked = 0
        self._total_spam = 0

    async def check(
        self,
        sender: str,
        subject: str,
        body: str,
        headers: Optional[Dict[str, str]] = None,
    ) -> float:
        """
        Check a message for spam.

        Args:
            sender: Sender email address
            subject: Email subject
            body: Email body (plain text)
            headers: Optional email headers

        Returns:
            Spam score (higher = more likely spam)
        """
        self._total_checked += 1
        score = 0.0

        # Check whitelist (exact address or domain match)
        sender_lower = sender.lower()
        sender_domain = sender_lower.rsplit("@", 1)[-1] if "@" in sender_lower else ""
        if any(sender_lower == w or sender_domain == w for w in self.whitelist):
            return 0.0

        # Check blacklist (exact address or domain match)
        if any(sender_lower == b or sender_domain == b for b in self.blacklist):
            score += 10.0

        # Apply rules
        for rule in self.rules:
            if rule.field == "body":
                text = body
            elif rule.field == "subject":
                text = subject
            elif rule.field == "sender":
                text = sender
            elif rule.field == "headers":
                text = str(headers) if headers else ""
            else:
                continue

            if rule.match(text):
                score += rule.score
                logger.debug(f"Spam rule matched: {rule.name} (+{rule.score})")

        # Header analysis
        if headers:
            score += self._analyze_headers(headers)

        # Track stats
        if score >= self.threshold:
            self._total_spam += 1

        return score

    def _analyze_headers(self, headers: Dict[str, str]) -> float:
        """Analyze headers for spam indicators."""
        score = 0.0

        # Check for missing headers
        required_headers = ["From", "Date", "Message-ID"]
        for h in required_headers:
            if h not in headers and h.lower() not in headers:
                score += 0.5

        # Check Received headers count (too few is suspicious)
        received_count = sum(1 for k in headers if k.lower() == "received")
        if received_count == 0:
            score += 1.0

        # Check for suspicious X-headers
        suspicious_x = ["X-Spam", "X-Virus", "X-Mailer: PHPMailer"]
        for h, v in headers.items():
            for sus in suspicious_x:
                if sus.lower() in f"{h}: {v}".lower():
                    score += 0.5

        # Check Content-Type for suspicious types
        content_type = headers.get("Content-Type", headers.get("content-type", ""))
        if "multipart/mixed" in content_type.lower():
            # Attachments - slightly higher scrutiny
            score += 0.3

        return score

    def add_whitelist(self, pattern: str):
        """Add a pattern to the whitelist."""
        self.whitelist.add(pattern.lower())

    def remove_whitelist(self, pattern: str):
        """Remove a pattern from the whitelist."""
        self.whitelist.discard(pattern.lower())

    def add_blacklist(self, pattern: str):
        """Add a pattern to the blacklist."""
        self.blacklist.add(pattern.lower())

    def remove_blacklist(self, pattern: str):
        """Remove a pattern from the blacklist."""
        self.blacklist.discard(pattern.lower())

    def add_rule(self, rule: SpamRule):
        """Add a custom spam rule."""
        self.rules.append(rule)

    def get_stats(self) -> Dict[str, Any]:
        """Get spam filter statistics."""
        return {
            "total_checked": self._total_checked,
            "total_spam": self._total_spam,
            "spam_rate": (
                self._total_spam / self._total_checked * 100 if self._total_checked > 0 else 0
            ),
            "threshold": self.threshold,
            "rule_count": len(self.rules),
            "whitelist_count": len(self.whitelist),
            "blacklist_count": len(self.blacklist),
        }

    def reset_stats(self):
        """Reset statistics."""
        self._total_checked = 0
        self._total_spam = 0


class BayesianSpamFilter:
    """
    Bayesian spam filter that learns from user feedback.

    Uses token frequency analysis to classify messages
    as spam or ham (not spam).
    """

    def __init__(self):
        self.spam_tokens: Dict[str, int] = {}
        self.ham_tokens: Dict[str, int] = {}
        self.spam_count = 0
        self.ham_count = 0

    def train_spam(self, text: str):
        """Train on a spam message."""
        tokens = self._tokenize(text)
        for token in tokens:
            self.spam_tokens[token] = self.spam_tokens.get(token, 0) + 1
        self.spam_count += 1

    def train_ham(self, text: str):
        """Train on a ham (not spam) message."""
        tokens = self._tokenize(text)
        for token in tokens:
            self.ham_tokens[token] = self.ham_tokens.get(token, 0) + 1
        self.ham_count += 1

    def classify(self, text: str) -> float:
        """
        Classify a message.

        Returns probability of being spam (0.0 - 1.0).
        """
        if self.spam_count == 0 or self.ham_count == 0:
            return 0.5  # Not enough training data

        import math

        tokens = self._tokenize(text)

        log_spam = 0.0
        log_ham = 0.0

        for token in tokens:
            # Get token probabilities (Laplace smoothing)
            p_token_spam = (self.spam_tokens.get(token, 0) + 1) / (self.spam_count + 2)
            p_token_ham = (self.ham_tokens.get(token, 0) + 1) / (self.ham_count + 2)

            log_spam += math.log(p_token_spam)
            log_ham += math.log(p_token_ham)

        # Log-sum-exp normalization to avoid underflow
        log_max = max(log_spam, log_ham)
        spam_exp = math.exp(log_spam - log_max)
        ham_exp = math.exp(log_ham - log_max)
        total = spam_exp + ham_exp

        if total == 0:
            return 0.5

        return spam_exp / total

    def _tokenize(self, text: str) -> List[str]:
        """Tokenize text for analysis."""
        # Simple word tokenization
        text = text.lower()
        tokens = re.findall(r"\b[a-z]{3,}\b", text)
        return tokens

    def save(self, path: str):
        """Save trained model."""
        import json

        data = {
            "spam_tokens": self.spam_tokens,
            "ham_tokens": self.ham_tokens,
            "spam_count": self.spam_count,
            "ham_count": self.ham_count,
        }
        with open(path, "w") as f:
            json.dump(data, f)

    def load(self, path: str):
        """Load trained model."""
        import json

        with open(path, "r") as f:
            data = json.load(f)
        self.spam_tokens = data["spam_tokens"]
        self.ham_tokens = data["ham_tokens"]
        self.spam_count = data["spam_count"]
        self.ham_count = data["ham_count"]
