"""Embedding Index - Semantic Search Across All Content.

Uses Rust-accelerated embeddings for fast similarity search.
Enables finding related memories even when exact keywords don't match.

Usage:
    from whitemagic.core.memory.memory_matrix import get_embedding_index

    index = get_embedding_index()
    index.add("memory_id", "This is the content to embed")

    similar = index.search("related content query", limit=10)
"""

from __future__ import annotations

import hashlib
from dataclasses import dataclass, field

from whitemagic.utils.fast_json import dumps_str as _json_dumps, loads as _json_loads
from datetime import datetime, timezone
from importlib.util import find_spec
from pathlib import Path
from threading import Lock
from typing import Any

try:
    import numpy as np
except ImportError:
    np = None  # type: ignore[assignment]

from whitemagic.config.paths import WM_ROOT
from whitemagic.utils.fileio import atomic_write, file_lock

# Try to import Rust acceleration
RUST_AVAILABLE = find_spec("whitemagic_rs") is not None

# Singleton
_index_instance: "EmbeddingIndex" | None = None
_index_lock = Lock()


@dataclass
class EmbeddingEntry:
    """An embedded document."""

    id: str
    content_preview: str  # First 200 chars
    content_hash: str
    embedding: list[float]
    metadata: dict[str, Any] = field(default_factory=dict)
    created: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


class EmbeddingIndex:
    """Semantic search index using embeddings.

    Features:
    - Fast similarity search
    - Rust-accelerated when available
    - Persistent storage
    - Automatic deduplication via content hashing
    """

    def __init__(self, storage_path: Path | None = None):
        self.storage_path = storage_path or (WM_ROOT / "matrix" / "embeddings")
        self.storage_path.mkdir(parents=True, exist_ok=True)
        self._entries: dict[str, EmbeddingEntry] = {}
        self._embeddings_matrix: np.ndarray | None = None
        self._id_list: list[str] = []
        self._load()

    def _load(self) -> None:
        """Load index from disk."""
        index_file = self.storage_path / "index.json"
        if index_file.exists():
            try:
                with file_lock(index_file):
                    data = _json_loads(index_file.read_text())
                for entry_data in data.get("entries", []):
                    entry = EmbeddingEntry(
                        id=entry_data["id"],
                        content_preview=entry_data["content_preview"],
                        content_hash=entry_data["content_hash"],
                        embedding=entry_data["embedding"],
                        metadata=entry_data.get("metadata", {}),
                        created=entry_data.get("created", ""),
                    )
                    self._entries[entry.id] = entry
                self._rebuild_matrix()
            except Exception:
                self._entries = {}

    def _save(self) -> None:
        """Save index to disk."""
        index_file = self.storage_path / "index.json"
        data = {
            "version": "1.0",
            "updated": datetime.now(timezone.utc).isoformat(),
            "total_entries": len(self._entries),
            "rust_available": RUST_AVAILABLE,
            "entries": [
                {
                    "id": e.id,
                    "content_preview": e.content_preview,
                    "content_hash": e.content_hash,
                    "embedding": e.embedding,
                    "metadata": e.metadata,
                    "created": e.created,
                }
                for e in self._entries.values()
            ],
        }
        with file_lock(index_file):
            atomic_write(index_file, _json_dumps(data))

    def _rebuild_matrix(self) -> None:
        """Rebuild numpy matrix for fast search."""
        if not self._entries:
            self._embeddings_matrix = None
            self._id_list = []
            return

        self._id_list = list(self._entries.keys())
        embeddings = [self._entries[id].embedding for id in self._id_list]
        self._embeddings_matrix = np.array(embeddings, dtype=np.float32)

    def _hash_content(self, content: str) -> str:
        """Generate content hash."""
        return hashlib.sha256(content.encode()).hexdigest()[:16]

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

        Uses simple TF-IDF style embedding if no model available.
        In production, would use sentence-transformers or similar.
        """
        # Simple bag-of-words style embedding (768 dimensions to match common models)
        # This is a placeholder - in production use sentence-transformers
        words = text.lower().split()
        embedding = [0.0] * 768

        for i, word in enumerate(words[:768]):
            # Simple hash-based embedding
            h = hash(word) % 768
            embedding[h] += 1.0 / (i + 1)  # Position-weighted

        # Normalize
        norm = sum(x*x for x in embedding) ** 0.5
        if norm > 0:
            embedding = [x / norm for x in embedding]

        return embedding

    def add(
        self,
        id: str,
        content: str,
        metadata: dict[str, Any] | None = None,
    ) -> bool:
        """Add content to the index.

        Args:
            id: Unique identifier (e.g., file path, memory ID)
            content: The text content to embed
            metadata: Optional metadata (tags, source, etc.)

        Returns:
            True if added, False if duplicate content

        """
        content_hash = self._hash_content(content)

        # Check for duplicate content
        for entry in self._entries.values():
            if entry.content_hash == content_hash:
                return False

        embedding = self._compute_embedding(content)

        entry = EmbeddingEntry(
            id=id,
            content_preview=content[:200],
            content_hash=content_hash,
            embedding=embedding,
            metadata=metadata or {},
        )

        self._entries[id] = entry
        self._rebuild_matrix()
        self._save()
        return True

    def search(
        self,
        query: str,
        limit: int = 10,
        min_score: float = 0.0,
    ) -> list[tuple[str, float, str]]:
        """Search for similar content.

        Args:
            query: Search query text
            limit: Maximum results to return
            min_score: Minimum similarity score (0-1)

        Returns:
            List of (id, score, preview) tuples

        """
        if self._embeddings_matrix is None or len(self._id_list) == 0:
            return []

        query_embedding = np.array(self._compute_embedding(query), dtype=np.float32)

        # Cosine similarity
        similarities = np.dot(self._embeddings_matrix, query_embedding)

        # Get top results
        indices = np.argsort(similarities)[::-1][:limit]

        results = []
        for idx in indices:
            score = float(similarities[idx])
            if score >= min_score:
                id = self._id_list[idx]
                preview = self._entries[id].content_preview
                results.append((id, score, preview))

        return results

    def get(self, id: str) -> EmbeddingEntry | None:
        """Get entry by ID."""
        return self._entries.get(id)

    def remove(self, id: str) -> bool:
        """Remove entry from index."""
        if id in self._entries:
            del self._entries[id]
            self._rebuild_matrix()
            self._save()
            return True
        return False

    def stats(self) -> dict[str, Any]:
        """Get index statistics."""
        return {
            "total_entries": len(self._entries),
            "embedding_dimensions": 768,
            "rust_available": RUST_AVAILABLE,
            "storage_path": str(self.storage_path),
        }

    def clear(self) -> None:
        """Clear all entries."""
        self._entries = {}
        self._embeddings_matrix = None
        self._id_list = []
        self._save()


def get_embedding_index() -> EmbeddingIndex:
    """Get singleton EmbeddingIndex instance."""
    global _index_instance
    with _index_lock:
        if _index_instance is None:
            _index_instance = EmbeddingIndex()
        return _index_instance
