# 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

"""
Filter Engine for Private Browser.

Blocks ads, trackers, and malicious content server-side
before it reaches the user.
"""

import logging
import re
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Set
from urllib.parse import urljoin, urlparse

logger = logging.getLogger(__name__)


class FilterLevel(Enum):
    """Filtering level."""

    STANDARD = "standard"  # Ads + common trackers
    STRICT = "strict"  # + all scripts from third parties
    PARANOID = "paranoid"  # + images from third parties, text only


# Built-in blocklists
AD_DOMAINS = {
    # Ad networks
    "doubleclick.net",
    "googlesyndication.com",
    "googleadservices.com",
    "googleads.g.doubleclick.net",
    "adservice.google.com",
    "pagead2.googlesyndication.com",
    "adsserver.com",
    "adnxs.com",
    "adsrvr.org",
    "advertising.com",
    "adobedtm.com",
    "adroll.com",
    "criteo.com",
    "criteo.net",
    "taboola.com",
    "outbrain.com",
    "revcontent.com",
    "mgid.com",
    "zergnet.com",
    "amazon-adsystem.com",
    "media.net",
    "pubmatic.com",
    "rubiconproject.com",
    "openx.net",
    "contextweb.com",
    "bidswitch.net",
    "casalemedia.com",
    "lijit.com",
    "sharethrough.com",
    "tribalfusion.com",
    "yieldmanager.com",
}

TRACKER_DOMAINS = {
    # Analytics & Tracking
    "google-analytics.com",
    "googletagmanager.com",
    "googletagservices.com",
    "facebook.com",
    "facebook.net",
    "fbcdn.net",
    "connect.facebook.net",
    "pixel.facebook.com",
    "analytics.twitter.com",
    "platform.twitter.com",
    "ads-twitter.com",
    "analytics.tiktok.com",
    "hotjar.com",
    "mixpanel.com",
    "segment.com",
    "segment.io",
    "amplitude.com",
    "heap.io",
    "heapanalytics.com",
    "fullstory.com",
    "mouseflow.com",
    "crazyegg.com",
    "luckyorange.com",
    "clicktale.net",
    "optimizely.com",
    "omtrdc.net",
    "demdex.net",
    "omniture.com",
    "scorecardresearch.com",
    "quantserve.com",
    "newrelic.com",
    "nr-data.net",
    "sentry.io",
    "bugsnag.com",
    "rollbar.com",
    "loggly.com",
    "sumologic.com",
    "datadoghq.com",
    # Fingerprinting
    "fingerprintjs.com",
    "fpjs.io",
    # Social widgets (tracking)
    "addthis.com",
    "sharethis.com",
    "addtoany.com",
    "po.st",
    # Marketing
    "hubspot.com",
    "hs-analytics.net",
    "marketo.com",
    "mktoresp.com",
    "pardot.com",
    "eloqua.com",
    "salesforce.com",
    "drift.com",
    "intercom.io",
    "intercomcdn.com",
    "zendesk.com",
    "zdassets.com",
}

MALWARE_DOMAINS = {
    # Known malware/phishing domains would go here
    # This would be updated from threat intelligence feeds
}

# Patterns to remove from HTML
REMOVAL_PATTERNS = [
    # Inline tracking pixels
    r"<img[^>]+(?:1x1|spacer|pixel|tracker|beacon)[^>]*>",
    r'<img[^>]+width=["\']?1["\']?[^>]+height=["\']?1["\']?[^>]*>',
    # Tracking iframes
    r"<iframe[^>]+(?:hidden|display:\s*none|visibility:\s*hidden)[^>]*>.*?</iframe>",
    # Inline scripts with tracking
    r"<script[^>]*>[^<]*(?:ga\(|gtag\(|fbq\(|_gaq|analytics|tracking)[^<]*</script>",
    # Tracking noscript tags
    r"<noscript[^>]*>[^<]*<img[^>]*>[^<]*</noscript>",
]


@dataclass
class FilterStats:
    """Statistics about filtered content."""

    ads_blocked: int = 0
    trackers_blocked: int = 0
    scripts_blocked: int = 0
    images_blocked: int = 0
    domains_blocked: Set[str] = None

    def __post_init__(self):
        if self.domains_blocked is None:
            self.domains_blocked = set()


class FilterEngine:
    """
    Content filtering engine.

    Processes HTML content to remove:
    - Advertisements
    - Tracking scripts and pixels
    - Third-party scripts (in strict mode)
    - Third-party images (in paranoid mode)
    """

    def __init__(self, config):
        self.config = config

        # Blocklists
        self.ad_domains: Set[str] = set(AD_DOMAINS)
        self.tracker_domains: Set[str] = set(TRACKER_DOMAINS)
        self.malware_domains: Set[str] = set(MALWARE_DOMAINS)

        # Custom allow/block lists
        self.custom_allow: Set[str] = set()
        self.custom_block: Set[str] = set()

        # Stats
        self._stats = FilterStats()
        self._lifetime_stats = FilterStats()

    async def load_blocklists(self):
        """Load additional blocklists from files/URLs."""
        # Could load EasyList, EasyPrivacy, etc.
        # For now, use built-in lists
        logger.info(
            f"Filter engine loaded: {len(self.ad_domains)} ad domains, "
            f"{len(self.tracker_domains)} tracker domains"
        )

    async def process(
        self,
        content: str,
        url: str,
        content_type: str,
    ) -> Dict[str, Any]:
        """
        Process and filter content.

        Args:
            content: HTML content
            url: Page URL (for resolving relative URLs)
            content_type: Content-Type header

        Returns:
            {
                'html': filtered HTML,
                'blocked': counts of blocked items,
            }
        """
        # Only process HTML
        if "text/html" not in content_type.lower():
            return {
                "html": content,
                "blocked": {},
            }

        # Reset per-page stats
        self._stats = FilterStats()

        # Get page domain for same-origin checks
        page_domain = urlparse(url).netloc

        # Apply filters
        html = content

        # 1. Remove tracking patterns
        html = self._remove_patterns(html)

        # 2. Filter external resources
        html = self._filter_scripts(html, page_domain)
        html = self._filter_images(html, page_domain)
        html = self._filter_iframes(html)
        html = self._filter_links(html)

        # 3. Remove empty containers
        html = self._cleanup_html(html)

        # 4. Rewrite remaining links to go through proxy
        html = self._rewrite_links(html, url)

        # Update lifetime stats
        self._lifetime_stats.ads_blocked += self._stats.ads_blocked
        self._lifetime_stats.trackers_blocked += self._stats.trackers_blocked
        self._lifetime_stats.scripts_blocked += self._stats.scripts_blocked
        self._lifetime_stats.images_blocked += self._stats.images_blocked
        self._lifetime_stats.domains_blocked.update(self._stats.domains_blocked)

        return {
            "html": html,
            "blocked": {
                "ads": self._stats.ads_blocked,
                "trackers": self._stats.trackers_blocked,
                "scripts": self._stats.scripts_blocked,
                "images": self._stats.images_blocked,
            },
        }

    def _remove_patterns(self, html: str) -> str:
        """Remove known tracking patterns."""
        for pattern in REMOVAL_PATTERNS:
            matches = re.findall(pattern, html, re.IGNORECASE | re.DOTALL)
            if matches:
                self._stats.trackers_blocked += len(matches)
            html = re.sub(pattern, "", html, flags=re.IGNORECASE | re.DOTALL)

        return html

    def _filter_scripts(self, html: str, page_domain: str) -> str:
        """Filter script tags."""

        def filter_script(match):
            tag = match.group(0)
            src_match = re.search(r'src=["\']([^"\']+)["\']', tag)

            if not src_match:
                # Inline script - check content
                if self.config.block_scripts:
                    self._stats.scripts_blocked += 1
                    return "<!-- script blocked -->"
                return tag

            src = src_match.group(1)
            domain = self._extract_domain(src)

            # Check if should block
            if self._should_block(domain, page_domain, "script"):
                self._stats.scripts_blocked += 1
                self._stats.domains_blocked.add(domain)
                return f"<!-- script blocked: {domain} -->"

            return tag

        return re.sub(
            r"<script[^>]*>.*?</script>", filter_script, html, flags=re.IGNORECASE | re.DOTALL
        )

    def _filter_images(self, html: str, page_domain: str) -> str:
        """Filter image tags."""

        def filter_image(match):
            tag = match.group(0)
            src_match = re.search(r'src=["\']([^"\']+)["\']', tag)

            if not src_match:
                return tag

            src = src_match.group(1)

            # Skip data URLs
            if src.startswith("data:"):
                return tag

            domain = self._extract_domain(src)

            # Check if should block
            if self._should_block(domain, page_domain, "image"):
                self._stats.images_blocked += 1
                self._stats.domains_blocked.add(domain)
                # Replace with placeholder
                alt = re.search(r'alt=["\']([^"\']*)["\']', tag)
                alt_text = alt.group(1) if alt else "Image"
                return f'<span class="blocked-image">[{alt_text}]</span>'

            return tag

        return re.sub(r"<img[^>]+>", filter_image, html, flags=re.IGNORECASE)

    def _filter_iframes(self, html: str) -> str:
        """Filter iframe tags."""

        def filter_iframe(match):
            tag = match.group(0)
            src_match = re.search(r'src=["\']([^"\']+)["\']', tag)

            if not src_match:
                return tag

            src = src_match.group(1)
            domain = self._extract_domain(src)

            # Block all iframes from ad/tracker domains
            if domain in self.ad_domains or domain in self.tracker_domains:
                self._stats.ads_blocked += 1
                self._stats.domains_blocked.add(domain)
                return f"<!-- iframe blocked: {domain} -->"

            return tag

        return re.sub(
            r"<iframe[^>]*>.*?</iframe>", filter_iframe, html, flags=re.IGNORECASE | re.DOTALL
        )

    def _filter_links(self, html: str) -> str:
        """Filter link tags (CSS, preload, etc.)."""

        def filter_link(match):
            tag = match.group(0)
            href_match = re.search(r'href=["\']([^"\']+)["\']', tag)

            if not href_match:
                return tag

            href = href_match.group(1)
            domain = self._extract_domain(href)

            # Block tracking pixels and beacons in link tags
            if domain in self.tracker_domains:
                self._stats.trackers_blocked += 1
                return ""

            return tag

        return re.sub(r"<link[^>]+>", filter_link, html, flags=re.IGNORECASE)

    def _should_block(
        self,
        domain: str,
        page_domain: str,
        resource_type: str,
    ) -> bool:
        """Determine if a domain should be blocked."""
        if not domain:
            return False

        # Custom allow list takes priority
        if domain in self.custom_allow:
            return False

        # Custom block list
        if domain in self.custom_block:
            return True

        # Malware always blocked
        if domain in self.malware_domains:
            return True

        # Ad domains
        if self.config.block_ads and domain in self.ad_domains:
            self._stats.ads_blocked += 1
            return True

        # Tracker domains
        if self.config.block_trackers and domain in self.tracker_domains:
            self._stats.trackers_blocked += 1
            return True

        # Strict mode: block third-party scripts
        if self.config.block_scripts and resource_type == "script":
            if not self._is_same_site(domain, page_domain):
                return True

        # Paranoid mode: block third-party images
        if self.config.block_images and resource_type == "image":
            if not self._is_same_site(domain, page_domain):
                return True

        return False

    def _is_same_site(self, domain1: str, domain2: str) -> bool:
        """Check if two domains are same-site."""
        if not domain1 or not domain2:
            return True

        # Extract base domain (e.g., example.com from www.example.com)
        def base_domain(d):
            parts = d.lower().split(".")
            if len(parts) >= 2:
                return ".".join(parts[-2:])
            return d

        return base_domain(domain1) == base_domain(domain2)

    def _extract_domain(self, url: str) -> str:
        """Extract domain from URL."""
        if not url or url.startswith("data:"):
            return ""

        # Handle protocol-relative URLs
        if url.startswith("//"):
            url = "https:" + url

        # Handle relative URLs (no domain)
        if not url.startswith(("http://", "https://")):
            return ""

        try:
            return urlparse(url).netloc.lower()
        except Exception:
            return ""

    def _cleanup_html(self, html: str) -> str:
        """Clean up HTML after filtering."""
        # Remove empty divs/spans that might have contained ads
        html = re.sub(r"<div[^>]*>\s*</div>", "", html)
        html = re.sub(r"<span[^>]*>\s*</span>", "", html)

        # Remove consecutive whitespace
        html = re.sub(r"\n\s*\n", "\n\n", html)

        return html

    def _rewrite_links(self, html: str, base_url: str) -> str:
        """Rewrite links to go through proxy."""
        # For now, just make relative URLs absolute
        # Full proxy rewriting would prepend /proxy?url=

        def make_absolute(match):
            attr = match.group(1)
            url = match.group(2)

            if url.startswith(("http://", "https://", "data:", "javascript:", "#", "mailto:")):
                return match.group(0)

            absolute = urljoin(base_url, url)
            return f'{attr}="{absolute}"'

        html = re.sub(r'(href|src)=["\']([^"\']+)["\']', make_absolute, html)

        return html

    def add_to_allowlist(self, domain: str):
        """Add domain to allow list."""
        self.custom_allow.add(domain.lower())

    def add_to_blocklist(self, domain: str):
        """Add domain to block list."""
        self.custom_block.add(domain.lower())

    def get_stats(self) -> Dict[str, Any]:
        """Get filter statistics."""
        return {
            "session": {
                "ads_blocked": self._stats.ads_blocked,
                "trackers_blocked": self._stats.trackers_blocked,
                "scripts_blocked": self._stats.scripts_blocked,
                "images_blocked": self._stats.images_blocked,
            },
            "lifetime": {
                "ads_blocked": self._lifetime_stats.ads_blocked,
                "trackers_blocked": self._lifetime_stats.trackers_blocked,
                "scripts_blocked": self._lifetime_stats.scripts_blocked,
                "images_blocked": self._lifetime_stats.images_blocked,
                "unique_domains_blocked": len(self._lifetime_stats.domains_blocked),
            },
            "blocklist_sizes": {
                "ad_domains": len(self.ad_domains),
                "tracker_domains": len(self.tracker_domains),
                "custom_block": len(self.custom_block),
                "custom_allow": len(self.custom_allow),
            },
        }
