"""
Unit tests for the SearchAgent implementation.
"""

import unittest

from src.agent import SearchAgent, Result


class TestSearchAgent(unittest.TestCase):
    def setUp(self):
        self.corpus = [
            "The quick brown fox jumps over the lazy dog.",
            "Deep learning models can capture complex patterns in data.",
            "PyTorch is a popular deep learning framework.",
            "Natural language processing involves understanding text.",
        ]
        self.agent = SearchAgent(self.corpus, embedding_dim=20)

    def test_vocab_size(self):
        # Vocabulary should contain all unique words
        vocab_size = self.agent.get_vocab_size()
        # Count unique words manually
        unique_words = set()
        for doc in self.corpus:
            unique_words.update(doc.lower().split())
        self.assertEqual(vocab_size, len(unique_words))

    def test_document_embeddings_shape(self):
        # Document embeddings should have shape (n_docs, dim)
        doc_emb = self.agent.doc_embeddings
        self.assertEqual(doc_emb.shape, (len(self.corpus), self.agent.embedding_dim))

    def test_query_encoding_shape(self):
        query = "deep learning"
        vec = self.agent._encode_text(query)
        self.assertEqual(vec.shape, (self.agent.embedding_dim,))

    def test_cosine_similarity(self):
        # Compute similarity between two identical vectors
        vec = self.agent._encode_text("deep learning")
        sims = self.agent._cosine_similarity(vec.unsqueeze(0), vec.unsqueeze(0))
        self.assertAlmostEqual(sims.item(), 1.0, places=5)

    def test_search_ranking(self):
        # Query that matches second document
        results = self.agent.search("deep learning", top_k=2)
        # The first result should be the second document (index 1)
        self.assertEqual(results[0].doc_id, 1)
        self.assertGreater(results[0].score, results[1].score)

    def test_unknown_words(self):
        # Query with unknown words should still return results
        results = self.agent.search("xyz abc", top_k=1)
        self.assertEqual(len(results), 1)
        self.assertIsInstance(results[0], Result)

    def test_empty_query(self):
        # Empty query should return top documents based on zero vector
        results = self.agent.search("", top_k=3)
        self.assertEqual(len(results), 3)
        # Scores should be non‑negative
        for res in results:
            self.assertGreaterEqual(res.score, 0.0)


if __name__ == "__main__":
    unittest.main()