# 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 mail storage for Familiar Email Server.

Stores emails in encrypted format using Fernet (AES-128-CBC).
Supports folders, search, and message management.
"""

import asyncio
import hashlib
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 support
try:
    from cryptography.fernet import Fernet

    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False
    logger.debug("cryptography not installed - storage encryption disabled")


@dataclass
class StoredMessage:
    """A stored email message."""

    id: str
    folder: str
    sender: str
    to: List[str]
    cc: List[str]
    subject: str
    date: datetime
    body: str
    html_body: Optional[str]
    attachments: List[Dict[str, Any]]
    headers: Dict[str, str]
    read: bool = False
    flagged: bool = False
    answered: bool = False
    deleted: bool = False
    message_id: Optional[str] = None  # RFC 2822 Message-ID header

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for storage."""
        data = asdict(self)
        data["date"] = self.date.isoformat()
        return data

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "StoredMessage":
        """Create from dictionary."""
        data["date"] = datetime.fromisoformat(data["date"])
        return cls(**data)


@dataclass
class Mailbox:
    """A user's mailbox containing folders and messages."""

    address: str
    password_hash: Optional[str]
    storage: "MailStorage"

    # Standard folders (class default — copied to instance in __post_init__)
    _DEFAULT_FOLDERS = ["INBOX", "Sent", "Drafts", "Trash", "Spam", "Archive"]

    def __post_init__(self):
        self.FOLDERS = list(self._DEFAULT_FOLDERS)

    async def get_messages(
        self,
        folder: str = "INBOX",
        limit: int = 50,
        offset: int = 0,
    ) -> List[StoredMessage]:
        """Get messages from a folder."""
        folder_path = self._get_folder_path(folder)

        if not folder_path.exists():
            return []

        messages = []
        files = sorted(folder_path.glob("*.json"), reverse=True)

        for file in files[offset : offset + limit]:
            try:
                msg = await self.storage._read_message(file)
                if msg:
                    messages.append(msg)
            except Exception as e:
                logger.error(f"Error reading message {file}: {e}")

        return messages

    async def get_message(self, message_id: str) -> Optional[StoredMessage]:
        """Get a specific message by ID."""
        # Search all folders
        for folder in self.FOLDERS:
            folder_path = self._get_folder_path(folder)
            msg_path = folder_path / f"{message_id}.json"

            if msg_path.exists():
                return await self.storage._read_message(msg_path)

        return None

    async def store_message(
        self,
        message: StoredMessage,
    ) -> str:
        """Store a message in the appropriate folder."""
        folder_path = self._get_folder_path(message.folder)
        folder_path.mkdir(parents=True, exist_ok=True)

        msg_path = folder_path / f"{message.id}.json"
        await self.storage._write_message(msg_path, message)

        return message.id

    async def move_message(
        self,
        message_id: str,
        to_folder: str,
    ) -> bool:
        """Move a message to a different folder."""
        # Find the message
        for folder in self.FOLDERS:
            folder_path = self._get_folder_path(folder)
            msg_path = folder_path / f"{message_id}.json"

            if msg_path.exists():
                # Read message
                msg = await self.storage._read_message(msg_path)
                if not msg:
                    return False

                # Update folder
                msg.folder = to_folder

                # Delete from old location
                msg_path.unlink()

                # Store in new location
                await self.store_message(msg)
                return True

        return False

    async def delete_message(self, message_id: str) -> bool:
        """Permanently delete a message."""
        for folder in self.FOLDERS:
            folder_path = self._get_folder_path(folder)
            msg_path = folder_path / f"{message_id}.json"

            if msg_path.exists():
                msg_path.unlink()
                return True

        return False

    async def mark_read(self, message_id: str, read: bool = True) -> bool:
        """Mark a message as read/unread."""
        msg = await self.get_message(message_id)
        if not msg:
            return False

        msg.read = read
        await self.store_message(msg)
        return True

    async def mark_flagged(self, message_id: str, flagged: bool = True) -> bool:
        """Mark a message as flagged/unflagged."""
        msg = await self.get_message(message_id)
        if not msg:
            return False

        msg.flagged = flagged
        await self.store_message(msg)
        return True

    async def search(
        self,
        query: str,
        folder: Optional[str] = None,
        limit: int = 50,
    ) -> List[StoredMessage]:
        """Search messages by text."""
        results = []
        query_lower = query.lower()

        folders = [folder] if folder else self.FOLDERS

        for f in folders:
            folder_path = self._get_folder_path(f)
            if not folder_path.exists():
                continue

            for file in folder_path.glob("*.json"):
                try:
                    msg = await self.storage._read_message(file)
                    if not msg:
                        continue

                    # Search in subject, sender, body
                    searchable = (f"{msg.subject} {msg.sender} {msg.body}").lower()

                    if query_lower in searchable:
                        results.append(msg)

                        if len(results) >= limit:
                            return results

                except Exception as e:
                    logger.error(f"Error searching message {file}: {e}")

        return results

    async def get_folder_counts(self) -> Dict[str, Dict[str, int]]:
        """Get message counts per folder."""
        counts = {}

        for folder in self.FOLDERS:
            folder_path = self._get_folder_path(folder)

            if not folder_path.exists():
                counts[folder] = {"total": 0, "unread": 0}
                continue

            total = 0
            unread = 0

            for file in folder_path.glob("*.json"):
                total += 1
                try:
                    msg = await self.storage._read_message(file)
                    if msg and not msg.read:
                        unread += 1
                except Exception:
                    pass

            counts[folder] = {"total": total, "unread": unread}

        return counts

    async def copy_message(self, message_id: str, to_folder: str) -> Optional[str]:
        """Copy a message to another folder. Returns the new message ID."""
        msg = await self.get_message(message_id)
        if not msg:
            return None

        # Create a copy with a new ID
        new_id = str(uuid.uuid4())
        copy = StoredMessage(
            id=new_id,
            folder=to_folder,
            sender=msg.sender,
            to=list(msg.to),
            cc=list(msg.cc),
            subject=msg.subject,
            date=msg.date,
            body=msg.body,
            html_body=msg.html_body,
            attachments=list(msg.attachments),
            headers=dict(msg.headers),
            read=msg.read,
            flagged=msg.flagged,
            answered=msg.answered,
            deleted=False,
            message_id=msg.message_id,
        )

        await self.store_message(copy)
        return new_id

    async def create_folder(self, folder_name: str) -> bool:
        """Create a new folder. Returns True on success."""
        safe_name = re.sub(r"[^\w\-]", "_", folder_name)
        if safe_name in self.FOLDERS:
            return False

        folder_path = self._get_folder_path(safe_name)
        folder_path.mkdir(parents=True, exist_ok=True)
        self.FOLDERS.append(safe_name)
        return True

    async def delete_folder(self, folder_name: str) -> bool:
        """Delete a folder. INBOX cannot be deleted."""
        if folder_name.upper() == "INBOX":
            return False
        if folder_name not in self.FOLDERS:
            return False

        import shutil

        folder_path = self._get_folder_path(folder_name)
        if folder_path.exists():
            shutil.rmtree(folder_path)
        self.FOLDERS.remove(folder_name)
        return True

    async def rename_folder(self, old_name: str, new_name: str) -> bool:
        """Rename a folder. INBOX cannot be renamed."""
        if old_name.upper() == "INBOX":
            return False
        if old_name not in self.FOLDERS:
            return False

        safe_new = re.sub(r"[^\w\-]", "_", new_name)
        if safe_new in self.FOLDERS:
            return False

        old_path = self._get_folder_path(old_name)
        new_path = self._get_folder_path(safe_new)

        if old_path.exists():
            old_path.rename(new_path)

        idx = self.FOLDERS.index(old_name)
        self.FOLDERS[idx] = safe_new
        return True

    def _get_folder_path(self, folder: str) -> Path:
        """Get the path for a folder."""
        # Sanitize folder name
        safe_folder = re.sub(r"[^\w\-]", "_", folder)
        mailbox_path = self.storage.storage_path / "mailboxes" / self.address
        return mailbox_path / safe_folder


class MailStorage:
    """
    Encrypted mail storage manager.

    Directory structure:
    storage_path/
    ├── mailboxes/
    │   └── user@domain.com/
    │       ├── INBOX/
    │       │   ├── msg-uuid-1.json
    │       │   └── msg-uuid-2.json
    │       ├── Sent/
    │       ├── Drafts/
    │       ├── Trash/
    │       ├── Spam/
    │       └── Archive/
    ├── users.json (encrypted user list)
    └── .encryption_key
    """

    def __init__(
        self,
        storage_path: Path,
        encryption_key: Optional[bytes] = None,
    ):
        self.storage_path = storage_path
        self.encryption_key = encryption_key
        self._fernet = None

        if encryption_key and HAS_CRYPTO:
            self._fernet = Fernet(encryption_key)

        self._mailboxes: Dict[str, Mailbox] = {}
        self._lock = asyncio.Lock()

    async def initialize(self):
        """Initialize storage."""
        self.storage_path.mkdir(parents=True, exist_ok=True)
        (self.storage_path / "mailboxes").mkdir(exist_ok=True)

        # Load existing mailboxes
        users_file = self.storage_path / "users.json"
        if users_file.exists():
            data = await self._read_encrypted_file(users_file)
            if data:
                users = json.loads(data)
                for user in users:
                    self._mailboxes[user["address"]] = Mailbox(
                        address=user["address"],
                        password_hash=user.get("password_hash"),
                        storage=self,
                    )

        logger.info(f"Mail storage initialized at {self.storage_path}")

    async def create_mailbox(
        self,
        address: str,
        password_hash: Optional[str] = None,
    ) -> bool:
        """Create a new mailbox."""
        async with self._lock:
            if address in self._mailboxes:
                return False

            # Create mailbox
            mailbox = Mailbox(
                address=address,
                password_hash=password_hash,
                storage=self,
            )

            # Create folder structure
            mailbox_path = self.storage_path / "mailboxes" / address
            mailbox_path.mkdir(parents=True, exist_ok=True)

            for folder in Mailbox._DEFAULT_FOLDERS:
                (mailbox_path / folder).mkdir(exist_ok=True)

            self._mailboxes[address] = mailbox

            # Save users list
            await self._save_users()

            logger.info(f"Created mailbox for {address}")
            return True

    async def delete_mailbox(self, address: str) -> bool:
        """Delete a mailbox and all its messages."""
        async with self._lock:
            if address not in self._mailboxes:
                return False

            # Remove from memory
            del self._mailboxes[address]

            # Remove from disk
            import shutil

            mailbox_path = self.storage_path / "mailboxes" / address
            if mailbox_path.exists():
                shutil.rmtree(mailbox_path)

            # Save users list
            await self._save_users()

            logger.info(f"Deleted mailbox for {address}")
            return True

    async def get_mailbox(self, address: str) -> Optional[Mailbox]:
        """Get a mailbox by address."""
        return self._mailboxes.get(address)

    async def list_mailboxes(self) -> List[str]:
        """List all mailbox addresses."""
        return list(self._mailboxes.keys())

    async def authenticate(self, address: str, password: str) -> bool:
        """Authenticate a user."""
        mailbox = self._mailboxes.get(address)
        if not mailbox or not mailbox.password_hash:
            return False

        stored = mailbox.password_hash
        if ":" in stored:
            # PBKDF2 format: salt:hash
            salt, expected = stored.split(":", 1)
            candidate = hashlib.pbkdf2_hmac(
                "sha256", password.encode(), salt.encode(), 100_000
            ).hex()
            import hmac as _hmac

            return _hmac.compare_digest(candidate, expected)
        else:
            # Legacy SHA-256 (backward compat)
            import hmac as _hmac

            return _hmac.compare_digest(hashlib.sha256(password.encode()).hexdigest(), stored)

    async def get_stats(self) -> Dict[str, Any]:
        """Get storage statistics."""
        total_messages = 0
        total_size = 0

        mailboxes_path = self.storage_path / "mailboxes"
        if mailboxes_path.exists():
            for file in mailboxes_path.rglob("*.json"):
                total_messages += 1
                total_size += file.stat().st_size

        return {
            "mailbox_count": len(self._mailboxes),
            "total_messages": total_messages,
            "total_size_bytes": total_size,
            "total_size_mb": round(total_size / (1024 * 1024), 2),
            "encryption_enabled": self._fernet is not None,
        }

    async def _save_users(self):
        """Save users list to disk."""
        users = [
            {
                "address": mb.address,
                "password_hash": mb.password_hash,
            }
            for mb in self._mailboxes.values()
        ]

        users_file = self.storage_path / "users.json"
        await self._write_encrypted_file(users_file, json.dumps(users))

    async def _read_message(self, path: Path) -> Optional[StoredMessage]:
        """Read and decrypt a message."""
        data = await self._read_encrypted_file(path)
        if not data:
            return None

        try:
            return StoredMessage.from_dict(json.loads(data))
        except Exception as e:
            logger.error(f"Error parsing message {path}: {e}")
            return None

    async def _write_message(self, path: Path, message: StoredMessage):
        """Encrypt and write a message."""
        data = json.dumps(message.to_dict())
        await self._write_encrypted_file(path, data)

    async def _read_encrypted_file(self, path: Path) -> Optional[str]:
        """Read and optionally decrypt a file."""
        if not path.exists():
            return None

        content = path.read_bytes()

        if self._fernet:
            try:
                content = self._fernet.decrypt(content)
            except Exception as e:
                logger.error(f"Decryption failed for {path}: {e}")
                return None

        return content.decode("utf-8")

    async def _write_encrypted_file(self, path: Path, data: str):
        """Optionally encrypt and write a file."""
        content = data.encode("utf-8")

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

        # Atomic write
        temp_path = path.with_suffix(".tmp")
        temp_path.write_bytes(content)
        temp_path.replace(path)
        path.chmod(0o600)


def generate_message_id(domain: str) -> str:
    """Generate a unique message ID."""
    unique = uuid.uuid4().hex[:16]
    timestamp = int(datetime.utcnow().timestamp())
    return f"{unique}.{timestamp}@{domain}"
