"""Tests for rbsp.indexer.segments -- SegmentInfo, Segment, SegmentManager."""

import os
import sys
import time
import unittest

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

from rbsp.indexer.segments import SegmentInfo, Segment, SegmentManager


# ---------------------------------------------------------------------------
# SegmentInfo
# ---------------------------------------------------------------------------

class TestSegmentInfo(unittest.TestCase):
    """Tests for the SegmentInfo frozen dataclass."""

    def test_defaults(self):
        info = SegmentInfo()
        self.assertIsNone(info.min_doc_id)
        self.assertIsNone(info.max_doc_id)
        self.assertEqual(info.term_count, 0)
        self.assertEqual(info.doc_count, 0)
        self.assertEqual(info.size_bytes, 0)

    def test_custom_values(self):
        info = SegmentInfo(
            min_doc_id="aaa",
            max_doc_id="zzz",
            term_count=100,
            doc_count=50,
            size_bytes=4096,
        )
        self.assertEqual(info.min_doc_id, "aaa")
        self.assertEqual(info.max_doc_id, "zzz")
        self.assertEqual(info.term_count, 100)
        self.assertEqual(info.doc_count, 50)
        self.assertEqual(info.size_bytes, 4096)

    def test_frozen_cannot_mutate(self):
        info = SegmentInfo()
        with self.assertRaises(AttributeError):
            info.doc_count = 10


# ---------------------------------------------------------------------------
# Segment
# ---------------------------------------------------------------------------

class TestSegment(unittest.TestCase):
    """Tests for the Segment dataclass."""

    def test_defaults(self):
        seg = Segment(id="abc123")
        self.assertEqual(seg.id, "abc123")
        self.assertEqual(seg.doc_count, 0)
        self.assertEqual(seg.size_bytes, 0)
        self.assertIsInstance(seg.created_at, float)
        self.assertEqual(seg.merged_from, [])
        self.assertEqual(seg.doc_ids, set())
        self.assertIsNone(seg.info)

    def test_custom_values(self):
        seg = Segment(
            id="seg1",
            doc_count=5,
            size_bytes=1024,
            doc_ids={"d1", "d2"},
            merged_from=["old1"],
        )
        self.assertEqual(seg.doc_count, 5)
        self.assertEqual(seg.size_bytes, 1024)
        self.assertEqual(seg.doc_ids, {"d1", "d2"})
        self.assertEqual(seg.merged_from, ["old1"])

    def test_mutable_fields(self):
        seg = Segment(id="seg1")
        seg.doc_count = 10
        seg.merged_from.append("parent_seg")
        self.assertEqual(seg.doc_count, 10)
        self.assertEqual(seg.merged_from, ["parent_seg"])


# ---------------------------------------------------------------------------
# SegmentManager -- create, lookup
# ---------------------------------------------------------------------------

class TestSegmentManagerCreate(unittest.TestCase):
    """Tests for create_segment and lookups."""

    def setUp(self):
        self.mgr = SegmentManager()

    def test_create_segment(self):
        seg = self.mgr.create_segment(["d1", "d2", "d3"])
        self.assertEqual(seg.doc_count, 3)
        self.assertEqual(seg.doc_ids, {"d1", "d2", "d3"})
        self.assertIsNotNone(seg.id)

    def test_create_segment_deduplicates_doc_ids(self):
        seg = self.mgr.create_segment(["d1", "d1", "d2"])
        self.assertEqual(seg.doc_count, 2)
        self.assertEqual(seg.doc_ids, {"d1", "d2"})

    def test_create_empty_segment(self):
        seg = self.mgr.create_segment([])
        self.assertEqual(seg.doc_count, 0)
        self.assertEqual(seg.doc_ids, set())

    def test_get_segment_found(self):
        seg = self.mgr.create_segment(["d1"])
        found = self.mgr.get_segment(seg.id)
        self.assertIs(found, seg)

    def test_get_segment_not_found(self):
        result = self.mgr.get_segment("nonexistent")
        self.assertIsNone(result)

    def test_get_segments_returns_all(self):
        self.mgr.create_segment(["d1"])
        self.mgr.create_segment(["d2"])
        self.mgr.create_segment(["d3"])
        segments = self.mgr.get_segments()
        self.assertEqual(len(segments), 3)

    def test_get_segments_sorted_by_creation_time(self):
        s1 = self.mgr.create_segment(["d1"])
        time.sleep(0.01)
        s2 = self.mgr.create_segment(["d2"])
        time.sleep(0.01)
        s3 = self.mgr.create_segment(["d3"])
        segments = self.mgr.get_segments()
        self.assertEqual([s.id for s in segments], [s1.id, s2.id, s3.id])


# ---------------------------------------------------------------------------
# SegmentManager -- mark_merged
# ---------------------------------------------------------------------------

class TestSegmentManagerMerge(unittest.TestCase):
    """Tests for mark_merged."""

    def setUp(self):
        self.mgr = SegmentManager()

    def test_mark_merged_removes_old_segments(self):
        s1 = self.mgr.create_segment(["d1"])
        s2 = self.mgr.create_segment(["d2"])
        s3 = self.mgr.create_segment(["d1", "d2"])

        self.mgr.mark_merged([s1.id, s2.id], s3.id)

        self.assertIsNone(self.mgr.get_segment(s1.id))
        self.assertIsNone(self.mgr.get_segment(s2.id))
        self.assertIsNotNone(self.mgr.get_segment(s3.id))

    def test_mark_merged_records_lineage(self):
        s1 = self.mgr.create_segment(["d1"])
        s2 = self.mgr.create_segment(["d2"])
        s3 = self.mgr.create_segment(["d1", "d2"])

        self.mgr.mark_merged([s1.id, s2.id], s3.id)
        merged = self.mgr.get_segment(s3.id)
        self.assertEqual(merged.merged_from, [s1.id, s2.id])

    def test_mark_merged_nonexistent_new_id_raises(self):
        s1 = self.mgr.create_segment(["d1"])
        with self.assertRaises(KeyError):
            self.mgr.mark_merged([s1.id], "nonexistent")


# ---------------------------------------------------------------------------
# SegmentManager -- delete, info methods
# ---------------------------------------------------------------------------

class TestSegmentManagerDeleteAndInfo(unittest.TestCase):
    """Tests for delete_segment, total_docs, segment_count, contains, clear."""

    def setUp(self):
        self.mgr = SegmentManager()

    def test_delete_existing_segment(self):
        seg = self.mgr.create_segment(["d1"])
        self.assertTrue(self.mgr.delete_segment(seg.id))
        self.assertIsNone(self.mgr.get_segment(seg.id))

    def test_delete_nonexistent_segment(self):
        self.assertFalse(self.mgr.delete_segment("nope"))

    def test_total_docs(self):
        self.mgr.create_segment(["d1", "d2"])
        self.mgr.create_segment(["d3"])
        self.assertEqual(self.mgr.total_docs(), 3)

    def test_total_docs_empty(self):
        self.assertEqual(self.mgr.total_docs(), 0)

    def test_segment_count(self):
        self.mgr.create_segment(["d1"])
        self.mgr.create_segment(["d2"])
        self.assertEqual(self.mgr.segment_count(), 2)

    def test_len_equals_segment_count(self):
        self.mgr.create_segment(["d1"])
        self.assertEqual(len(self.mgr), self.mgr.segment_count())

    def test_contains_existing(self):
        seg = self.mgr.create_segment(["d1"])
        self.assertIn(seg.id, self.mgr)

    def test_contains_nonexistent(self):
        self.assertNotIn("fake_id", self.mgr)

    def test_clear(self):
        self.mgr.create_segment(["d1"])
        self.mgr.create_segment(["d2"])
        self.mgr.clear()
        self.assertEqual(self.mgr.segment_count(), 0)
        self.assertEqual(self.mgr.total_docs(), 0)


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