# 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

"""
MemoryManager — extracted from agent.py (PR 4).

Handles episodic memory lifecycle:
- Per-user episodic memory stores with double-checked locking
- Async persistence to disk (plaintext or encrypted)
- remember/recall/clear_history convenience methods
- User history load/save for Phase 1 multi-user
- Episodic encryption singleton
"""

import json
import logging
import os
import threading
from pathlib import Path
from typing import Any, Dict, Optional

from .episodic_memory import (
    FileMemoryPersistence as EpisodicPersistence,
)
from .episodic_memory import (
    MemorySystem as EpisodicMemorySystem,
)

# Phase 1: Multi-user support
try:
    from .data_store import get_data_store
    from .users import UserContext

    HAS_USER_MANAGEMENT = True
except ImportError:
    HAS_USER_MANAGEMENT = False
    UserContext = None
    get_data_store = None

logger = logging.getLogger(__name__)


# ── EpisodicMemory encryption singleton ──────────────────────────────────────
try:
    from .encryption import CRYPTOGRAPHY_AVAILABLE as _EP_CRYPTO
    from .encryption import EncryptedStorage as _EpisodicEncStorage
except ImportError:
    _EpisodicEncStorage = None
    _EP_CRYPTO = False

_episodic_encryption = None
_episodic_encryption_initialized = False
_episodic_encryption_lock = threading.Lock()


def _get_episodic_encryption():
    """
    Return the EncryptedStorage singleton for episodic memory, or None.
    """
    global _episodic_encryption, _episodic_encryption_initialized
    if _episodic_encryption_initialized:
        return _episodic_encryption
    with _episodic_encryption_lock:
        if _episodic_encryption_initialized:
            return _episodic_encryption
        key = os.environ.get("FAMILIAR_ENCRYPTION_KEY")
        if key and _EP_CRYPTO and _EpisodicEncStorage:
            try:
                enc = _EpisodicEncStorage()
                if enc.is_available:
                    _episodic_encryption = enc
                    logger.info("EpisodicMemory encryption enabled")
                else:
                    logger.warning(
                        "FAMILIAR_ENCRYPTION_KEY set but encryption unavailable — "
                        "EpisodicMemory will use plaintext storage"
                    )
            except Exception as e:
                logger.warning(f"EpisodicMemory encryption init failed: {e}")
        _episodic_encryption_initialized = True
    return _episodic_encryption


class MemoryManager:
    """Manages episodic memory, remember/recall, and user history."""

    def __init__(self, agent_ref):
        self._agent = agent_ref
        self._episodic_stores: Dict[str, EpisodicMemorySystem] = {}
        self._episodic_lock = threading.Lock()

        config = agent_ref.config
        self._episodic_enabled = getattr(config.agent, "memory_enabled", True)
        self._episodic_dir = (
            Path.home() / ".familiar" / "episodic" if self._episodic_enabled else None
        )
        if self._episodic_enabled and self._episodic_dir:
            self._episodic_dir.mkdir(parents=True, exist_ok=True)
            logger.debug(f"EpisodicMemory active: store={self._episodic_dir}")

    def get_episodic_memory(self, user_id: str) -> Optional[EpisodicMemorySystem]:
        """
        Return the EpisodicMemorySystem for this user, creating and loading
        it from disk on first access. Thread-safe: double-checked locking.
        """
        if not self._episodic_enabled or self._episodic_dir is None:
            return None

        uid = str(user_id)
        if uid in self._episodic_stores:
            return self._episodic_stores[uid]

        with self._episodic_lock:
            if uid in self._episodic_stores:
                return self._episodic_stores[uid]

            config = self._agent.config
            ms = EpisodicMemorySystem(
                max_memories=getattr(config.agent, "episodic_max_memories", 5000),
                working_memory_size=getattr(config.agent, "episodic_working_size", 20),
                decay_rate=getattr(config.agent, "episodic_decay_rate", 0.05),
            )

            import re

            safe_uid = re.sub(r"[^a-zA-Z0-9_\-.]", "_", uid)
            enc_path = self._episodic_dir / f"{safe_uid}.json.enc"
            plain_path = self._episodic_dir / f"{safe_uid}.json"
            loaded = False

            # Try encrypted file first
            encryption = _get_episodic_encryption()
            if encryption and encryption.is_available and enc_path.exists():
                try:
                    encrypted_data = enc_path.read_text()
                    json_data = encryption.decrypt(encrypted_data)
                    data = json.loads(json_data)
                    ms.import_memories(data)
                    loaded = True
                    logger.debug(
                        f"EpisodicMemory loaded (encrypted) "
                        f"{ms.get_stats().get('total_stored', 0)} memories for user {uid}"
                    )
                except Exception as e:
                    logger.error(f"EpisodicMemory failed to load encrypted file for {uid}: {e}")
                    loaded = False

            # Fall back to plaintext
            if not loaded and plain_path.exists():
                try:
                    persistence = EpisodicPersistence(str(plain_path))
                    loaded = persistence.load(ms)
                    if loaded:
                        logger.debug(
                            f"EpisodicMemory loaded (plaintext) "
                            f"{ms.get_stats().get('total_stored', 0)} memories for user {uid}"
                        )
                        if encryption and encryption.is_available:
                            logger.info(
                                f"EpisodicMemory: plaintext found for {uid}, "
                                f"will migrate to encrypted on next save"
                            )
                except Exception as e:
                    logger.error(f"EpisodicMemory failed to load plaintext file for {uid}: {e}")

            self._episodic_stores[uid] = ms
            return ms

    def save_episodic_memory_async(self, user_id: str):
        """Persist the user's EpisodicMemorySystem to disk in a daemon thread."""
        ms = self._episodic_stores.get(str(user_id))
        if ms is None or self._episodic_dir is None:
            return

        import re

        safe_uid = re.sub(r"[^a-zA-Z0-9_\-.]", "_", str(user_id))
        plaintext_path = self._episodic_dir / f"{safe_uid}.json"
        enc_path = self._episodic_dir / f"{safe_uid}.json.enc"

        try:
            snapshot = ms.export()
        except Exception as e:
            logger.warning(f"EpisodicMemory export failed for {user_id}: {e}")
            return

        def _save():
            try:
                json_data = json.dumps(snapshot, indent=2, default=str)
                encryption = _get_episodic_encryption()

                if encryption and encryption.is_available:
                    if enc_path.exists():
                        backup = enc_path.with_suffix(".enc.bak")
                        try:
                            import shutil

                            shutil.copy2(enc_path, backup)
                        except Exception:
                            pass

                    encrypted_data = encryption.encrypt(json_data)
                    tmp_path = enc_path.with_suffix(".enc.tmp")
                    tmp_path.write_text(encrypted_data)
                    tmp_path.chmod(0o600)
                    tmp_path.replace(enc_path)

                    if plaintext_path.exists():
                        plaintext_path.unlink(missing_ok=True)

                    logger.debug(f"EpisodicMemory saved (encrypted) for user {user_id}")
                else:
                    tmp_path = plaintext_path.with_suffix(".json.tmp")
                    tmp_path.write_text(json_data)
                    tmp_path.replace(plaintext_path)
                    logger.debug(f"EpisodicMemory saved (plaintext) for user {user_id}")

            except Exception as e:
                logger.warning(f"EpisodicMemory save failed for {user_id}: {e}")

        t = threading.Thread(target=_save, daemon=True)
        t.start()

    # ── Convenience methods ────────────────────────────────────────

    def remember(self, key: str, value: str, category: str = "fact"):
        """Store something in memory."""
        agent = self._agent
        if agent.memory:
            agent.memory.remember(key, value, category=category)

    def recall(self, query: str) -> list:
        """Search memory."""
        agent = self._agent
        if agent.memory:
            return agent.memory.search(query)
        return []

    def clear_history(self, user_id: Any = "default", channel: str = "default"):
        """Clear conversation history for a user."""
        self._agent.history.clear(user_id, channel)

    # ── Phase 1: User history helpers ─────────────────────────────

    def get_user_history(self, user_context: "UserContext", channel: str, limit: int = 20) -> list:
        """Get conversation history from user's isolated data store."""
        agent = self._agent
        if not HAS_USER_MANAGEMENT or not get_data_store:
            return agent.history.get_history(user_context.user.id, channel)

        try:
            store = get_data_store(
                user_id=user_context.user.id, user_role=user_context.user.role.value
            )
            conv_id = f"{channel}_current"
            messages = store.load_history(conv_id)
            return messages[-limit:] if messages else []
        except Exception as e:
            logger.warning(f"Could not load user history: {e}")
            return agent.history.get_history(user_context.user.id, channel)

    def save_user_history(
        self,
        user_context: "UserContext",
        channel: str,
        user_message: str,
        assistant_response: str,
    ):
        """Save conversation to user's isolated data store."""
        agent = self._agent
        user_id = user_context.user.id

        agent.history.add_message(user_id, "user", user_message, channel)
        agent.history.add_message(user_id, "assistant", assistant_response, channel)

        if HAS_USER_MANAGEMENT and get_data_store:
            try:
                store = get_data_store(user_id=user_id, user_role=user_context.user.role.value)
                conv_id = f"{channel}_current"
                messages = store.load_history(conv_id)
                messages.append({"role": "user", "content": user_message})
                messages.append({"role": "assistant", "content": assistant_response})
                if len(messages) > 100:
                    messages = messages[-100:]
                store.save_history(conv_id, messages)
            except Exception as e:
                logger.warning(f"Could not save user history: {e}")
