# 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

"""
Long-term Memory (v1.4.0 Phase 4 - Enhancement #18)

Episodic memory system for LLM agents with:
- Memory encoding and storage
- Relevance-based retrieval
- Memory consolidation and decay
- Importance scoring
- Semantic clustering
- Working memory management
- Persistence support

Usage:
    from familiar.core import MemorySystem, Memory, MemoryType

    memory = MemorySystem()

    # Store memories
    memory.store(Memory(
        content="User prefers Python over JavaScript",
        memory_type=MemoryType.PREFERENCE,
        importance=0.8
    ))

    # Retrieve relevant memories
    relevant = memory.retrieve("What programming language?", top_k=5)

    # Consolidate and decay over time
    memory.consolidate()
"""

import json
import logging
import math
import os
import re
import tempfile
import threading
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple

logger = logging.getLogger(__name__)


# ============================================================
# CONSTANTS
# ============================================================


def _utcnow() -> datetime:
    """Get current UTC time as timezone-aware datetime."""
    return datetime.now(timezone.utc)


# Default decay parameters
DEFAULT_DECAY_RATE = 0.1  # Per day
DEFAULT_MIN_IMPORTANCE = 0.1
DEFAULT_CONSOLIDATION_THRESHOLD = 0.3

# Memory limits
DEFAULT_WORKING_MEMORY_SIZE = 10
DEFAULT_MAX_MEMORIES = 10000


# ============================================================
# ENUMS
# ============================================================


class MemoryType(str, Enum):
    """Types of memories."""

    # Episodic - specific events/interactions
    EPISODIC = "episodic"
    CONVERSATION = "conversation"
    EVENT = "event"

    # Semantic - facts and knowledge
    SEMANTIC = "semantic"
    FACT = "fact"
    PREFERENCE = "preference"

    # Procedural - how to do things
    PROCEDURAL = "procedural"
    SKILL = "skill"

    # Meta - about the memory system itself
    META = "meta"
    SUMMARY = "summary"


class MemoryStatus(str, Enum):
    """Status of a memory."""

    ACTIVE = "active"
    CONSOLIDATED = "consolidated"
    ARCHIVED = "archived"
    FORGOTTEN = "forgotten"


class RetrievalStrategy(str, Enum):
    """Strategy for memory retrieval."""

    RECENCY = "recency"  # Most recent first
    IMPORTANCE = "importance"  # Most important first
    RELEVANCE = "relevance"  # Most relevant to query
    COMBINED = "combined"  # Weighted combination
    ASSOCIATIVE = "associative"  # Follow associations


# ============================================================
# MEMORY DATACLASS
# ============================================================


@dataclass
class Memory:
    """
    A single memory unit.

    Attributes:
        id: Unique identifier
        content: The memory content (text)
        memory_type: Type of memory

        # Importance and decay
        importance: Base importance (0-1)
        current_strength: Current strength after decay

        # Timing
        created_at: When the memory was created
        last_accessed: When last retrieved
        access_count: Number of times retrieved

        # Associations
        tags: Categorical tags
        associations: IDs of related memories
        source: Where the memory came from

        # Metadata
        metadata: Additional context
        embedding: Optional vector embedding
    """

    # Core
    content: str
    memory_type: MemoryType = MemoryType.EPISODIC
    id: str = field(default_factory=lambda: str(uuid.uuid4()))

    # Importance
    importance: float = 0.5
    current_strength: float = field(default=None)

    # Timing
    created_at: datetime = field(default_factory=_utcnow)
    last_accessed: datetime = field(default_factory=_utcnow)
    access_count: int = 0

    # Status
    status: MemoryStatus = MemoryStatus.ACTIVE

    # Associations
    tags: Set[str] = field(default_factory=set)
    associations: Set[str] = field(default_factory=set)
    source: Optional[str] = None

    # Context
    context: Optional[str] = None
    conversation_id: Optional[str] = None

    # Metadata
    metadata: Dict[str, Any] = field(default_factory=dict)
    embedding: Optional[List[float]] = None

    def __post_init__(self):
        """Initialize computed fields."""
        if self.current_strength is None:
            self.current_strength = self.importance
        # Convert tags/associations to sets if needed
        if isinstance(self.tags, list):
            self.tags = set(self.tags)
        if isinstance(self.associations, list):
            self.associations = set(self.associations)

    def decay(self, decay_rate: float = DEFAULT_DECAY_RATE) -> float:
        """
        Apply time-based decay to the memory.

        Args:
            decay_rate: Decay rate per day

        Returns:
            New strength after decay
        """
        days_since_access = (_utcnow() - self.last_accessed).total_seconds() / 86400

        # Exponential decay
        decay_factor = math.exp(-decay_rate * days_since_access)
        self.current_strength = self.importance * decay_factor

        return self.current_strength

    def reinforce(self, boost: float = 0.1):
        """
        Reinforce the memory (on access).

        Args:
            boost: Amount to boost importance
        """
        self.last_accessed = _utcnow()
        self.access_count += 1

        # Boost importance (with diminishing returns)
        boost_factor = boost / (1 + self.access_count * 0.1)
        self.importance = min(1.0, self.importance + boost_factor)
        self.current_strength = self.importance

    def add_association(self, memory_id: str):
        """Add an association to another memory."""
        self.associations.add(memory_id)

    def add_tag(self, tag: str):
        """Add a tag to this memory."""
        self.tags.add(tag.lower())

    @property
    def age_days(self) -> float:
        """Get age in days."""
        return (_utcnow() - self.created_at).total_seconds() / 86400

    @property
    def recency_score(self) -> float:
        """Get recency score (0-1, higher is more recent)."""
        days_since_access = (_utcnow() - self.last_accessed).total_seconds() / 86400
        return math.exp(-0.1 * days_since_access)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "id": self.id,
            "content": self.content,
            "memory_type": self.memory_type.value,
            "importance": self.importance,
            "current_strength": self.current_strength,
            "created_at": self.created_at.isoformat(),
            "last_accessed": self.last_accessed.isoformat(),
            "access_count": self.access_count,
            "status": self.status.value,
            "tags": list(self.tags),
            "associations": list(self.associations),
            "source": self.source,
            "context": self.context,
            "conversation_id": self.conversation_id,
            "metadata": self.metadata,
            "embedding": self.embedding,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Memory":
        """Create from dictionary."""
        data = data.copy()
        data["memory_type"] = MemoryType(data["memory_type"])
        data["status"] = MemoryStatus(data["status"])
        data["created_at"] = datetime.fromisoformat(data["created_at"])
        data["last_accessed"] = datetime.fromisoformat(data["last_accessed"])
        data["tags"] = set(data.get("tags", []))
        data["associations"] = set(data.get("associations", []))
        return cls(**data)

    def __hash__(self):
        return hash(self.id)

    def __eq__(self, other):
        if isinstance(other, Memory):
            return self.id == other.id
        return False


# ============================================================
# MEMORY INDEX
# ============================================================


class MemoryIndex:
    """
    Index for efficient memory retrieval.

    Supports:
    - Tag-based lookup
    - Type-based lookup
    - Time-range queries
    - Text search (basic)
    """

    def __init__(self):
        self._by_id: Dict[str, Memory] = {}
        self._by_tag: Dict[str, Set[str]] = {}  # tag -> memory IDs
        self._by_type: Dict[MemoryType, Set[str]] = {}  # type -> memory IDs
        self._by_conversation: Dict[str, Set[str]] = {}  # conv_id -> memory IDs
        self._lock = threading.RLock()

    def add(self, memory: Memory):
        """Add a memory to the index."""
        with self._lock:
            self._by_id[memory.id] = memory

            # Index by tags
            for tag in memory.tags:
                if tag not in self._by_tag:
                    self._by_tag[tag] = set()
                self._by_tag[tag].add(memory.id)

            # Index by type
            if memory.memory_type not in self._by_type:
                self._by_type[memory.memory_type] = set()
            self._by_type[memory.memory_type].add(memory.id)

            # Index by conversation
            if memory.conversation_id:
                if memory.conversation_id not in self._by_conversation:
                    self._by_conversation[memory.conversation_id] = set()
                self._by_conversation[memory.conversation_id].add(memory.id)

    def remove(self, memory_id: str):
        """Remove a memory from the index."""
        with self._lock:
            memory = self._by_id.get(memory_id)
            if not memory:
                return

            del self._by_id[memory_id]

            # Remove from tag index
            for tag in memory.tags:
                if tag in self._by_tag:
                    self._by_tag[tag].discard(memory_id)

            # Remove from type index
            if memory.memory_type in self._by_type:
                self._by_type[memory.memory_type].discard(memory_id)

            # Remove from conversation index
            if memory.conversation_id and memory.conversation_id in self._by_conversation:
                self._by_conversation[memory.conversation_id].discard(memory_id)

    def get(self, memory_id: str) -> Optional[Memory]:
        """Get a memory by ID."""
        with self._lock:
            return self._by_id.get(memory_id)

    def get_by_tag(self, tag: str) -> List[Memory]:
        """Get memories by tag."""
        with self._lock:
            ids = self._by_tag.get(tag.lower(), set())
            return [self._by_id[mid] for mid in ids if mid in self._by_id]

    def get_by_type(self, memory_type: MemoryType) -> List[Memory]:
        """Get memories by type."""
        with self._lock:
            ids = self._by_type.get(memory_type, set())
            return [self._by_id[mid] for mid in ids if mid in self._by_id]

    def get_by_conversation(self, conversation_id: str) -> List[Memory]:
        """Get memories by conversation ID."""
        with self._lock:
            ids = self._by_conversation.get(conversation_id, set())
            return [self._by_id[mid] for mid in ids if mid in self._by_id]

    def get_all(self) -> List[Memory]:
        """Get all memories."""
        with self._lock:
            return list(self._by_id.values())

    def search_text(self, query: str, limit: int = 100) -> List[Memory]:
        """
        Basic text search in memory content.

        Args:
            query: Search query
            limit: Maximum results

        Returns:
            Matching memories
        """
        query_lower = query.lower()
        query_words = set(query_lower.split())

        results = []
        with self._lock:
            for memory in self._by_id.values():
                content_lower = memory.content.lower()

                # Check for word matches
                content_words = set(content_lower.split())
                matches = len(query_words & content_words)

                if matches > 0 or query_lower in content_lower:
                    results.append((memory, matches))

        # Sort by match count
        results.sort(key=lambda x: x[1], reverse=True)
        return [m for m, _ in results[:limit]]

    def __len__(self) -> int:
        return len(self._by_id)

    def __contains__(self, memory_id: str) -> bool:
        return memory_id in self._by_id


# ============================================================
# SIMILARITY FUNCTIONS
# ============================================================


def cosine_similarity(a: List[float], b: List[float]) -> float:
    """Compute cosine similarity between two vectors."""
    if not a or not b or len(a) != len(b):
        return 0.0

    dot_product = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))

    if norm_a == 0 or norm_b == 0:
        return 0.0

    return dot_product / (norm_a * norm_b)


def jaccard_similarity(a: Set[str], b: Set[str]) -> float:
    """Compute Jaccard similarity between two sets."""
    if not a and not b:
        return 0.0

    intersection = len(a & b)
    union = len(a | b)

    return intersection / union if union > 0 else 0.0


def text_similarity(a: str, b: str) -> float:
    """
    Compute basic text similarity.

    Uses word overlap as a simple metric.
    """
    words_a = set(a.lower().split())
    words_b = set(b.lower().split())
    return jaccard_similarity(words_a, words_b)


# ============================================================
# WORKING MEMORY
# ============================================================


class WorkingMemory:
    """
    Short-term working memory buffer.

    Holds the most relevant memories for current context.
    Limited capacity with automatic eviction.
    """

    def __init__(self, capacity: int = DEFAULT_WORKING_MEMORY_SIZE):
        self.capacity = capacity
        self._memories: List[Memory] = []
        self._lock = threading.Lock()

    def add(self, memory: Memory):
        """Add a memory to working memory."""
        with self._lock:
            # Remove if already present
            self._memories = [m for m in self._memories if m.id != memory.id]

            # Add at front (most recent)
            self._memories.insert(0, memory)

            # Evict if over capacity
            if len(self._memories) > self.capacity:
                self._memories = self._memories[: self.capacity]

    def get_all(self) -> List[Memory]:
        """Get all memories in working memory."""
        with self._lock:
            return list(self._memories)

    def get_context_string(self, max_tokens: int = 2000) -> str:
        """
        Get working memory as a context string.

        Args:
            max_tokens: Approximate token limit

        Returns:
            Formatted context string
        """
        with self._lock:
            lines = []
            total_chars = 0
            char_limit = max_tokens * 4  # Rough estimate

            for memory in self._memories:
                line = f"- [{memory.memory_type.value}] {memory.content}"
                if total_chars + len(line) > char_limit:
                    break
                lines.append(line)
                total_chars += len(line)

            return "\n".join(lines)

    def clear(self):
        """Clear working memory."""
        with self._lock:
            self._memories.clear()

    def __len__(self) -> int:
        return len(self._memories)

    def __iter__(self) -> Iterator[Memory]:
        return iter(self._memories)


# ============================================================
# MEMORY SYSTEM
# ============================================================


class MemorySystem:
    """
    Complete long-term memory system for agents.

    Features:
    - Memory encoding and storage
    - Multiple retrieval strategies
    - Automatic decay and consolidation
    - Working memory management
    - Persistence support

    Usage:
        memory = MemorySystem()

        # Store a memory
        memory.store(Memory(
            content="User's favorite color is blue",
            memory_type=MemoryType.PREFERENCE,
            importance=0.7
        ))

        # Retrieve relevant memories
        relevant = memory.retrieve("What color does the user like?", top_k=5)

        # Get working memory context
        context = memory.working_memory.get_context_string()
    """

    def __init__(
        self,
        embedding_fn: Optional[Callable[[str], List[float]]] = None,
        max_memories: int = DEFAULT_MAX_MEMORIES,
        working_memory_size: int = DEFAULT_WORKING_MEMORY_SIZE,
        decay_rate: float = DEFAULT_DECAY_RATE,
        consolidation_threshold: float = DEFAULT_CONSOLIDATION_THRESHOLD,
    ):
        """
        Initialize memory system.

        Args:
            embedding_fn: Function to generate embeddings (optional)
            max_memories: Maximum number of memories to store
            working_memory_size: Size of working memory buffer
            decay_rate: Daily decay rate for memory strength
            consolidation_threshold: Strength threshold for consolidation
        """
        self.embedding_fn = embedding_fn
        self.max_memories = max_memories
        self.decay_rate = decay_rate
        self.consolidation_threshold = consolidation_threshold

        # Storage
        self._index = MemoryIndex()
        self._working_memory = WorkingMemory(working_memory_size)

        # Statistics
        self._stats = {
            "total_stored": 0,
            "total_retrieved": 0,
            "total_forgotten": 0,
            "total_consolidated": 0,
        }

        self._lock = threading.RLock()

    @property
    def working_memory(self) -> WorkingMemory:
        """Get working memory buffer."""
        return self._working_memory

    def store(
        self,
        memory: Memory,
        add_to_working: bool = True,
        compute_embedding: bool = True,
    ) -> Memory:
        """
        Store a memory.

        Args:
            memory: Memory to store
            add_to_working: Whether to add to working memory
            compute_embedding: Whether to compute embedding

        Returns:
            Stored memory (with computed fields)
        """
        with self._lock:
            # Check working memory for near-duplicates (word-overlap Jaccard)
            for existing in self._working_memory.get_all():
                if text_similarity(memory.content, existing.content) > 0.9:
                    existing.strength = min(1.0, existing.strength + 0.1)
                    logger.debug(f"Near-duplicate memory suppressed, reinforced {existing.id}")
                    return existing

            # Compute embedding if available
            if compute_embedding and self.embedding_fn and not memory.embedding:
                try:
                    memory.embedding = self.embedding_fn(memory.content)
                except Exception as e:
                    logger.warning(f"Failed to compute embedding: {e}")

            # Extract tags from content if none provided
            if not memory.tags:
                memory.tags = self._extract_tags(memory.content)

            # Add to index
            self._index.add(memory)
            self._stats["total_stored"] += 1

            # Add to working memory
            if add_to_working:
                self._working_memory.add(memory)

            # Enforce capacity limit
            self._enforce_capacity()

            logger.debug(f"Stored memory {memory.id}: {memory.content[:50]}...")

            return memory

    def store_conversation(
        self,
        messages: List[Dict[str, str]],
        conversation_id: str,
        importance: float = 0.5,
    ) -> List[Memory]:
        """
        Store a conversation as memories.

        Args:
            messages: List of messages with 'role' and 'content'
            conversation_id: Conversation identifier
            importance: Base importance for memories

        Returns:
            List of created memories
        """
        memories = []

        for i, msg in enumerate(messages):
            role = msg.get("role", "unknown")
            content = msg.get("content", "")

            if not content.strip():
                continue

            memory = Memory(
                content=f"[{role}]: {content}",
                memory_type=MemoryType.CONVERSATION,
                importance=importance,
                conversation_id=conversation_id,
                metadata={"message_index": i, "role": role},
            )

            self.store(memory, add_to_working=False)
            memories.append(memory)

        # Link sequential memories
        for i in range(len(memories) - 1):
            memories[i].add_association(memories[i + 1].id)
            memories[i + 1].add_association(memories[i].id)

        return memories

    def retrieve(
        self,
        query: str,
        top_k: int = 5,
        strategy: RetrievalStrategy = RetrievalStrategy.COMBINED,
        memory_types: Optional[List[MemoryType]] = None,
        tags: Optional[List[str]] = None,
        min_strength: float = 0.0,
        include_associations: bool = True,
    ) -> List[Memory]:
        """
        Retrieve relevant memories.

        Args:
            query: Search query
            top_k: Number of memories to retrieve
            strategy: Retrieval strategy
            memory_types: Filter by memory types
            tags: Filter by tags
            min_strength: Minimum strength threshold
            include_associations: Whether to include associated memories

        Returns:
            List of relevant memories
        """
        with self._lock:
            # Get candidate memories
            candidates = self._get_candidates(memory_types, tags)

            # Apply decay and filter by strength
            for memory in candidates:
                memory.decay(self.decay_rate)

            candidates = [m for m in candidates if m.current_strength >= min_strength]

            if not candidates:
                return []

            # Score and rank
            scored = self._score_memories(query, candidates, strategy)

            # Select top-k
            scored.sort(key=lambda x: x[1], reverse=True)
            results = [m for m, _ in scored[:top_k]]

            # Include associations if requested
            if include_associations:
                results = self._expand_associations(results, top_k)

            # Reinforce retrieved memories
            for memory in results:
                memory.reinforce()

            # Update working memory
            for memory in results:
                self._working_memory.add(memory)

            self._stats["total_retrieved"] += len(results)

            return results

    def retrieve_by_embedding(
        self,
        embedding: List[float],
        top_k: int = 5,
        min_similarity: float = 0.5,
    ) -> List[Tuple[Memory, float]]:
        """
        Retrieve memories by embedding similarity.

        Args:
            embedding: Query embedding
            top_k: Number of results
            min_similarity: Minimum similarity threshold

        Returns:
            List of (memory, similarity) tuples
        """
        results = []

        with self._lock:
            for memory in self._index.get_all():
                if memory.embedding:
                    sim = cosine_similarity(embedding, memory.embedding)
                    if sim >= min_similarity:
                        results.append((memory, sim))

        results.sort(key=lambda x: x[1], reverse=True)
        return results[:top_k]

    def forget(self, memory_id: str) -> bool:
        """
        Forget a memory.

        Args:
            memory_id: ID of memory to forget

        Returns:
            True if forgotten, False if not found
        """
        with self._lock:
            memory = self._index.get(memory_id)
            if not memory:
                return False

            memory.status = MemoryStatus.FORGOTTEN
            self._index.remove(memory_id)
            self._stats["total_forgotten"] += 1

            logger.debug(f"Forgot memory {memory_id}")
            return True

    def consolidate(
        self,
        threshold: Optional[float] = None,
        summarize_fn: Optional[Callable[[List[Memory]], str]] = None,
    ) -> int:
        """
        Consolidate weak memories.

        Groups similar weak memories and creates summaries.

        Args:
            threshold: Strength threshold for consolidation
            summarize_fn: Function to create summaries

        Returns:
            Number of memories consolidated
        """
        threshold = threshold or self.consolidation_threshold
        consolidated_count = 0

        with self._lock:
            # Find weak memories
            weak_memories = []
            for memory in self._index.get_all():
                memory.decay(self.decay_rate)
                if memory.current_strength < threshold:
                    weak_memories.append(memory)

            if not weak_memories:
                return 0

            # Group by type and tags
            groups: Dict[str, List[Memory]] = {}
            for memory in weak_memories:
                key = f"{memory.memory_type.value}:{','.join(sorted(memory.tags))}"
                if key not in groups:
                    groups[key] = []
                groups[key].append(memory)

            # Consolidate each group
            for key, memories in groups.items():
                if len(memories) < 2:
                    continue

                # Create summary
                if summarize_fn:
                    summary_content = summarize_fn(memories)
                else:
                    summary_content = self._create_summary(memories)

                # Create summary memory
                summary = Memory(
                    content=summary_content,
                    memory_type=MemoryType.SUMMARY,
                    importance=max(m.importance for m in memories),
                    tags=set.union(*[m.tags for m in memories]),
                    metadata={
                        "consolidated_from": [m.id for m in memories],
                        "consolidated_at": _utcnow().isoformat(),
                    },
                )

                self.store(summary, add_to_working=False)

                # Mark originals as consolidated
                for memory in memories:
                    memory.status = MemoryStatus.CONSOLIDATED
                    self._index.remove(memory.id)
                    consolidated_count += 1

            self._stats["total_consolidated"] += consolidated_count

        logger.info(f"Consolidated {consolidated_count} memories")
        return consolidated_count

    def get_stats(self) -> Dict[str, Any]:
        """Get memory system statistics."""
        with self._lock:
            type_counts = {}
            for memory_type in MemoryType:
                count = len(self._index.get_by_type(memory_type))
                if count > 0:
                    type_counts[memory_type.value] = count

            return {
                "total_memories": len(self._index),
                "working_memory_size": len(self._working_memory),
                "by_type": type_counts,
                **self._stats,
            }

    def clear(self):
        """Clear all memories."""
        with self._lock:
            self._index = MemoryIndex()
            self._working_memory.clear()
            self._stats = {
                "total_stored": 0,
                "total_retrieved": 0,
                "total_forgotten": 0,
                "total_consolidated": 0,
            }

    def export(self) -> Dict[str, Any]:
        """Export all memories to a dictionary."""
        with self._lock:
            return {
                "memories": [m.to_dict() for m in self._index.get_all()],
                "stats": self._stats.copy(),
                "exported_at": _utcnow().isoformat(),
            }

    def import_memories(self, data: Dict[str, Any]):
        """Import memories from a dictionary."""
        with self._lock:
            for mem_data in data.get("memories", []):
                memory = Memory.from_dict(mem_data)
                self._index.add(memory)

            self._stats = data.get("stats", self._stats)

    # --------------------------------------------------------
    # Private methods
    # --------------------------------------------------------

    def _get_candidates(
        self,
        memory_types: Optional[List[MemoryType]],
        tags: Optional[List[str]],
    ) -> List[Memory]:
        """Get candidate memories for retrieval."""
        candidates = []

        if memory_types:
            for mt in memory_types:
                candidates.extend(self._index.get_by_type(mt))
        elif tags:
            for tag in tags:
                candidates.extend(self._index.get_by_tag(tag))
        else:
            candidates = self._index.get_all()

        # Filter by status
        return [m for m in candidates if m.status == MemoryStatus.ACTIVE]

    def _score_memories(
        self,
        query: str,
        memories: List[Memory],
        strategy: RetrievalStrategy,
    ) -> List[Tuple[Memory, float]]:
        """Score memories based on retrieval strategy."""
        scored = []

        # Compute query embedding if available
        query_embedding = None
        if self.embedding_fn:
            try:
                query_embedding = self.embedding_fn(query)
            except Exception:
                pass

        for memory in memories:
            if strategy == RetrievalStrategy.RECENCY:
                score = memory.recency_score

            elif strategy == RetrievalStrategy.IMPORTANCE:
                score = memory.current_strength

            elif strategy == RetrievalStrategy.RELEVANCE:
                score = self._compute_relevance(query, memory, query_embedding)

            elif strategy == RetrievalStrategy.COMBINED:
                relevance = self._compute_relevance(query, memory, query_embedding)
                recency = memory.recency_score
                strength = memory.current_strength

                # Weighted combination
                score = 0.5 * relevance + 0.3 * strength + 0.2 * recency

            elif strategy == RetrievalStrategy.ASSOCIATIVE:
                # Score based on associations with working memory
                wm_ids = {m.id for m in self._working_memory.get_all()}
                association_score = len(memory.associations & wm_ids) / max(1, len(wm_ids))
                relevance = self._compute_relevance(query, memory, query_embedding)
                score = 0.6 * relevance + 0.4 * association_score

            else:
                score = memory.current_strength

            scored.append((memory, score))

        return scored

    def _compute_relevance(
        self,
        query: str,
        memory: Memory,
        query_embedding: Optional[List[float]] = None,
    ) -> float:
        """Compute relevance score for a memory."""
        # Use embedding similarity if available
        if query_embedding and memory.embedding:
            return cosine_similarity(query_embedding, memory.embedding)

        # Fall back to text similarity
        return text_similarity(query, memory.content)

    def _expand_associations(
        self,
        memories: List[Memory],
        limit: int,
    ) -> List[Memory]:
        """Expand results to include associated memories."""
        result_ids = {m.id for m in memories}
        expanded = list(memories)

        for memory in memories:
            for assoc_id in memory.associations:
                if assoc_id not in result_ids and len(expanded) < limit:
                    assoc_memory = self._index.get(assoc_id)
                    if assoc_memory and assoc_memory.status == MemoryStatus.ACTIVE:
                        expanded.append(assoc_memory)
                        result_ids.add(assoc_id)

        return expanded[:limit]

    def _extract_tags(self, content: str) -> Set[str]:
        """Extract tags from content."""
        tags = set()

        # Extract hashtags
        hashtags = re.findall(r"#(\w+)", content)
        tags.update(tag.lower() for tag in hashtags)

        # Extract common keywords (simple approach)
        keywords = [
            "python",
            "javascript",
            "code",
            "error",
            "help",
            "question",
            "answer",
            "important",
            "remember",
        ]
        content_lower = content.lower()
        for keyword in keywords:
            if keyword in content_lower:
                tags.add(keyword)

        return tags

    def _create_summary(self, memories: List[Memory]) -> str:
        """Create a basic summary of memories."""
        contents = [m.content for m in memories]

        if len(contents) <= 3:
            return " | ".join(contents)

        # Just take first few and note count
        summary = " | ".join(contents[:3])
        return f"{summary} (and {len(contents) - 3} more related items)"

    def _enforce_capacity(self):
        """Enforce maximum memory capacity."""
        while len(self._index) > self.max_memories:
            # Find weakest memory
            weakest = None
            weakest_strength = float("inf")

            for memory in self._index.get_all():
                memory.decay(self.decay_rate)
                if memory.current_strength < weakest_strength:
                    weakest = memory
                    weakest_strength = memory.current_strength

            if weakest:
                self.forget(weakest.id)


# ============================================================
# MEMORY PERSISTENCE
# ============================================================


class MemoryPersistence(ABC):
    """Abstract base for memory persistence backends."""

    @abstractmethod
    def save(self, memory_system: MemorySystem) -> bool:
        """Save memory system state."""
        pass

    @abstractmethod
    def load(self, memory_system: MemorySystem) -> bool:
        """Load memory system state."""
        pass


class FileMemoryPersistence(MemoryPersistence):
    """File-based memory persistence."""

    def __init__(self, filepath: str):
        self.filepath = filepath

    def save(self, memory_system: MemorySystem) -> bool:
        """Save to JSON file atomically via temp file + rename."""
        tmp_path = None
        try:
            data = memory_system.export()
            dir_name = os.path.dirname(self.filepath) or "."
            fd, tmp_path = tempfile.mkstemp(dir=dir_name, suffix=".tmp", prefix=".memories_")
            try:
                with os.fdopen(fd, "w") as f:
                    json.dump(data, f, indent=2, default=str)
            except BaseException:
                os.close(fd)
                raise
            os.replace(tmp_path, self.filepath)
            tmp_path = None  # Rename succeeded, don't unlink
            return True
        except Exception as e:
            logger.error(f"Failed to save memories: {e}")
            if tmp_path:
                try:
                    os.unlink(tmp_path)
                except OSError:
                    pass
            return False

    def load(self, memory_system: MemorySystem) -> bool:
        """Load from JSON file."""
        try:
            with open(self.filepath, "r") as f:
                data = json.load(f)
            memory_system.import_memories(data)
            return True
        except FileNotFoundError:
            logger.info(f"No existing memory file: {self.filepath}")
            return False
        except Exception as e:
            logger.error(f"Failed to load memories: {e}")
            return False


# ============================================================
# CONVENIENCE FUNCTIONS
# ============================================================

# Global memory system
_memory_system: Optional[MemorySystem] = None
_memory_lock = threading.Lock()


def get_memory_system() -> MemorySystem:
    """Get the global MemorySystem instance."""
    global _memory_system
    with _memory_lock:
        if _memory_system is None:
            _memory_system = MemorySystem()
        return _memory_system


def set_memory_system(memory_system: MemorySystem):
    """Set the global MemorySystem instance."""
    global _memory_system
    with _memory_lock:
        _memory_system = memory_system


def reset_memory_system():
    """Reset the global MemorySystem instance."""
    global _memory_system
    with _memory_lock:
        _memory_system = None


def create_memory_system(
    embedding_fn: Optional[Callable[[str], List[float]]] = None,
    max_memories: int = DEFAULT_MAX_MEMORIES,
    **kwargs,
) -> MemorySystem:
    """
    Create a configured MemorySystem.

    Args:
        embedding_fn: Optional embedding function
        max_memories: Maximum memories to store
        **kwargs: Additional MemorySystem options

    Returns:
        Configured MemorySystem
    """
    return MemorySystem(embedding_fn=embedding_fn, max_memories=max_memories, **kwargs)


def store_memory(
    content: str, memory_type: MemoryType = MemoryType.EPISODIC, importance: float = 0.5, **kwargs
) -> Memory:
    """
    Store a memory in the global system.

    Convenience function for quick memory storage.
    """
    memory_system = get_memory_system()
    memory = Memory(content=content, memory_type=memory_type, importance=importance, **kwargs)
    return memory_system.store(memory)


def retrieve_memories(query: str, top_k: int = 5, **kwargs) -> List[Memory]:
    """
    Retrieve memories from the global system.

    Convenience function for quick memory retrieval.
    """
    memory_system = get_memory_system()
    return memory_system.retrieve(query, top_k=top_k, **kwargs)


# ============================================================
# EXPORTS
# ============================================================

__all__ = [
    # Core classes
    "MemorySystem",
    "Memory",
    "MemoryIndex",
    "WorkingMemory",
    # Persistence
    "MemoryPersistence",
    "FileMemoryPersistence",
    # Enums
    "MemoryType",
    "MemoryStatus",
    "RetrievalStrategy",
    # Similarity functions
    "cosine_similarity",
    "jaccard_similarity",
    "text_similarity",
    # Convenience functions
    "get_memory_system",
    "set_memory_system",
    "reset_memory_system",
    "create_memory_system",
    "store_memory",
    "retrieve_memories",
    # Constants
    "DEFAULT_DECAY_RATE",
    "DEFAULT_MIN_IMPORTANCE",
    "DEFAULT_CONSOLIDATION_THRESHOLD",
    "DEFAULT_WORKING_MEMORY_SIZE",
    "DEFAULT_MAX_MEMORIES",
]
