"""
Vector Memory System - Semantic search and persistent memory storage

Provides vector database memory for agents with:
- Semantic similarity search
- Multiple memory areas/namespaces
- AI-powered memory consolidation
- Persistent storage via SQLite and FAISS

Performance Features (PERF-011):
- LRU cache with configurable max size
- Embedding cache for repeated queries
"""

import hashlib
import json
import sqlite3
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable, Optional

import numpy as np


# PERF-011: LRU Cache configuration
DEFAULT_CACHE_MAX_SIZE = 1000
EMBEDDING_CACHE_MAX_SIZE = 500


class LRUCache:
    """
    Thread-safe LRU cache with configurable max size.

    PERF-011: Prevents unbounded memory growth from caching.
    """

    def __init__(self, max_size: int = DEFAULT_CACHE_MAX_SIZE):
        self.max_size = max_size
        self._cache: OrderedDict = OrderedDict()

    def get(self, key: str) -> Optional[Any]:
        """Get item from cache, moves to end (most recently used)."""
        if key in self._cache:
            self._cache.move_to_end(key)
            return self._cache[key]
        return None

    def set(self, key: str, value: Any) -> None:
        """Set item in cache, evicts oldest if at capacity."""
        if key in self._cache:
            self._cache.move_to_end(key)
        else:
            if len(self._cache) >= self.max_size:
                self._cache.popitem(last=False)  # Remove oldest
        self._cache[key] = value

    def delete(self, key: str) -> bool:
        """Delete item from cache."""
        if key in self._cache:
            del self._cache[key]
            return True
        return False

    def clear(self) -> None:
        """Clear all items from cache."""
        self._cache.clear()

    def __contains__(self, key: str) -> bool:
        return key in self._cache

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


class MemoryArea(Enum):
    """Memory storage areas for different types of content."""

    MAIN = "main"  # General memories
    FRAGMENTS = "fragments"  # Code/text fragments
    SOLUTIONS = "solutions"  # Problem solutions
    INSTRUMENTS = "instruments"  # Reusable procedures
    OSINT = "osint"  # OSINT findings
    CONVERSATIONS = "conversations"  # Conversation history


@dataclass
class MemoryItem:
    """A single memory item."""

    id: str
    content: str
    area: MemoryArea
    embedding: Optional[list[float]] = None
    metadata: dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.utcnow)
    accessed_at: Optional[datetime] = None
    access_count: int = 0
    relevance_score: float = 0.0

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "id": self.id,
            "content": self.content,
            "area": self.area.value,
            "metadata": self.metadata,
            "created_at": self.created_at.isoformat(),
            "accessed_at": self.accessed_at.isoformat() if self.accessed_at else None,
            "access_count": self.access_count,
            "relevance_score": self.relevance_score,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "MemoryItem":
        """Create from dictionary."""
        return cls(
            id=data["id"],
            content=data["content"],
            area=MemoryArea(data["area"]),
            metadata=data.get("metadata", {}),
            created_at=datetime.fromisoformat(data["created_at"]),
            accessed_at=datetime.fromisoformat(data["accessed_at"]) if data.get("accessed_at") else None,
            access_count=data.get("access_count", 0),
            relevance_score=data.get("relevance_score", 0.0),
        )


class EmbeddingProvider:
    """
    Base class for embedding providers.

    Override get_embedding() to use different embedding models.

    PERF-011: Includes embedding cache to avoid recomputing.
    """

    def __init__(self, dimension: int = 384, cache_size: int = EMBEDDING_CACHE_MAX_SIZE):
        self.dimension = dimension
        # PERF-011: Cache for embeddings
        self._embedding_cache = LRUCache(max_size=cache_size)

    def _compute_embedding(self, text: str) -> list[float]:
        """
        Compute embedding vector for text (without caching).

        Default implementation uses simple hash-based embeddings.
        Override for real embeddings (OpenAI, sentence-transformers, etc.)
        """
        # Simple hash-based embedding for fallback
        # In production, use real embedding model
        hash_bytes = hashlib.sha256(text.encode()).digest()
        embedding = []
        for i in range(0, min(len(hash_bytes), self.dimension * 4), 4):
            value = int.from_bytes(hash_bytes[i : i + 4], "big")
            embedding.append((value / (2**32)) * 2 - 1)  # Normalize to [-1, 1]

        # Pad or truncate to dimension
        if len(embedding) < self.dimension:
            embedding.extend([0.0] * (self.dimension - len(embedding)))
        embedding = embedding[: self.dimension]

        # Normalize
        norm = np.linalg.norm(embedding)
        if norm > 0:
            embedding = [e / norm for e in embedding]

        return embedding

    def get_embedding(self, text: str) -> list[float]:
        """
        Get embedding vector for text with caching.

        PERF-011: Uses LRU cache to avoid recomputing embeddings.
        """
        # Create cache key from text hash
        cache_key = hashlib.md5(text.encode()).hexdigest()

        # Check cache
        cached = self._embedding_cache.get(cache_key)
        if cached is not None:
            return cached

        # Compute and cache
        embedding = self._compute_embedding(text)
        self._embedding_cache.set(cache_key, embedding)

        return embedding

    async def get_embedding_async(self, text: str) -> list[float]:
        """Async version of get_embedding."""
        return self.get_embedding(text)

    def clear_cache(self) -> None:
        """Clear the embedding cache."""
        self._embedding_cache.clear()


class OpenAIEmbedding(EmbeddingProvider):
    """OpenAI embedding provider."""

    def __init__(self, model: str = "text-embedding-3-small"):
        super().__init__(dimension=1536)
        self.model = model
        self._client = None

    def _get_client(self):
        if self._client is None:
            try:
                from openai import OpenAI

                self._client = OpenAI()
            except ImportError:
                raise ImportError("openai package required for OpenAI embeddings")
        return self._client

    def get_embedding(self, text: str) -> list[float]:
        """Get embedding from OpenAI."""
        client = self._get_client()
        response = client.embeddings.create(
            input=text,
            model=self.model,
        )
        return response.data[0].embedding


class VectorIndex:
    """
    Simple vector index for similarity search.

    Uses numpy for basic operations. For production,
    consider FAISS, Annoy, or other optimized libraries.
    """

    def __init__(self, dimension: int = 384):
        self.dimension = dimension
        self._vectors: list[np.ndarray] = []
        self._ids: list[str] = []

    def add(self, id: str, vector: list[float]) -> None:
        """Add a vector to the index."""
        vec = np.array(vector, dtype=np.float32)
        if id in self._ids:
            # Update existing
            idx = self._ids.index(id)
            self._vectors[idx] = vec
        else:
            self._ids.append(id)
            self._vectors.append(vec)

    def remove(self, id: str) -> bool:
        """Remove a vector from the index."""
        if id in self._ids:
            idx = self._ids.index(id)
            del self._ids[idx]
            del self._vectors[idx]
            return True
        return False

    def search(
        self,
        query_vector: list[float],
        top_k: int = 10,
        threshold: float = 0.5,
    ) -> list[tuple[str, float]]:
        """
        Search for similar vectors.

        Args:
            query_vector: Query embedding
            top_k: Maximum results
            threshold: Minimum similarity score

        Returns:
            List of (id, similarity) tuples
        """
        if not self._vectors:
            return []

        query = np.array(query_vector, dtype=np.float32)

        # Calculate cosine similarities
        similarities = []
        for i, vec in enumerate(self._vectors):
            sim = np.dot(query, vec) / (np.linalg.norm(query) * np.linalg.norm(vec) + 1e-8)
            if sim >= threshold:
                similarities.append((self._ids[i], float(sim)))

        # Sort by similarity
        similarities.sort(key=lambda x: x[1], reverse=True)
        return similarities[:top_k]

    def save(self, path: Path) -> None:
        """Save index to file."""
        data = {
            "dimension": self.dimension,
            "ids": self._ids,
            "vectors": [v.tolist() for v in self._vectors],
        }
        with open(path, "w") as f:
            json.dump(data, f)

    def load(self, path: Path) -> None:
        """Load index from file."""
        if not path.exists():
            return

        with open(path) as f:
            data = json.load(f)

        self.dimension = data["dimension"]
        self._ids = data["ids"]
        self._vectors = [np.array(v, dtype=np.float32) for v in data["vectors"]]

    @property
    def size(self) -> int:
        """Get number of vectors in index."""
        return len(self._ids)


class VectorMemory:
    """
    Vector-based memory system for agents.

    Features:
    - Semantic similarity search
    - Multiple memory areas
    - SQLite persistence
    - Vector index for fast search
    - Memory consolidation
    """

    def __init__(
        self,
        storage_path: Path,
        embedding_provider: Optional[EmbeddingProvider] = None,
        default_area: MemoryArea = MemoryArea.MAIN,
    ):
        """
        Initialize vector memory.

        Args:
            storage_path: Directory for persistent storage
            embedding_provider: Provider for text embeddings
            default_area: Default memory area
        """
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True)

        self.embedding = embedding_provider or EmbeddingProvider()
        self.default_area = default_area

        # Initialize storage
        self._db_path = self.storage_path / "memory.db"
        self._index_path = self.storage_path / "vectors.json"
        self._init_database()

        # Load vector index
        self._index = VectorIndex(dimension=self.embedding.dimension)
        self._index.load(self._index_path)

        # PERF-011: LRU Memory cache with bounded size
        self._cache = LRUCache(max_size=DEFAULT_CACHE_MAX_SIZE)

    def _init_database(self) -> None:
        """Initialize SQLite database."""
        conn = sqlite3.connect(self._db_path)
        cursor = conn.cursor()

        cursor.execute("""
            CREATE TABLE IF NOT EXISTS memories (
                id TEXT PRIMARY KEY,
                content TEXT NOT NULL,
                area TEXT NOT NULL,
                metadata TEXT,
                embedding BLOB,
                created_at TEXT NOT NULL,
                accessed_at TEXT,
                access_count INTEGER DEFAULT 0
            )
        """)

        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_area ON memories(area)
        """)

        conn.commit()
        conn.close()

    def save(
        self,
        content: str,
        area: Optional[MemoryArea] = None,
        metadata: Optional[dict[str, Any]] = None,
        id: Optional[str] = None,
    ) -> MemoryItem:
        """
        Save a memory.

        Args:
            content: Memory content
            area: Memory area
            metadata: Additional metadata
            id: Optional custom ID

        Returns:
            Saved MemoryItem
        """
        area = area or self.default_area
        metadata = metadata or {}

        # Generate ID if not provided
        if not id:
            id = hashlib.md5(f"{content}{time.time()}".encode()).hexdigest()[:16]

        # Get embedding
        embedding = self.embedding.get_embedding(content)

        # Create memory item
        item = MemoryItem(
            id=id,
            content=content,
            area=area,
            embedding=embedding,
            metadata=metadata,
        )

        # Save to database
        conn = sqlite3.connect(self._db_path)
        cursor = conn.cursor()

        cursor.execute(
            """
            INSERT OR REPLACE INTO memories
            (id, content, area, metadata, embedding, created_at, accessed_at, access_count)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                item.id,
                item.content,
                item.area.value,
                json.dumps(item.metadata),
                json.dumps(embedding),
                item.created_at.isoformat(),
                None,
                0,
            ),
        )

        conn.commit()
        conn.close()

        # Add to vector index
        self._index.add(id, embedding)
        self._index.save(self._index_path)

        # PERF-011: Cache using LRU
        self._cache.set(id, item)

        return item

    def search(
        self,
        query: str,
        area: Optional[MemoryArea] = None,
        limit: int = 10,
        threshold: float = 0.5,
    ) -> list[MemoryItem]:
        """
        Search for similar memories.

        Args:
            query: Search query
            area: Optional area filter
            limit: Maximum results
            threshold: Minimum similarity

        Returns:
            List of matching MemoryItems
        """
        # Get query embedding
        query_embedding = self.embedding.get_embedding(query)

        # Search vector index
        results = self._index.search(query_embedding, top_k=limit * 2, threshold=threshold)

        # Load memories and filter by area
        memories = []
        for id, score in results:
            item = self.load(id)
            if item:
                if area is None or item.area == area:
                    item.relevance_score = score
                    memories.append(item)

                    # Update access stats
                    self._update_access(id)

        return memories[:limit]

    def load(self, id: str) -> Optional[MemoryItem]:
        """
        Load a memory by ID.

        PERF-011: Uses LRU cache for frequently accessed items.

        Args:
            id: Memory ID

        Returns:
            MemoryItem or None
        """
        # PERF-011: Check LRU cache first
        cached = self._cache.get(id)
        if cached is not None:
            return cached

        # Load from database
        conn = sqlite3.connect(self._db_path)
        cursor = conn.cursor()

        cursor.execute("SELECT * FROM memories WHERE id = ?", (id,))
        row = cursor.fetchone()
        conn.close()

        if not row:
            return None

        item = MemoryItem(
            id=row[0],
            content=row[1],
            area=MemoryArea(row[2]),
            metadata=json.loads(row[3]) if row[3] else {},
            embedding=json.loads(row[4]) if row[4] else None,
            created_at=datetime.fromisoformat(row[5]),
            accessed_at=datetime.fromisoformat(row[6]) if row[6] else None,
            access_count=row[7],
        )

        self._cache.set(id, item)
        return item

    def _update_access(self, id: str) -> None:
        """Update access statistics for a memory."""
        conn = sqlite3.connect(self._db_path)
        cursor = conn.cursor()

        cursor.execute(
            """
            UPDATE memories
            SET accessed_at = ?, access_count = access_count + 1
            WHERE id = ?
            """,
            (datetime.utcnow().isoformat(), id),
        )

        conn.commit()
        conn.close()

        # PERF-011: Update LRU cache
        cached = self._cache.get(id)
        if cached is not None:
            cached.accessed_at = datetime.utcnow()
            cached.access_count += 1
            self._cache.set(id, cached)

    def delete(self, id: str) -> bool:
        """
        Delete a memory.

        Args:
            id: Memory ID

        Returns:
            True if deleted
        """
        conn = sqlite3.connect(self._db_path)
        cursor = conn.cursor()

        cursor.execute("DELETE FROM memories WHERE id = ?", (id,))
        deleted = cursor.rowcount > 0

        conn.commit()
        conn.close()

        if deleted:
            self._index.remove(id)
            self._index.save(self._index_path)
            self._cache.delete(id)  # PERF-011: Use LRU cache delete

        return deleted

    def list_area(self, area: MemoryArea, limit: int = 100) -> list[MemoryItem]:
        """
        List all memories in an area.

        Args:
            area: Memory area
            limit: Maximum results

        Returns:
            List of MemoryItems
        """
        conn = sqlite3.connect(self._db_path)
        cursor = conn.cursor()

        cursor.execute(
            "SELECT id FROM memories WHERE area = ? ORDER BY created_at DESC LIMIT ?",
            (area.value, limit),
        )

        ids = [row[0] for row in cursor.fetchall()]
        conn.close()

        return [self.load(id) for id in ids if self.load(id)]

    def consolidate(
        self,
        area: Optional[MemoryArea] = None,
        similarity_threshold: float = 0.9,
        consolidate_fn: Optional[Callable[[list[str]], str]] = None,
    ) -> int:
        """
        Consolidate similar memories.

        Args:
            area: Optional area to consolidate
            similarity_threshold: Threshold for considering memories similar
            consolidate_fn: Function to merge memory contents

        Returns:
            Number of memories consolidated
        """
        # Default consolidation: keep the longest
        if not consolidate_fn:
            consolidate_fn = lambda contents: max(contents, key=len)

        # Get all memories in area
        if area:
            memories = self.list_area(area, limit=1000)
        else:
            memories = []
            for a in MemoryArea:
                memories.extend(self.list_area(a, limit=1000))

        # Find similar groups
        groups: list[list[MemoryItem]] = []
        processed = set()

        for mem in memories:
            if mem.id in processed:
                continue

            # Find similar memories
            similar = self.search(
                mem.content,
                area=area,
                limit=20,
                threshold=similarity_threshold,
            )

            if len(similar) > 1:
                group = [m for m in similar if m.id not in processed]
                if len(group) > 1:
                    groups.append(group)
                    for m in group:
                        processed.add(m.id)
            else:
                processed.add(mem.id)

        # Consolidate groups
        consolidated = 0
        for group in groups:
            # Merge contents
            contents = [m.content for m in group]
            merged_content = consolidate_fn(contents)

            # Merge metadata
            merged_metadata = {}
            for m in group:
                merged_metadata.update(m.metadata)
            merged_metadata["consolidated_from"] = [m.id for m in group]

            # Delete old memories
            for m in group:
                self.delete(m.id)

            # Save merged memory
            self.save(
                content=merged_content,
                area=group[0].area,
                metadata=merged_metadata,
            )

            consolidated += len(group) - 1

        return consolidated

    def get_stats(self) -> dict[str, Any]:
        """Get memory statistics."""
        conn = sqlite3.connect(self._db_path)
        cursor = conn.cursor()

        stats = {
            "total_memories": 0,
            "by_area": {},
            "total_accesses": 0,
        }

        for area in MemoryArea:
            cursor.execute(
                "SELECT COUNT(*), SUM(access_count) FROM memories WHERE area = ?",
                (area.value,),
            )
            row = cursor.fetchone()
            count = row[0] or 0
            accesses = row[1] or 0

            stats["by_area"][area.value] = {
                "count": count,
                "accesses": accesses,
            }
            stats["total_memories"] += count
            stats["total_accesses"] += accesses

        conn.close()

        stats["vector_index_size"] = self._index.size
        stats["cache_size"] = len(self._cache)

        return stats

    def clear(self, area: Optional[MemoryArea] = None) -> int:
        """
        Clear memories.

        Args:
            area: Optional area to clear (clears all if None)

        Returns:
            Number of memories deleted
        """
        conn = sqlite3.connect(self._db_path)
        cursor = conn.cursor()

        if area:
            # Get IDs to delete
            cursor.execute("SELECT id FROM memories WHERE area = ?", (area.value,))
            ids = [row[0] for row in cursor.fetchall()]

            cursor.execute("DELETE FROM memories WHERE area = ?", (area.value,))
        else:
            cursor.execute("SELECT id FROM memories")
            ids = [row[0] for row in cursor.fetchall()]

            cursor.execute("DELETE FROM memories")

        deleted = cursor.rowcount
        conn.commit()
        conn.close()

        # Update index and cache
        for id in ids:
            self._index.remove(id)
            self._cache.delete(id)

        self._index.save(self._index_path)

        return deleted


# Global memory instance
_memory: Optional[VectorMemory] = None


def get_memory(storage_path: Optional[Path] = None) -> VectorMemory:
    """
    Get the global memory instance.

    Args:
        storage_path: Optional storage path for initialization

    Returns:
        VectorMemory instance
    """
    global _memory

    if _memory is None:
        path = storage_path or Path.home() / ".groknroll" / "memory"
        _memory = VectorMemory(path)

    return _memory
