# 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

"""
Persistent Memory System

Stores:
- Long-term facts about the user (preferences, context, important info)
- Conversation history per channel/user
- Learned behaviors and patterns

Memory is saved to disk and survives restarts.
Uses atomic writes to prevent corruption on crash.
Supports optional encryption at rest via FAMILIAR_ENCRYPTION_KEY.
"""

import json
import logging
import os
import threading
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, List, Optional

from .config import HISTORY_FILE, MEMORY_FILE
from .utils import atomic_write_json

# Embedding imports (graceful — embeddings are optional)
try:
    from .embeddings import (
        EmbeddingProvider,
        get_embedding_provider,
        top_k_similar,
    )

    EMBEDDINGS_AVAILABLE = True
except ImportError:
    EMBEDDINGS_AVAILABLE = False
    EmbeddingProvider = None

logger = logging.getLogger(__name__)

# Import encryption module (graceful degradation)
try:
    from .encryption import CRYPTOGRAPHY_AVAILABLE, EncryptedStorage
except ImportError:
    CRYPTOGRAPHY_AVAILABLE = False
    EncryptedStorage = None

# Global encryption instance for memory
_memory_encryption: Optional["EncryptedStorage"] = None
_memory_encryption_initialized: bool = False
_memory_encryption_lock = threading.Lock()


def _get_memory_encryption() -> Optional["EncryptedStorage"]:
    """Get or initialize memory encryption."""
    global _memory_encryption, _memory_encryption_initialized

    if _memory_encryption_initialized:
        return _memory_encryption

    with _memory_encryption_lock:
        if _memory_encryption_initialized:
            return _memory_encryption

        encryption_key = os.environ.get("FAMILIAR_ENCRYPTION_KEY")
        if encryption_key and CRYPTOGRAPHY_AVAILABLE and EncryptedStorage:
            try:
                _memory_encryption = EncryptedStorage()
                if _memory_encryption.is_available:
                    logger.info("Memory encryption enabled")
                else:
                    _memory_encryption = None
            except Exception as e:
                logger.warning(f"Failed to initialize memory encryption: {e}")
                _memory_encryption = None

        _memory_encryption_initialized = True

    return _memory_encryption


@dataclass
class MemoryEntry:
    """A single memory/fact."""

    key: str
    value: str
    category: str  # "user_info", "preference", "fact", "task", "relationship"
    created_at: str
    updated_at: str
    source: str  # Which conversation/channel it came from
    importance: int = 5  # 1-10, higher = more important


class Memory:
    """
    Persistent memory store with thread-safe atomic writes.

    Usage:
        memory = Memory()
        memory.remember("user_name", "George", category="user_info")
        memory.remember("cooking_style", "French technique with indigenous influences", category="preference")

        name = memory.recall("user_name")
        all_prefs = memory.recall_category("preference")
    """

    def __init__(self, memory_file: Path = MEMORY_FILE, prefer_local_embeddings: bool = False):
        self.memory_file = memory_file
        self.memories: dict[str, MemoryEntry] = {}
        self._lock = threading.RLock()  # Use RLock for reentrant locking

        # Semantic search: embedding provider and vector cache
        self._embedder: Optional["EmbeddingProvider"] = None
        self._vectors: dict[str, List[float]] = {}  # key -> embedding vector
        self._vectors_dirty = False
        self._prefer_local = prefer_local_embeddings

        self._load()
        self._init_embeddings()

    def _load(self):
        """Load memories from disk, with encryption support."""
        encryption = _get_memory_encryption()
        encrypted_file = self.memory_file.with_suffix(".json.enc")

        # Try encrypted file first
        if encrypted_file.exists():
            if encryption and encryption.is_available:
                try:
                    encrypted_data = encrypted_file.read_text()
                    decrypted_data = encryption.decrypt(encrypted_data)
                    data = json.loads(decrypted_data)
                    for key, entry_data in data.items():
                        self.memories[key] = MemoryEntry(**entry_data)
                    logger.debug("Loaded encrypted memory file")
                    return
                except Exception as e:
                    logger.error(f"Failed to load encrypted memory: {e}")
                    # Try backup
                    backup = encrypted_file.with_suffix(".enc.bak")
                    if backup.exists():
                        try:
                            encrypted_data = backup.read_text()
                            decrypted_data = encryption.decrypt(encrypted_data)
                            data = json.loads(decrypted_data)
                            for key, entry_data in data.items():
                                self.memories[key] = MemoryEntry(**entry_data)
                            logger.warning("Loaded encrypted memory from backup")
                            return
                        except Exception as e:
                            logger.debug("Error suppressed: %s", e)
                    return  # Don't fall back to unencrypted
            else:
                logger.warning("Encrypted memory exists but encryption not available")
                return

        # Try unencrypted file
        if self.memory_file.exists():
            try:
                with open(self.memory_file, "r") as f:
                    data = json.load(f)
                    for key, entry_data in data.items():
                        self.memories[key] = MemoryEntry(**entry_data)

                # Auto-migrate to encrypted if encryption is now enabled
                if encryption and encryption.is_available:
                    logger.info("Migrating memory to encrypted storage")
                    self._save()  # Will save encrypted
                    # Remove unencrypted file after successful migration
                    try:
                        self.memory_file.unlink()
                    except Exception as e:
                        logger.debug("Error suppressed: %s", e)
            except json.JSONDecodeError as e:
                # Try to load backup if main file is corrupted
                backup = self.memory_file.with_suffix(".json.bak")
                if backup.exists():
                    try:
                        with open(backup, "r") as f:
                            data = json.load(f)
                            for key, entry_data in data.items():
                                self.memories[key] = MemoryEntry(**entry_data)
                        logger.warning(f"Loaded memory from backup due to corruption: {e}")
                    except Exception:
                        logger.error(f"Could not load memory or backup: {e}")
                else:
                    logger.error(f"Could not load memory (corrupted): {e}")
            except Exception as e:
                logger.error(f"Could not load memory: {e}")

    # ── Embedding Layer ────────────────────────────────────────

    def _init_embeddings(self):
        """Initialize the embedding provider (cascading fallback)."""
        if not EMBEDDINGS_AVAILABLE:
            logger.debug("Embeddings module not available — using substring search only")
            return

        try:
            self._embedder = get_embedding_provider(prefer_local=self._prefer_local)
            logger.info(f"Memory embeddings: {self._embedder.name}")
            self._load_vectors()
            self._sync_vectors()
        except Exception as e:
            logger.warning(f"Could not initialize embeddings: {e}. Using substring search.")
            self._embedder = None

    def _vectors_file(self) -> Path:
        """Path to the vector cache file."""
        return self.memory_file.with_name("memory_vectors.json")

    def _load_vectors(self):
        """Load cached vectors from disk."""
        vf = self._vectors_file()
        if not vf.exists():
            return
        try:
            data = json.loads(vf.read_text())
            cached_provider = data.get("provider", "")
            # Invalidate cache if provider changed (different dimensions)
            if self._embedder and cached_provider != self._embedder.name:
                logger.info(
                    f"Embedding provider changed ({cached_provider} → {self._embedder.name}), "
                    "re-embedding all memories"
                )
                self._vectors = {}
                return
            self._vectors = data.get("vectors", {})
            logger.debug(f"Loaded {len(self._vectors)} cached vectors")
        except Exception as e:
            logger.warning(f"Could not load vector cache: {e}")
            self._vectors = {}

    def _save_vectors(self):
        """Persist vector cache to disk."""
        if not self._vectors_dirty or not self._embedder:
            return
        vf = self._vectors_file()
        try:
            vf.parent.mkdir(parents=True, exist_ok=True)
            data = {
                "provider": self._embedder.name,
                "dimensions": self._embedder.dimensions
                if hasattr(self._embedder, "_dimensions") and self._embedder._dimensions
                else len(next(iter(self._vectors.values()), [])),
                "count": len(self._vectors),
                "vectors": self._vectors,
            }
            # Use atomic write pattern
            tmp = vf.with_suffix(".json.tmp")
            tmp.write_text(json.dumps(data))
            tmp.replace(vf)
            self._vectors_dirty = False
            logger.debug(f"Saved {len(self._vectors)} vectors")
        except Exception as e:
            logger.warning(f"Could not save vector cache: {e}")

    def _sync_vectors(self):
        """Ensure all memories have embeddings, remove orphans."""
        if not self._embedder:
            return

        # Find memories missing vectors
        missing = [key for key in self.memories if key not in self._vectors]

        # Remove orphaned vectors (memory was deleted)
        orphans = [key for key in self._vectors if key not in self.memories]
        for key in orphans:
            del self._vectors[key]
            self._vectors_dirty = True

        # Embed missing entries
        if missing:
            texts = [f"{self.memories[k].key}: {self.memories[k].value}" for k in missing]
            try:
                vectors = self._embedder.embed_batch(texts)
                for key, vec in zip(missing, vectors):
                    self._vectors[key] = vec
                self._vectors_dirty = True
                logger.info(f"Embedded {len(missing)} new memories")
            except Exception as e:
                logger.warning(f"Failed to embed {len(missing)} memories: {e}")

        if self._vectors_dirty:
            self._save_vectors()

    def _embed_entry(self, key: str):
        """Embed a single memory entry and cache it."""
        if not self._embedder or key not in self.memories:
            return
        entry = self.memories[key]
        try:
            vec = self._embedder.embed(f"{entry.key}: {entry.value}")
            self._vectors[key] = vec
            self._vectors_dirty = True
        except Exception as e:
            logger.warning(f"Failed to embed memory '{key}': {e}")

    def _save(self):
        """Save memories to disk atomically, with encryption support."""
        data = {k: asdict(v) for k, v in self.memories.items()}
        json_data = json.dumps(data, indent=2)

        encryption = _get_memory_encryption()

        if encryption and encryption.is_available:
            # Save encrypted
            encrypted_file = self.memory_file.with_suffix(".json.enc")

            # Create backup of current encrypted file
            if encrypted_file.exists():
                backup = encrypted_file.with_suffix(".enc.bak")
                try:
                    import shutil

                    shutil.copy2(encrypted_file, backup)
                except Exception as e:
                    logger.debug("Error suppressed: %s", e)
            try:
                encrypted_data = encryption.encrypt(json_data)

                # Atomic write
                tmp_file = encrypted_file.with_suffix(".enc.tmp")
                tmp_file.write_text(encrypted_data)
                tmp_file.chmod(0o600)
                tmp_file.replace(encrypted_file)

            except Exception as e:
                logger.error(f"Failed to save encrypted memory: {e}")
                # NEVER fall back to plaintext — the user expects encryption.
                raise RuntimeError(
                    f"Encrypted memory save failed and plaintext fallback is "
                    f"disabled for safety: {e}"
                ) from e
        else:
            # Save unencrypted
            # Create backup of current file before overwriting
            if self.memory_file.exists():
                backup = self.memory_file.with_suffix(".json.bak")
                try:
                    import shutil

                    shutil.copy2(self.memory_file, backup)
                except Exception:
                    pass  # Backup failure is not critical

            atomic_write_json(self.memory_file, data)

    def remember(
        self,
        key: str,
        value: str,
        category: str = "fact",
        source: str = "conversation",
        importance: int = 5,
    ):
        """Store a memory."""
        now = datetime.now().isoformat()

        with self._lock:
            if key in self.memories:
                # Update existing
                self.memories[key].value = value
                self.memories[key].updated_at = now
                self.memories[key].importance = max(self.memories[key].importance, importance)
            else:
                # Create new
                self.memories[key] = MemoryEntry(
                    key=key,
                    value=value,
                    category=category,
                    created_at=now,
                    updated_at=now,
                    source=source,
                    importance=importance,
                )
            self._save()
            self._embed_entry(key)
            if self._vectors_dirty:
                self._save_vectors()

    def recall(self, key: str) -> Optional[str]:
        """Retrieve a specific memory."""
        entry = self.memories.get(key)
        return entry.value if entry else None

    def recall_category(self, category: str) -> dict[str, str]:
        """Get all memories in a category."""
        return {k: v.value for k, v in self.memories.items() if v.category == category}

    def recall_all(self) -> dict[str, MemoryEntry]:
        """Get all memories."""
        return self.memories.copy()

    def forget(self, key: str) -> bool:
        """Remove a memory and its embedding vector."""
        with self._lock:
            if key in self.memories:
                del self.memories[key]
                if key in self._vectors:
                    del self._vectors[key]
                    self._vectors_dirty = True
                self._save()
                if self._vectors_dirty:
                    self._save_vectors()
                return True
            return False

    def search(self, query: str, max_results: int = 10) -> list[MemoryEntry]:
        """
        Search memories using semantic similarity with substring fallback.

        When embeddings are available, computes cosine similarity between
        the query and all memory vectors. Combines with substring matching
        to ensure exact-match queries still work perfectly.

        Args:
            query: Search query (natural language)
            max_results: Maximum number of results to return

        Returns:
            List of MemoryEntry sorted by relevance
        """
        semantic_results = []
        substring_results = self._substring_search(query, max_results)

        # Try semantic search
        if self._embedder and self._vectors:
            try:
                query_vec = self._embedder.embed(query)
                vector_pairs = [
                    (key, vec) for key, vec in self._vectors.items() if key in self.memories
                ]
                if vector_pairs:
                    ranked = top_k_similar(query_vec, vector_pairs, k=max_results)
                    for key, score in ranked:
                        if score > 0.15 and key in self.memories:
                            semantic_results.append(self.memories[key])
            except Exception as e:
                logger.warning(f"Semantic search failed: {e}")

        # Merge: substring matches first (exact > semantic), then semantic
        seen = set()
        merged = []

        # Substring hits are highest confidence
        for entry in substring_results:
            if entry.key not in seen:
                seen.add(entry.key)
                merged.append(entry)

        # Semantic results fill in what substring missed
        for entry in semantic_results:
            if entry.key not in seen:
                seen.add(entry.key)
                merged.append(entry)

        return merged[:max_results]

    def _substring_search(self, query: str, max_results: int = 10) -> list[MemoryEntry]:
        """Original substring search — always available as fallback."""
        query = query.lower()
        results = []
        for entry in self.memories.values():
            if query in entry.key.lower() or query in entry.value.lower():
                results.append(entry)
        return sorted(results, key=lambda x: x.importance, reverse=True)[:max_results]

    def get_relevant_context(self, query: str, max_results: int = 10) -> str:
        """
        Get memories relevant to a specific query, formatted for LLM context.

        Unlike get_context_string() which dumps the top-N by importance globally,
        this method returns memories semantically relevant to the current user message.

        Args:
            query: The user's current message
            max_results: Maximum memories to include

        Returns:
            Formatted string for injection into system prompt, or empty string
        """
        relevant = self.search(query, max_results=max_results)
        if not relevant:
            return ""

        lines = ["<relevant_memories>"]

        # Group by category
        categories: dict[str, list[MemoryEntry]] = {}
        for m in relevant:
            if m.category not in categories:
                categories[m.category] = []
            categories[m.category].append(m)

        for cat, entries in categories.items():
            lines.append(f"\n### {cat.replace('_', ' ').title()}")
            for e in entries:
                lines.append(f"- {e.key}: {e.value}")

        lines.append("\n</relevant_memories>")
        return "\n".join(lines)

    def get_context_string(self, max_entries: int = 20) -> str:
        """Get memories formatted for LLM context."""
        if not self.memories:
            return ""

        # Sort by importance and recency
        sorted_memories = sorted(
            self.memories.values(), key=lambda x: (x.importance, x.updated_at), reverse=True
        )[:max_entries]

        lines = ["<user_knowledge>"]

        # Group by category
        categories: dict[str, list[MemoryEntry]] = {}
        for m in sorted_memories:
            if m.category not in categories:
                categories[m.category] = []
            categories[m.category].append(m)

        for cat, entries in categories.items():
            lines.append(f"\n### {cat.replace('_', ' ').title()}")
            for e in entries:
                lines.append(f"- {e.key}: {e.value}")

        lines.append("\n</user_knowledge>")
        return "\n".join(lines)


class ConversationHistory:
    """
    Stores conversation history per user/channel.
    Persists across restarts with atomic writes.

    Limits:
        - max_per_chat: Maximum messages per individual chat
        - max_total_chats: Maximum number of chats to keep (oldest pruned first)
    """

    # Maximum number of chats to keep to prevent unbounded growth
    MAX_TOTAL_CHATS = 500

    def __init__(
        self,
        history_file: Path = HISTORY_FILE,
        max_per_chat: int = 100,
        max_total_chats: Optional[int] = None,
    ):
        self.history_file = history_file
        self.max_per_chat = max_per_chat
        self.max_total_chats = max_total_chats or self.MAX_TOTAL_CHATS
        self.histories: dict[str, list] = {}  # chat_id -> messages
        self._lock = threading.RLock()
        self._dirty = False  # Track if save is needed
        self._save_interval = 5  # Batch saves to avoid excessive disk writes
        self._last_save = 0
        self._load()

    def _prune_old_chats(self):
        """Remove oldest chats if we exceed the maximum."""
        if len(self.histories) <= self.max_total_chats:
            return

        # Sort chats by their last message timestamp (oldest first)
        def get_last_timestamp(chat_key):
            messages = self.histories.get(chat_key, [])
            if not messages:
                return ""
            last_msg = messages[-1]
            return last_msg.get("timestamp", "")

        sorted_keys = sorted(self.histories.keys(), key=get_last_timestamp)

        # Remove oldest chats
        chats_to_remove = len(self.histories) - self.max_total_chats
        for key in sorted_keys[:chats_to_remove]:
            del self.histories[key]

    def _load(self):
        """Load history from disk, with encryption support."""
        encryption = _get_memory_encryption()
        encrypted_file = self.history_file.with_suffix(".json.enc")

        # Try encrypted file first
        if encrypted_file.exists():
            if encryption and encryption.is_available:
                try:
                    encrypted_data = encrypted_file.read_text()
                    decrypted_data = encryption.decrypt(encrypted_data)
                    self.histories = json.loads(decrypted_data)
                    logger.debug("Loaded encrypted history file")
                    return
                except Exception as e:
                    logger.error(f"Failed to load encrypted history: {e}")
                    # Try backup
                    backup = encrypted_file.with_suffix(".enc.bak")
                    if backup.exists():
                        try:
                            encrypted_data = backup.read_text()
                            decrypted_data = encryption.decrypt(encrypted_data)
                            self.histories = json.loads(decrypted_data)
                            logger.warning("Loaded encrypted history from backup")
                            return
                        except Exception as e:
                            logger.debug("Error suppressed: %s", e)
                    self.histories = {}
                    return
            else:
                logger.warning("Encrypted history exists but encryption not available")
                self.histories = {}
                return

        # Try unencrypted file
        if self.history_file.exists():
            try:
                with open(self.history_file, "r") as f:
                    self.histories = json.load(f)

                # Auto-migrate to encrypted if encryption is now enabled
                if encryption and encryption.is_available:
                    logger.info("Migrating history to encrypted storage")
                    self._save(force=True)  # Will save encrypted
                    # Remove unencrypted file
                    try:
                        self.history_file.unlink()
                    except Exception as e:
                        logger.debug("Error suppressed: %s", e)
            except json.JSONDecodeError as e:
                # Try backup
                backup = self.history_file.with_suffix(".json.bak")
                if backup.exists():
                    try:
                        with open(backup, "r") as f:
                            self.histories = json.load(f)
                        logger.warning(f"Loaded history from backup: {e}")
                    except Exception:
                        logger.error(f"Could not load history or backup: {e}")
                        self.histories = {}
                else:
                    logger.error(f"Could not load history (corrupted): {e}")
                    self.histories = {}
            except Exception as e:
                logger.error(f"Could not load history: {e}")

    def _save(self, force: bool = False):
        """Save history to disk atomically with optional batching, encryption support."""
        import time

        now = time.time()
        if not force and (now - self._last_save) < self._save_interval:
            self._dirty = True
            return

        encryption = _get_memory_encryption()
        json_data = json.dumps(self.histories)  # Compact format

        if encryption and encryption.is_available:
            # Save encrypted
            encrypted_file = self.history_file.with_suffix(".json.enc")

            # Create backup
            if encrypted_file.exists():
                backup = encrypted_file.with_suffix(".enc.bak")
                try:
                    import shutil

                    shutil.copy2(encrypted_file, backup)
                except Exception as e:
                    logger.debug("Error suppressed: %s", e)
            try:
                encrypted_data = encryption.encrypt(json_data)

                # Atomic write
                tmp_file = encrypted_file.with_suffix(".enc.tmp")
                tmp_file.write_text(encrypted_data)
                tmp_file.chmod(0o600)
                tmp_file.replace(encrypted_file)

            except Exception as e:
                logger.error(f"Failed to save encrypted history: {e}")
                # Fall back to unencrypted
                atomic_write_json(self.history_file, self.histories, indent=None)
        else:
            # Save unencrypted
            # Create backup
            if self.history_file.exists():
                backup = self.history_file.with_suffix(".json.bak")
                try:
                    import shutil

                    shutil.copy2(self.history_file, backup)
                except Exception as e:
                    logger.debug("Error suppressed: %s", e)
            atomic_write_json(self.history_file, self.histories, indent=None)  # Compact format

        self._last_save = now
        self._dirty = False

    def flush(self):
        """Force save any pending changes."""
        with self._lock:
            if self._dirty:
                self._save(force=True)

    def _chat_key(self, chat_id: Any, channel: str = "default") -> str:
        """Generate unique key for chat."""
        return f"{channel}:{chat_id}"

    def add_message(self, chat_id: Any, role: str, content: Any, channel: str = "default"):
        """Add a message to history."""
        key = self._chat_key(chat_id, channel)

        with self._lock:
            if key not in self.histories:
                self.histories[key] = []
                # Check if we need to prune old chats when adding new ones
                self._prune_old_chats()

            self.histories[key].append(
                {"role": role, "content": content, "timestamp": datetime.now().isoformat()}
            )

            # Trim if too long
            if len(self.histories[key]) > self.max_per_chat:
                self.histories[key] = self.histories[key][-self.max_per_chat :]

            self._save()

    def get_history(self, chat_id: Any, channel: str = "default", limit: int = 50) -> list:
        """Get conversation history for LLM (without timestamps)."""
        key = self._chat_key(chat_id, channel)
        history = self.histories.get(key, [])

        # Return just role/content for LLM
        return [{"role": m["role"], "content": m["content"]} for m in history[-limit:]]

    def clear(self, chat_id: Any, channel: str = "default"):
        """Clear history for a chat."""
        key = self._chat_key(chat_id, channel)
        with self._lock:
            if key in self.histories:
                self.histories[key] = []
                self._save()

    def get_all_chats(self) -> list[str]:
        """Get all chat IDs."""
        return list(self.histories.keys())


# Memory extraction prompt for LLM
MEMORY_EXTRACTION_PROMPT = """
Based on this conversation, extract any important facts I should remember about the user.
Return a JSON array of objects with: key, value, category, importance (1-10)

Categories: user_info, preference, fact, task, relationship, skill, interest

Only extract genuinely useful information. Don't over-extract trivial things.
If nothing notable, return empty array: []

Example output:
[
  {"key": "name", "value": "George", "category": "user_info", "importance": 9},
  {"key": "occupation", "value": "chef and baker", "category": "user_info", "importance": 8},
  {"key": "interest_physics", "value": "working on unified field theory", "category": "interest", "importance": 7}
]
"""
