"""
Cross-dataset similarity search using video embeddings.

Provides functionality to:
- Store and index video embeddings
- Find similar videos across datasets using cosine similarity
- Cache embeddings for fast retrieval
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import h5py
import numpy as np

logger = logging.getLogger(__name__)


def compute_cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """
    Compute cosine similarity between two vectors.

    Args:
        a: First vector
        b: Second vector

    Returns:
        Cosine similarity (0 to 1 for non-negative vectors)
    """
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)

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

    return float(np.dot(a, b) / (norm_a * norm_b))


def compute_cosine_distance_matrix(
    queries: np.ndarray,
    database: np.ndarray,
) -> np.ndarray:
    """
    Compute cosine distance matrix between query and database embeddings.

    Args:
        queries: Query embeddings (N, D)
        database: Database embeddings (M, D)

    Returns:
        Distance matrix (N, M) where lower = more similar
    """
    # Normalize
    queries_norm = queries / (np.linalg.norm(queries, axis=1, keepdims=True) + 1e-8)
    database_norm = database / (np.linalg.norm(database, axis=1, keepdims=True) + 1e-8)

    # Cosine similarity
    similarities = queries_norm @ database_norm.T

    # Convert to distance (1 - similarity)
    return 1.0 - similarities


@dataclass
class SimilarityMatch:
    """A single similarity match result."""

    video_id: str
    video_path: str | None
    dataset: str
    gloss: str | None
    similarity: float
    rank: int
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "video_id": self.video_id,
            "video_path": self.video_path,
            "dataset": self.dataset,
            "gloss": self.gloss,
            "similarity": self.similarity,
            "rank": self.rank,
            "metadata": self.metadata,
        }


class EmbeddingStore:
    """
    Storage for video embeddings with metadata.

    Supports:
    - In-memory storage
    - HDF5 file persistence
    - Fast nearest neighbor search
    """

    def __init__(self, store_path: str | Path | None = None):
        """
        Initialize embedding store.

        Args:
            store_path: Optional HDF5 file path for persistence
        """
        self.store_path = Path(store_path) if store_path else None

        # In-memory storage
        self._embeddings: dict[str, np.ndarray] = {}
        self._metadata: dict[str, dict[str, Any]] = {}

        # Cached matrix for fast search
        self._embedding_matrix: np.ndarray | None = None
        self._id_list: list[str] = []

        if self.store_path and self.store_path.exists():
            self._load_from_file()

    def add(
        self,
        video_id: str,
        embedding: np.ndarray,
        dataset: str,
        gloss: str | None = None,
        video_path: str | None = None,
        **metadata,
    ):
        """
        Add an embedding to the store.

        Args:
            video_id: Unique identifier for the video
            embedding: Embedding vector
            dataset: Dataset name
            gloss: Optional gloss/label
            video_path: Optional path to video file
            **metadata: Additional metadata
        """
        self._embeddings[video_id] = np.asarray(embedding)
        self._metadata[video_id] = {
            "dataset": dataset,
            "gloss": gloss,
            "video_path": video_path,
            **metadata,
        }

        # Invalidate cache
        self._embedding_matrix = None
        self._id_list = []

    def get(self, video_id: str) -> tuple[np.ndarray, dict[str, Any]] | None:
        """Get embedding and metadata for a video."""
        if video_id not in self._embeddings:
            return None
        return self._embeddings[video_id], self._metadata[video_id]

    def remove(self, video_id: str) -> bool:
        """Remove a video from the store."""
        if video_id not in self._embeddings:
            return False

        del self._embeddings[video_id]
        del self._metadata[video_id]

        # Invalidate cache
        self._embedding_matrix = None
        self._id_list = []

        return True

    def _build_index(self):
        """Build search index from embeddings."""
        if not self._embeddings:
            self._embedding_matrix = np.array([]).reshape(0, 0)
            self._id_list = []
            return

        self._id_list = list(self._embeddings.keys())
        self._embedding_matrix = np.stack([self._embeddings[vid] for vid in self._id_list])

    def search(
        self,
        query: np.ndarray,
        top_k: int = 10,
        datasets: list[str] | None = None,
        exclude_ids: list[str] | None = None,
    ) -> list[SimilarityMatch]:
        """
        Find most similar videos to query embedding.

        Args:
            query: Query embedding vector
            top_k: Number of results to return
            datasets: Optional list of datasets to search (None = all)
            exclude_ids: Video IDs to exclude from results

        Returns:
            List of SimilarityMatch objects sorted by similarity (descending)
        """
        if self._embedding_matrix is None:
            self._build_index()

        if len(self._id_list) == 0:
            return []

        query = np.asarray(query)
        exclude_set = set(exclude_ids or [])

        # Filter by dataset if specified
        if datasets:
            datasets_set = set(datasets)
            valid_indices = [
                i
                for i, vid in enumerate(self._id_list)
                if self._metadata[vid]["dataset"] in datasets_set and vid not in exclude_set
            ]
        else:
            valid_indices = [i for i, vid in enumerate(self._id_list) if vid not in exclude_set]

        if not valid_indices:
            return []

        assert self._embedding_matrix is not None
        # Compute similarities for valid indices
        valid_embeddings = self._embedding_matrix[valid_indices]
        distances = compute_cosine_distance_matrix(
            query.reshape(1, -1),
            valid_embeddings,
        )[0]

        similarities = 1.0 - distances

        # Get top K
        num_results = min(top_k, len(valid_indices))
        top_indices = np.argsort(similarities)[::-1][:num_results]

        results = []
        for rank, idx in enumerate(top_indices):
            original_idx = valid_indices[idx]
            video_id = self._id_list[original_idx]
            meta = self._metadata[video_id]

            results.append(
                SimilarityMatch(
                    video_id=video_id,
                    video_path=meta.get("video_path"),
                    dataset=meta["dataset"],
                    gloss=meta.get("gloss"),
                    similarity=float(similarities[idx]),
                    rank=rank + 1,
                    metadata={
                        k: v for k, v in meta.items() if k not in ["dataset", "gloss", "video_path"]
                    },
                )
            )

        return results

    def search_batch(
        self,
        queries: np.ndarray,
        top_k: int = 10,
        datasets: list[str] | None = None,
    ) -> list[list[SimilarityMatch]]:
        """
        Search for multiple queries at once.

        Args:
            queries: Query embeddings (N, D)
            top_k: Number of results per query
            datasets: Optional list of datasets to search

        Returns:
            List of result lists, one per query
        """
        return [self.search(q, top_k, datasets) for q in queries]

    def get_datasets(self) -> list[str]:
        """Get list of unique datasets in store."""
        datasets = set()
        for meta in self._metadata.values():
            datasets.add(meta["dataset"])
        return sorted(datasets)

    def get_statistics(self) -> dict[str, Any]:
        """Get store statistics."""
        datasets: dict[str, int] = {}
        for meta in self._metadata.values():
            ds = meta["dataset"]
            datasets[ds] = datasets.get(ds, 0) + 1

        embedding_dim = 0
        if self._embeddings:
            first_emb = next(iter(self._embeddings.values()))
            embedding_dim = len(first_emb)

        return {
            "total_videos": len(self._embeddings),
            "datasets": datasets,
            "embedding_dim": embedding_dim,
        }

    def save(self, path: str | Path | None = None):
        """
        Save store to HDF5 file.

        Args:
            path: Save path (uses store_path if not provided)
        """
        save_path = Path(path) if path else self.store_path
        if save_path is None:
            raise ValueError("No save path specified")

        save_path.parent.mkdir(parents=True, exist_ok=True)

        with h5py.File(save_path, "w") as f:
            # Create groups
            emb_group = f.create_group("embeddings")
            meta_group = f.create_group("metadata")

            for video_id, embedding in self._embeddings.items():
                emb_group.create_dataset(video_id, data=embedding)

            for video_id, meta in self._metadata.items():
                meta_group.attrs[video_id] = json.dumps(meta)

        logger.info(f"Saved {len(self._embeddings)} embeddings to {save_path}")

    def _load_from_file(self):
        """Load store from HDF5 file."""
        if not self.store_path or not self.store_path.exists():
            return

        with h5py.File(self.store_path, "r") as f:
            if "embeddings" in f:
                for video_id in f["embeddings"].keys():
                    self._embeddings[video_id] = f["embeddings"][video_id][()]

            if "metadata" in f:
                for video_id in f["metadata"].attrs.keys():
                    self._metadata[video_id] = json.loads(f["metadata"].attrs[video_id])

        logger.info(f"Loaded {len(self._embeddings)} embeddings from {self.store_path}")

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

    def __contains__(self, video_id: str) -> bool:
        return video_id in self._embeddings


class SimilaritySearch:
    """
    High-level interface for cross-dataset similarity search.

    Example:
        embedder = SignRepEmbedder()
        search = SimilaritySearch(embedder, store_path="embeddings.h5")

        # Index a dataset
        search.index_dataset(my_dataset)

        # Find similar videos
        matches = search.find_similar("query.mp4", top_k=10)
    """

    def __init__(
        self,
        embedder,
        store: EmbeddingStore | None = None,
        store_path: str | Path | None = None,
    ):
        """
        Initialize similarity search.

        Args:
            embedder: Video embedder (e.g., SignRepEmbedder)
            store: Optional pre-existing embedding store
            store_path: Optional path for embedding store persistence
        """
        self.embedder = embedder
        self.store = store or EmbeddingStore(store_path)

    def index_video(
        self,
        video_path: str | Path,
        video_id: str | None = None,
        dataset: str = "unknown",
        gloss: str | None = None,
        **metadata,
    ) -> str:
        """
        Compute and store embedding for a video.

        Args:
            video_path: Path to video file
            video_id: Optional custom ID (uses filename if not provided)
            dataset: Dataset name for filtering
            gloss: Optional sign gloss/label
            **metadata: Additional metadata

        Returns:
            Video ID used in store
        """
        video_path = Path(video_path)
        video_id = video_id or video_path.stem

        # Compute embedding
        embedding = self.embedder.embed_video(video_path)

        # Store
        self.store.add(
            video_id=video_id,
            embedding=embedding,
            dataset=dataset,
            gloss=gloss,
            video_path=str(video_path),
            **metadata,
        )

        return video_id

    def index_dataset(
        self,
        dataset,
        max_samples: int | None = None,
        show_progress: bool = True,
    ) -> int:
        """
        Index all videos in a dataset.

        Args:
            dataset: SLTK dataset object with video samples
            max_samples: Maximum samples to index (None = all)
            show_progress: Show progress bar

        Returns:
            Number of videos indexed
        """
        count = 0

        if show_progress:
            try:
                from tqdm import tqdm
            except ImportError:
                tqdm = None
        else:
            tqdm = None

        samples = list(dataset)
        if max_samples:
            samples = samples[:max_samples]

        if tqdm:
            samples = tqdm(samples, desc=f"Indexing {dataset.NAME}")

        for sample in samples:
            if sample.video_path is None or not Path(sample.video_path).exists():
                continue

            try:
                self.index_video(
                    video_path=sample.video_path,
                    video_id=sample.id,
                    dataset=dataset.NAME,
                    gloss=sample.glosses[0] if sample.glosses else None,
                    signer_id=getattr(sample, "signer_id", None),
                )
                count += 1
            except Exception as e:
                logger.warning(f"Failed to index {sample.id}: {e}")

        return count

    def find_similar(
        self,
        query: str | Path | np.ndarray,
        top_k: int = 10,
        datasets: list[str] | None = None,
        exclude_self: bool = True,
    ) -> list[SimilarityMatch]:
        """
        Find videos similar to a query.

        Args:
            query: Query video path or embedding vector
            top_k: Number of results to return
            datasets: Optional list of datasets to search (None = all)
            exclude_self: Exclude query video from results if it's in the store

        Returns:
            List of SimilarityMatch objects
        """
        # Get query embedding
        if isinstance(query, (str, Path)):
            query_path = Path(query)
            query_embedding = self.embedder.embed_video(query_path)
            exclude_ids = [query_path.stem] if exclude_self else None
        else:
            query_embedding = query
            exclude_ids = None

        return self.store.search(
            query=query_embedding,
            top_k=top_k,
            datasets=datasets,
            exclude_ids=exclude_ids,
        )

    def find_similar_in_dataset(
        self,
        query: str | Path | np.ndarray,
        target_dataset: str,
        top_k: int = 10,
    ) -> list[SimilarityMatch]:
        """
        Find similar videos in a specific dataset.

        Useful for cross-dataset search (e.g., find WLASL matches for a BSLCP video).

        Args:
            query: Query video path or embedding
            target_dataset: Dataset to search in
            top_k: Number of results

        Returns:
            List of matches from target dataset
        """
        return self.find_similar(
            query=query,
            top_k=top_k,
            datasets=[target_dataset],
        )

    def cross_dataset_search(
        self,
        query: str | Path | np.ndarray,
        source_dataset: str,
        target_datasets: list[str],
        top_k: int = 10,
    ) -> dict[str, list[SimilarityMatch]]:
        """
        Find similar videos across multiple datasets.

        Args:
            query: Query video path or embedding
            source_dataset: Dataset the query is from (for exclusion)
            target_datasets: Datasets to search in
            top_k: Number of results per dataset

        Returns:
            Dictionary mapping dataset name to matches
        """
        results = {}

        for dataset in target_datasets:
            if dataset == source_dataset:
                continue

            matches = self.find_similar_in_dataset(
                query=query,
                target_dataset=dataset,
                top_k=top_k,
            )
            results[dataset] = matches

        return results

    def save(self, path: str | Path | None = None):
        """Save the embedding store."""
        self.store.save(path)

    def get_statistics(self) -> dict[str, Any]:
        """Get search statistics."""
        stats = self.store.get_statistics()
        stats["embedder_loaded"] = self.embedder.is_loaded()
        stats["embedding_dim"] = self.embedder.embedding_dim
        return stats
