"""Tests for rbsp.runtime.parallel -- chunk_list and ParallelSearcher.merge_results."""

import os
import sys
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from rbsp.runtime.parallel import chunk_list, ParallelSearcher


# ---------------------------------------------------------------------------
# chunk_list tests
# ---------------------------------------------------------------------------

class TestChunkList(unittest.TestCase):
    """Tests for the chunk_list utility function."""

    def test_even_split(self):
        result = chunk_list([1, 2, 3, 4], 2)
        self.assertEqual(result, [[1, 2], [3, 4]])

    def test_uneven_split(self):
        result = chunk_list([1, 2, 3, 4, 5], 2)
        # First chunk gets the remainder element
        self.assertEqual(result, [[1, 2, 3], [4, 5]])

    def test_more_chunks_than_elements(self):
        result = chunk_list([1, 2], 5)
        # Should produce at most len(lst) chunks
        self.assertEqual(len(result), 2)
        # Each chunk has one element
        self.assertEqual(result, [[1], [2]])

    def test_single_chunk(self):
        result = chunk_list([1, 2, 3], 1)
        self.assertEqual(result, [[1, 2, 3]])

    def test_empty_list(self):
        result = chunk_list([], 3)
        self.assertEqual(result, [])

    def test_single_element(self):
        result = chunk_list([42], 3)
        self.assertEqual(result, [[42]])

    def test_n_equals_length(self):
        result = chunk_list([1, 2, 3], 3)
        self.assertEqual(result, [[1], [2], [3]])

    def test_n_zero_raises(self):
        with self.assertRaises(ValueError):
            chunk_list([1, 2], 0)

    def test_n_negative_raises(self):
        with self.assertRaises(ValueError):
            chunk_list([1, 2], -1)

    def test_all_elements_preserved(self):
        original = list(range(17))
        chunks = chunk_list(original, 4)
        flattened = [x for c in chunks for x in c]
        self.assertEqual(flattened, original)

    def test_no_empty_chunks(self):
        for n in range(1, 10):
            for size in range(0, 10):
                lst = list(range(size))
                chunks = chunk_list(lst, n)
                for chunk in chunks:
                    self.assertTrue(len(chunk) > 0, f"Empty chunk for lst={lst}, n={n}")


# ---------------------------------------------------------------------------
# ParallelSearcher.merge_results tests
# ---------------------------------------------------------------------------

class TestMergeResults(unittest.TestCase):
    """Tests for ParallelSearcher.merge_results static method."""

    def test_merge_two_shards(self):
        shard1 = [("a", 0.9), ("c", 0.5)]
        shard2 = [("b", 0.8), ("d", 0.4)]
        result = ParallelSearcher.merge_results([shard1, shard2], top_k=10)
        # Should be sorted descending by score
        scores = [score for _, score in result]
        self.assertEqual(scores, sorted(scores, reverse=True))
        self.assertEqual(len(result), 4)

    def test_merge_respects_top_k(self):
        shard1 = [("a", 0.9), ("c", 0.7), ("e", 0.5)]
        shard2 = [("b", 0.8), ("d", 0.6), ("f", 0.4)]
        result = ParallelSearcher.merge_results([shard1, shard2], top_k=3)
        self.assertEqual(len(result), 3)
        doc_ids = [doc_id for doc_id, _ in result]
        self.assertEqual(doc_ids, ["a", "b", "c"])

    def test_merge_empty_shard_results(self):
        result = ParallelSearcher.merge_results([], top_k=10)
        self.assertEqual(result, [])

    def test_merge_single_shard(self):
        shard = [("x", 1.0), ("y", 0.5)]
        result = ParallelSearcher.merge_results([shard], top_k=10)
        self.assertEqual(result, [("x", 1.0), ("y", 0.5)])

    def test_merge_with_empty_shards(self):
        shard1 = [("a", 0.9)]
        shard2 = []
        result = ParallelSearcher.merge_results([shard1, shard2], top_k=10)
        self.assertEqual(result, [("a", 0.9)])

    def test_merge_all_empty_shards(self):
        result = ParallelSearcher.merge_results([[], [], []], top_k=10)
        self.assertEqual(result, [])

    def test_merge_default_top_k_is_10(self):
        shard = [(f"doc{i}", 1.0 - i * 0.01) for i in range(20)]
        result = ParallelSearcher.merge_results([shard])
        self.assertEqual(len(result), 10)

    def test_merge_preserves_scores(self):
        shard1 = [("a", 0.95)]
        shard2 = [("b", 0.85)]
        result = ParallelSearcher.merge_results([shard1, shard2], top_k=10)
        result_dict = dict(result)
        self.assertAlmostEqual(result_dict["a"], 0.95)
        self.assertAlmostEqual(result_dict["b"], 0.85)

    def test_merge_three_shards(self):
        shard1 = [("a", 0.9)]
        shard2 = [("b", 0.8)]
        shard3 = [("c", 0.7)]
        result = ParallelSearcher.merge_results([shard1, shard2, shard3], top_k=10)
        doc_ids = [doc_id for doc_id, _ in result]
        self.assertEqual(doc_ids, ["a", "b", "c"])


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