"""Tests for embedding and similarity search functionality."""

import tempfile
from pathlib import Path

import numpy as np
import pytest

from sltk.embedding.similarity import (
    EmbeddingStore,
    SimilarityMatch,
    compute_cosine_distance_matrix,
    compute_cosine_similarity,
)


class TestCosineSimilarity:
    """Tests for cosine similarity computation."""

    def test_identical_vectors(self):
        """Test similarity of identical vectors is 1."""
        a = np.array([1.0, 2.0, 3.0])
        assert compute_cosine_similarity(a, a) == pytest.approx(1.0)

    def test_orthogonal_vectors(self):
        """Test similarity of orthogonal vectors is 0."""
        a = np.array([1.0, 0.0])
        b = np.array([0.0, 1.0])
        assert compute_cosine_similarity(a, b) == pytest.approx(0.0)

    def test_opposite_vectors(self):
        """Test similarity of opposite vectors is -1."""
        a = np.array([1.0, 0.0])
        b = np.array([-1.0, 0.0])
        assert compute_cosine_similarity(a, b) == pytest.approx(-1.0)

    def test_zero_vector(self):
        """Test similarity with zero vector is 0."""
        a = np.array([1.0, 2.0, 3.0])
        b = np.array([0.0, 0.0, 0.0])
        assert compute_cosine_similarity(a, b) == 0.0

    def test_normalized_vectors(self):
        """Test similarity of unit vectors."""
        a = np.array([1.0, 0.0, 0.0])
        b = np.array([0.5, 0.5, 0.0])
        b = b / np.linalg.norm(b)  # Normalize

        # cos(45 degrees) = sqrt(2)/2
        expected = np.sqrt(2) / 2
        assert compute_cosine_similarity(a, b) == pytest.approx(expected, rel=1e-5)


class TestCosineDistanceMatrix:
    """Tests for cosine distance matrix computation."""

    def test_single_query(self):
        """Test distance matrix with single query."""
        queries = np.array([[1.0, 0.0, 0.0]])
        database = np.array(
            [
                [1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 0.0, 1.0],
            ]
        )

        distances = compute_cosine_distance_matrix(queries, database)

        assert distances.shape == (1, 3)
        assert distances[0, 0] == pytest.approx(0.0, abs=1e-6)  # Same vector
        assert distances[0, 1] == pytest.approx(1.0, abs=1e-6)  # Orthogonal
        assert distances[0, 2] == pytest.approx(1.0, abs=1e-6)  # Orthogonal

    def test_multiple_queries(self):
        """Test distance matrix with multiple queries."""
        queries = np.array(
            [
                [1.0, 0.0],
                [0.0, 1.0],
            ]
        )
        database = np.array(
            [
                [1.0, 0.0],
                [0.0, 1.0],
            ]
        )

        distances = compute_cosine_distance_matrix(queries, database)

        assert distances.shape == (2, 2)
        # Identity matches
        assert distances[0, 0] == pytest.approx(0.0, abs=1e-6)
        assert distances[1, 1] == pytest.approx(0.0, abs=1e-6)
        # Cross matches (orthogonal)
        assert distances[0, 1] == pytest.approx(1.0, abs=1e-6)
        assert distances[1, 0] == pytest.approx(1.0, abs=1e-6)


class TestSimilarityMatch:
    """Tests for SimilarityMatch dataclass."""

    def test_to_dict(self):
        """Test conversion to dictionary."""
        match = SimilarityMatch(
            video_id="test_video",
            video_path="/path/to/video.mp4",
            dataset="wlasl",
            gloss="HELLO",
            similarity=0.95,
            rank=1,
            metadata={"signer": "person1"},
        )

        d = match.to_dict()

        assert d["video_id"] == "test_video"
        assert d["video_path"] == "/path/to/video.mp4"
        assert d["dataset"] == "wlasl"
        assert d["gloss"] == "HELLO"
        assert d["similarity"] == 0.95
        assert d["rank"] == 1
        assert d["metadata"]["signer"] == "person1"


class TestEmbeddingStore:
    """Tests for EmbeddingStore."""

    @pytest.fixture
    def store(self):
        """Create a fresh embedding store."""
        return EmbeddingStore()

    @pytest.fixture
    def populated_store(self, store):
        """Create store with some embeddings."""
        # Add embeddings from different datasets
        store.add(
            video_id="wlasl_001",
            embedding=np.array([1.0, 0.0, 0.0]),
            dataset="wlasl",
            gloss="HELLO",
        )
        store.add(
            video_id="wlasl_002",
            embedding=np.array([0.9, 0.1, 0.0]),
            dataset="wlasl",
            gloss="HI",
        )
        store.add(
            video_id="bslcp_001",
            embedding=np.array([0.8, 0.2, 0.0]),
            dataset="bslcp",
            gloss="GREET",
        )
        store.add(
            video_id="how2sign_001",
            embedding=np.array([0.0, 1.0, 0.0]),
            dataset="how2sign",
            gloss="WAVE",
        )
        return store

    def test_add_and_get(self, store):
        """Test adding and retrieving embeddings."""
        embedding = np.array([1.0, 2.0, 3.0])
        store.add(
            video_id="test_001",
            embedding=embedding,
            dataset="test",
            gloss="SIGN",
        )

        result = store.get("test_001")
        assert result is not None

        retrieved_emb, metadata = result
        np.testing.assert_array_almost_equal(retrieved_emb, embedding)
        assert metadata["dataset"] == "test"
        assert metadata["gloss"] == "SIGN"

    def test_contains(self, store):
        """Test containment check."""
        store.add("vid1", np.array([1, 2, 3]), "test")

        assert "vid1" in store
        assert "vid2" not in store

    def test_remove(self, store):
        """Test removing embeddings."""
        store.add("vid1", np.array([1, 2, 3]), "test")
        assert "vid1" in store

        result = store.remove("vid1")
        assert result is True
        assert "vid1" not in store

        # Removing non-existent returns False
        result = store.remove("vid1")
        assert result is False

    def test_len(self, store):
        """Test length."""
        assert len(store) == 0

        store.add("vid1", np.array([1, 2, 3]), "test")
        assert len(store) == 1

        store.add("vid2", np.array([4, 5, 6]), "test")
        assert len(store) == 2

    def test_search_basic(self, populated_store):
        """Test basic similarity search."""
        query = np.array([1.0, 0.0, 0.0])
        matches = populated_store.search(query, top_k=2)

        assert len(matches) == 2
        assert matches[0].video_id == "wlasl_001"  # Exact match
        assert matches[0].similarity == pytest.approx(1.0)
        assert matches[0].rank == 1

    def test_search_filter_by_dataset(self, populated_store):
        """Test filtering search by dataset."""
        query = np.array([1.0, 0.0, 0.0])

        # Only search in bslcp
        matches = populated_store.search(query, top_k=10, datasets=["bslcp"])

        assert len(matches) == 1
        assert matches[0].video_id == "bslcp_001"

    def test_search_exclude_ids(self, populated_store):
        """Test excluding specific IDs from search."""
        query = np.array([1.0, 0.0, 0.0])

        matches = populated_store.search(query, top_k=10, exclude_ids=["wlasl_001"])

        # wlasl_001 should be excluded
        video_ids = [m.video_id for m in matches]
        assert "wlasl_001" not in video_ids

    def test_search_empty_store(self, store):
        """Test searching empty store."""
        query = np.array([1.0, 0.0, 0.0])
        matches = store.search(query)
        assert len(matches) == 0

    def test_get_datasets(self, populated_store):
        """Test getting list of datasets."""
        datasets = populated_store.get_datasets()

        assert sorted(datasets) == ["bslcp", "how2sign", "wlasl"]

    def test_get_statistics(self, populated_store):
        """Test getting store statistics."""
        stats = populated_store.get_statistics()

        assert stats["total_videos"] == 4
        assert stats["datasets"]["wlasl"] == 2
        assert stats["datasets"]["bslcp"] == 1
        assert stats["datasets"]["how2sign"] == 1
        assert stats["embedding_dim"] == 3

    def test_save_and_load(self, populated_store):
        """Test saving and loading store."""
        with tempfile.TemporaryDirectory() as tmpdir:
            save_path = Path(tmpdir) / "embeddings.h5"

            # Save
            populated_store.save(save_path)
            assert save_path.exists()

            # Load into new store
            loaded_store = EmbeddingStore(save_path)

            assert len(loaded_store) == len(populated_store)
            assert loaded_store.get_datasets() == populated_store.get_datasets()

            # Verify content
            result = loaded_store.get("wlasl_001")
            assert result is not None
            emb, meta = result
            assert meta["gloss"] == "HELLO"


class TestEmbeddingStoreSearchQuality:
    """Tests for search quality and ranking."""

    @pytest.fixture
    def store_with_gradients(self):
        """Create store with embeddings at various similarities."""
        store = EmbeddingStore()

        # Reference embedding
        store.add("ref", np.array([1.0, 0.0, 0.0]), "ref", "REF")

        # Embeddings with decreasing similarity to reference
        similarities = [0.99, 0.9, 0.8, 0.7, 0.5, 0.3, 0.1, 0.0]
        for i, sim in enumerate(similarities):
            # Create embedding with approximately this similarity
            angle = np.arccos(sim)
            emb = np.array([np.cos(angle), np.sin(angle), 0.0])
            store.add(f"vid_{i}", emb, "test", f"SIM_{sim}")

        return store

    def test_ranking_order(self, store_with_gradients):
        """Test that results are ranked by similarity."""
        query = np.array([1.0, 0.0, 0.0])
        matches = store_with_gradients.search(query, top_k=10, exclude_ids=["ref"])

        # Check ranking
        for i in range(len(matches) - 1):
            assert matches[i].similarity >= matches[i + 1].similarity
            assert matches[i].rank < matches[i + 1].rank

    def test_top_k_limits(self, store_with_gradients):
        """Test that top_k limits results."""
        query = np.array([1.0, 0.0, 0.0])

        matches_5 = store_with_gradients.search(query, top_k=5)
        matches_3 = store_with_gradients.search(query, top_k=3)

        assert len(matches_5) == 5
        assert len(matches_3) == 3

        # Top 3 should be subset of top 5
        top3_ids = {m.video_id for m in matches_3}
        top5_ids = {m.video_id for m in matches_5}
        assert top3_ids.issubset(top5_ids)


def _get_signrep_root() -> Path:
    """Get SignRep root for tests."""
    import os

    env_root = os.environ.get("SLTK_SIGNREP_ROOT")
    if env_root:
        return Path(env_root)
    return Path.home() / "Projects" / "SignRep"


def _signrep_available() -> bool:
    """Check if SignRep is available."""
    signrep_path = _get_signrep_root()
    ckpt_path = signrep_path / "v11F4_m80_rgba_disc3_kl" / "ckpt.pt"
    return signrep_path.exists() and ckpt_path.exists()


class TestSignRepIntegration:
    """Integration tests for SignRep (skipped if not available)."""

    @pytest.fixture
    def signrep_available(self):
        """Check if SignRep is available."""
        return _signrep_available()

    @pytest.mark.skipif(not _signrep_available(), reason="SignRep not available")
    def test_signrep_config(self):
        """Test SignRepConfig creation."""
        from sltk.embedding import SignRepConfig

        config = SignRepConfig()
        assert config.segment_length == 16
        assert config.stride == 2
        assert config.out_dim == 768

    @pytest.mark.skipif(not _signrep_available(), reason="SignRep not available")
    def test_signrep_lazy_load(self):
        """Test SignRepEmbedder with lazy loading."""
        from sltk.embedding import SignRepEmbedder

        embedder = SignRepEmbedder(lazy_load=True)
        assert not embedder.is_loaded()
        assert embedder.embedding_dim == 768


class TestSimilaritySearchMock:
    """Tests for SimilaritySearch with mocked embedder."""

    class MockEmbedder:
        """Mock embedder for testing."""

        def __init__(self):
            self._embeddings = {}

        def embed_video(self, video_path):
            # Return deterministic embedding based on path
            path = str(video_path)
            if path not in self._embeddings:
                # Generate random but deterministic embedding
                np.random.seed(hash(path) % (2**32))
                self._embeddings[path] = np.random.randn(768)
            return self._embeddings[path]

        @property
        def embedding_dim(self):
            return 768

        def is_loaded(self):
            return True

    @pytest.fixture
    def mock_search(self):
        """Create SimilaritySearch with mock embedder."""
        from sltk.embedding.similarity import EmbeddingStore, SimilaritySearch

        embedder = self.MockEmbedder()
        store = EmbeddingStore()
        return SimilaritySearch(embedder, store=store)

    def test_index_and_search(self, mock_search, tmp_path):
        """Test indexing videos and searching."""
        # Create fake video files
        for i in range(5):
            (tmp_path / f"video_{i}.mp4").touch()

        # Index videos
        for i in range(5):
            mock_search.index_video(
                tmp_path / f"video_{i}.mp4",
                dataset="test",
                gloss=f"SIGN_{i}",
            )

        # Search
        query_emb = mock_search.embedder.embed_video(tmp_path / "video_0.mp4")
        matches = mock_search.find_similar(query_emb, top_k=3)

        assert len(matches) == 3
        assert matches[0].rank == 1

    def test_cross_dataset_search(self, mock_search, tmp_path):
        """Test cross-dataset search."""
        # Create videos for two datasets
        for ds in ["dataset_a", "dataset_b"]:
            ds_dir = tmp_path / ds
            ds_dir.mkdir()
            for i in range(3):
                (ds_dir / f"video_{i}.mp4").touch()
                mock_search.index_video(
                    ds_dir / f"video_{i}.mp4",
                    dataset=ds,
                    gloss=f"SIGN_{i}",
                )

        # Cross-dataset search
        query = mock_search.embedder.embed_video(tmp_path / "dataset_a" / "video_0.mp4")
        results = mock_search.cross_dataset_search(
            query=query,
            source_dataset="dataset_a",
            target_datasets=["dataset_a", "dataset_b"],
            top_k=2,
        )

        # Should only have results from dataset_b (source excluded)
        assert "dataset_a" not in results
        assert "dataset_b" in results
        assert len(results["dataset_b"]) == 2

    def test_get_statistics(self, mock_search, tmp_path):
        """Test getting search statistics."""
        # Index some videos
        for i in range(3):
            (tmp_path / f"video_{i}.mp4").touch()
            mock_search.index_video(tmp_path / f"video_{i}.mp4", dataset="test")

        stats = mock_search.get_statistics()

        assert stats["total_videos"] == 3
        assert stats["embedder_loaded"] is True
        assert stats["embedding_dim"] == 768
