# 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

"""
Encrypted Page Archive for Private Browser.

Stores browsed pages in encrypted format for later access.
Supports full-text search across archived pages.
"""

import gzip
import json
import logging
import re
import uuid
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Optional encryption
try:
    from cryptography.fernet import Fernet

    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False


@dataclass
class ArchivedPage:
    """An archived web page."""

    id: str
    url: str
    title: str
    content_html: str
    content_text: str
    reader_html: Optional[str]
    summary: Optional[str]
    archived_at: str
    size_bytes: int

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ArchivedPage":
        return cls(**data)


class PageArchive:
    """
    Encrypted page archive.

    Stores pages compressed and encrypted for privacy.
    Maintains a searchable index for full-text search.

    Directory structure:
    archive_path/
    ├── pages/
    │   ├── abc123.gz.enc  (compressed, encrypted page)
    │   └── ...
    ├── index.json.enc     (encrypted search index)
    └── .key               (encryption key)
    """

    def __init__(self, archive_path: Path):
        self.archive_path = Path(archive_path)
        self.pages_path = self.archive_path / "pages"

        self._fernet: Optional[Fernet] = None
        self._index: Dict[str, Dict[str, Any]] = {}  # id -> metadata
        self._search_index: Dict[str, List[str]] = {}  # word -> [ids]

    async def initialize(self):
        """Initialize the archive."""
        self.archive_path.mkdir(parents=True, exist_ok=True)
        self.pages_path.mkdir(exist_ok=True)

        # Get or create encryption key
        if HAS_CRYPTO:
            key = await self._get_encryption_key()
            self._fernet = Fernet(key)

        # Load index
        await self._load_index()

        logger.info(f"Page archive initialized at {self.archive_path} ({len(self._index)} pages)")

    async def _get_encryption_key(self) -> bytes:
        """Get or generate encryption key."""
        key_file = self.archive_path / ".key"

        if key_file.exists():
            return key_file.read_bytes()

        key = Fernet.generate_key()
        key_file.write_bytes(key)
        key_file.chmod(0o600)

        return key

    async def _load_index(self):
        """Load the page index."""
        index_file = self.archive_path / "index.json"

        if index_file.exists():
            try:
                data = index_file.read_bytes()

                # Try encrypted first
                if self._fernet:
                    try:
                        data = self._fernet.decrypt(data)
                    except Exception:
                        pass  # Might be unencrypted

                index_data = json.loads(data.decode("utf-8"))
                self._index = index_data.get("pages", {})
                self._search_index = index_data.get("search", {})
            except Exception as e:
                logger.error(f"Error loading index: {e}")
                self._index = {}
                self._search_index = {}

    async def _save_index(self):
        """Save the page index."""
        index_file = self.archive_path / "index.json"

        data = json.dumps(
            {
                "pages": self._index,
                "search": self._search_index,
            }
        ).encode("utf-8")

        if self._fernet:
            data = self._fernet.encrypt(data)

        # Atomic write
        temp_file = index_file.with_suffix(".tmp")
        temp_file.write_bytes(data)
        temp_file.replace(index_file)

    async def save(self, browse_result) -> str:
        """
        Save a page to the archive.

        Args:
            browse_result: BrowseResult from browser

        Returns:
            Archive ID
        """
        # Generate ID
        archive_id = str(uuid.uuid4())[:8]

        # Create archived page
        page = ArchivedPage(
            id=archive_id,
            url=browse_result.url,
            title=browse_result.title,
            content_html=browse_result.content_html,
            content_text=browse_result.content_text,
            reader_html=browse_result.reader_html,
            summary=browse_result.summary,
            archived_at=datetime.utcnow().isoformat(),
            size_bytes=len(browse_result.content_html),
        )

        # Compress
        compressed = gzip.compress(json.dumps(page.to_dict()).encode("utf-8"), compresslevel=9)

        # Encrypt
        if self._fernet:
            compressed = self._fernet.encrypt(compressed)

        # Save to file
        page_file = self.pages_path / f"{archive_id}.gz.enc"
        page_file.write_bytes(compressed)

        # Update index
        self._index[archive_id] = {
            "url": page.url,
            "title": page.title,
            "archived_at": page.archived_at,
            "size_bytes": page.size_bytes,
            "has_summary": page.summary is not None,
        }

        # Update search index
        self._index_text(archive_id, f"{page.title} {page.content_text}")

        await self._save_index()

        logger.info(f"Archived page {archive_id}: {page.title}")

        return archive_id

    async def load(self, archive_id: str) -> Optional[ArchivedPage]:
        """
        Load a page from the archive.

        Args:
            archive_id: Archive ID

        Returns:
            ArchivedPage or None
        """
        page_file = self.pages_path / f"{archive_id}.gz.enc"

        if not page_file.exists():
            return None

        try:
            data = page_file.read_bytes()

            # Decrypt
            if self._fernet:
                data = self._fernet.decrypt(data)

            # Decompress
            data = gzip.decompress(data)

            return ArchivedPage.from_dict(json.loads(data.decode("utf-8")))

        except Exception as e:
            logger.error(f"Error loading archived page {archive_id}: {e}")
            return None

    async def delete(self, archive_id: str) -> bool:
        """
        Delete a page from the archive.

        Args:
            archive_id: Archive ID

        Returns:
            True if deleted
        """
        page_file = self.pages_path / f"{archive_id}.gz.enc"

        if page_file.exists():
            page_file.unlink()

        if archive_id in self._index:
            del self._index[archive_id]

        # Remove from search index
        for word, ids in list(self._search_index.items()):
            if archive_id in ids:
                ids.remove(archive_id)
                if not ids:
                    del self._search_index[word]

        await self._save_index()

        return True

    async def list(
        self,
        limit: int = 50,
        search: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        List archived pages.

        Args:
            limit: Maximum results
            search: Optional search query

        Returns:
            List of page metadata
        """
        if search:
            return await self.search(search, limit)

        # Return most recent
        items = sorted(self._index.items(), key=lambda x: x[1].get("archived_at", ""), reverse=True)

        return [{"id": id, **meta} for id, meta in items[:limit]]

    async def search(
        self,
        query: str,
        limit: int = 50,
    ) -> List[Dict[str, Any]]:
        """
        Full-text search across archived pages.

        Args:
            query: Search query
            limit: Maximum results

        Returns:
            List of matching page metadata
        """
        # Tokenize query
        words = self._tokenize(query.lower())

        if not words:
            return []

        # Find pages matching all words
        matching_ids = None
        for word in words:
            word_ids = set()

            # Match word prefixes
            for indexed_word, ids in self._search_index.items():
                if indexed_word.startswith(word):
                    word_ids.update(ids)

            if matching_ids is None:
                matching_ids = word_ids
            else:
                matching_ids &= word_ids

        if not matching_ids:
            return []

        # Get metadata for matches
        results = []
        for id in matching_ids:
            if id in self._index:
                results.append({"id": id, **self._index[id]})

        # Sort by recency
        results.sort(key=lambda x: x.get("archived_at", ""), reverse=True)

        return results[:limit]

    def _index_text(self, archive_id: str, text: str):
        """Add text to search index."""
        words = self._tokenize(text.lower())

        # Limit to 1000 unique words per page
        words = list(set(words))[:1000]

        for word in words:
            if word not in self._search_index:
                self._search_index[word] = []
            if archive_id not in self._search_index[word]:
                self._search_index[word].append(archive_id)

    def _tokenize(self, text: str) -> List[str]:
        """Tokenize text for indexing."""
        # Extract words (3+ chars)
        words = re.findall(r"\b[a-z]{3,}\b", text)
        return words

    async def get_stats(self) -> Dict[str, Any]:
        """Get archive statistics."""
        total_size = 0
        for file in self.pages_path.glob("*.gz.enc"):
            total_size += file.stat().st_size

        return {
            "page_count": len(self._index),
            "index_words": len(self._search_index),
            "total_size_bytes": total_size,
            "total_size_mb": round(total_size / (1024 * 1024), 2),
            "encrypted": self._fernet is not None,
        }

    async def cleanup(self, max_age_days: int = 365):
        """
        Remove old archived pages.

        Args:
            max_age_days: Maximum age in days
        """
        from datetime import timedelta

        cutoff = datetime.utcnow() - timedelta(days=max_age_days)
        cutoff_str = cutoff.isoformat()

        to_delete = []
        for id, meta in self._index.items():
            if meta.get("archived_at", "") < cutoff_str:
                to_delete.append(id)

        for id in to_delete:
            await self.delete(id)

        if to_delete:
            logger.info(f"Cleaned up {len(to_delete)} old archived pages")

    async def export(
        self,
        archive_id: str,
        format: str = "html",
    ) -> Optional[bytes]:
        """
        Export an archived page.

        Args:
            archive_id: Archive ID
            format: Export format (html, txt, pdf)

        Returns:
            Exported content
        """
        page = await self.load(archive_id)
        if not page:
            return None

        if format == "txt":
            return page.content_text.encode("utf-8")
        elif format == "html":
            return (
                page.reader_html.encode("utf-8")
                if page.reader_html
                else page.content_html.encode("utf-8")
            )
        else:
            return page.content_html.encode("utf-8")
