"""Tests for rbsp.indexer.merge — Tiered and Log merge policies."""

from __future__ import annotations

import pytest

from rbsp.indexer.merge import (
    LogMergePolicy,
    MergePolicy,
    MergeSegment,
    TieredMergePolicy,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _seg(id: str, size: int, doc_count: int = 10) -> MergeSegment:
    return MergeSegment(id=id, doc_count=doc_count, size_bytes=size)


def _ids(groups):
    """Extract lists of segment IDs from merge groups."""
    return [[s.id for s in g] for g in groups]


# ---------------------------------------------------------------------------
# MergeSegment dataclass
# ---------------------------------------------------------------------------

class TestMergeSegment:
    def test_fields(self):
        seg = MergeSegment(id="a", doc_count=5, size_bytes=1024)
        assert seg.id == "a"
        assert seg.doc_count == 5
        assert seg.size_bytes == 1024

    def test_equality(self):
        a = MergeSegment(id="a", doc_count=5, size_bytes=1024)
        b = MergeSegment(id="a", doc_count=5, size_bytes=1024)
        assert a == b

    def test_slots(self):
        seg = MergeSegment(id="x", doc_count=1, size_bytes=0)
        with pytest.raises(AttributeError):
            seg.extra = 42  # type: ignore[attr-defined]


# ---------------------------------------------------------------------------
# MergePolicy (abstract)
# ---------------------------------------------------------------------------

class TestMergePolicyAbstract:
    def test_cannot_instantiate(self):
        with pytest.raises(TypeError):
            MergePolicy()  # type: ignore[abstract]


# ---------------------------------------------------------------------------
# TieredMergePolicy
# ---------------------------------------------------------------------------

class TestTieredMergePolicy:
    def test_no_merges_single_segment(self):
        policy = TieredMergePolicy(segments_per_tier=3)
        segs = [_seg("a", 1000)]
        assert policy.find_merges(segs) == []

    def test_no_merges_empty(self):
        policy = TieredMergePolicy()
        assert policy.find_merges([]) == []

    def test_no_merges_under_threshold(self):
        policy = TieredMergePolicy(segments_per_tier=5)
        segs = [_seg(f"s{i}", 1000) for i in range(5)]
        assert policy.find_merges(segs) == []

    def test_merges_when_tier_exceeds(self):
        # 6 segments all at floor size, segments_per_tier=3 -> should merge
        policy = TieredMergePolicy(segments_per_tier=3, max_merge_at_once=10)
        segs = [_seg(f"s{i}", 500) for i in range(6)]
        merges = policy.find_merges(segs)
        assert len(merges) >= 1
        # Total merged segments should reduce the tier below threshold
        total_merged = sum(len(g) for g in merges)
        assert total_merged >= 3

    def test_respects_max_merge_at_once(self):
        policy = TieredMergePolicy(
            segments_per_tier=2, max_merge_at_once=3
        )
        segs = [_seg(f"s{i}", 500) for i in range(10)]
        merges = policy.find_merges(segs)
        for group in merges:
            assert len(group) <= 3

    def test_respects_max_merged_segment_size(self):
        policy = TieredMergePolicy(
            segments_per_tier=2,
            max_merge_at_once=10,
            max_merged_segment_size=3000,
            floor_segment_size=1024,
        )
        # Each segment effective size = 1024 (floor)
        segs = [_seg(f"s{i}", 500) for i in range(20)]
        merges = policy.find_merges(segs)
        for group in merges:
            total = sum(max(s.size_bytes, policy.floor_segment_size) for s in group)
            assert total <= policy.max_merged_segment_size

    def test_floor_segment_size(self):
        policy = TieredMergePolicy(
            floor_segment_size=2048,
            segments_per_tier=2,
        )
        # Tiny segments treated as floor size
        segs = [_seg(f"s{i}", 10) for i in range(5)]
        merges = policy.find_merges(segs)
        assert len(merges) >= 1

    def test_tiered_grouping(self):
        # Mix of sizes — small and large segments should be in different tiers
        policy = TieredMergePolicy(
            segments_per_tier=2,
            floor_segment_size=100,
            max_merge_at_once=10,
        )
        segs = [
            _seg("small1", 50),
            _seg("small2", 80),
            _seg("small3", 90),
            _seg("big1", 100_000),
            _seg("big2", 120_000),
            _seg("big3", 130_000),
        ]
        merges = policy.find_merges(segs)
        # At least the small tier should merge (3 > 2)
        assert len(merges) >= 1

    def test_validation_max_merge_at_once(self):
        with pytest.raises(ValueError, match="max_merge_at_once"):
            TieredMergePolicy(max_merge_at_once=1)

    def test_validation_segments_per_tier(self):
        with pytest.raises(ValueError, match="segments_per_tier"):
            TieredMergePolicy(segments_per_tier=1)

    def test_default_params(self):
        policy = TieredMergePolicy()
        assert policy.max_merge_at_once == 10
        assert policy.segments_per_tier == 10
        assert policy.floor_segment_size == 1024
        assert policy.max_merged_segment_size == 5_000_000

    def test_works_with_segment_from_segments_module(self):
        """Merge policies should work with segments.Segment via duck typing."""
        from rbsp.indexer.segments import Segment

        policy = TieredMergePolicy(segments_per_tier=2)
        segs = [
            Segment(id="a", doc_count=5, size_bytes=500),
            Segment(id="b", doc_count=5, size_bytes=600),
            Segment(id="c", doc_count=5, size_bytes=700),
            Segment(id="d", doc_count=5, size_bytes=800),
        ]
        merges = policy.find_merges(segs)
        assert len(merges) >= 1
        for group in merges:
            for seg in group:
                assert hasattr(seg, "id")


# ---------------------------------------------------------------------------
# LogMergePolicy
# ---------------------------------------------------------------------------

class TestLogMergePolicy:
    def test_no_merges_single_segment(self):
        policy = LogMergePolicy(merge_factor=3)
        segs = [_seg("a", 1000)]
        assert policy.find_merges(segs) == []

    def test_no_merges_empty(self):
        policy = LogMergePolicy()
        assert policy.find_merges([]) == []

    def test_no_merges_under_factor(self):
        policy = LogMergePolicy(merge_factor=5)
        segs = [_seg(f"s{i}", 1024) for i in range(4)]
        assert policy.find_merges(segs) == []

    def test_merges_at_factor(self):
        policy = LogMergePolicy(merge_factor=3)
        # All same size -> same tier, should merge groups of 3
        segs = [_seg(f"s{i}", 1024) for i in range(6)]
        merges = policy.find_merges(segs)
        assert len(merges) == 2
        assert all(len(g) == 3 for g in merges)

    def test_merges_groups_by_size_tier(self):
        policy = LogMergePolicy(merge_factor=2)
        # Two tiers: ~1024 (tier 10) and ~1_000_000 (tier 19)
        segs = [
            _seg("small1", 1024),
            _seg("small2", 1024),
            _seg("big1", 1_000_000),
            _seg("big2", 1_000_000),
        ]
        merges = policy.find_merges(segs)
        assert len(merges) == 2
        ids = _ids(merges)
        # Each group should have segments of the same tier
        assert {"small1", "small2"} in [set(g) for g in ids]
        assert {"big1", "big2"} in [set(g) for g in ids]

    def test_leftover_not_merged(self):
        policy = LogMergePolicy(merge_factor=3)
        segs = [_seg(f"s{i}", 1024) for i in range(5)]
        merges = policy.find_merges(segs)
        # Only one group of 3; 2 leftover
        assert len(merges) == 1
        assert len(merges[0]) == 3

    def test_validation_merge_factor(self):
        with pytest.raises(ValueError, match="merge_factor"):
            LogMergePolicy(merge_factor=1)

    def test_default_merge_factor(self):
        policy = LogMergePolicy()
        assert policy.merge_factor == 10

    def test_zero_size_segments(self):
        policy = LogMergePolicy(merge_factor=2)
        segs = [_seg(f"s{i}", 0) for i in range(4)]
        merges = policy.find_merges(segs)
        assert len(merges) == 2
        assert all(len(g) == 2 for g in merges)

    def test_works_with_segment_from_segments_module(self):
        from rbsp.indexer.segments import Segment

        policy = LogMergePolicy(merge_factor=2)
        segs = [
            Segment(id="a", doc_count=5, size_bytes=1024),
            Segment(id="b", doc_count=5, size_bytes=1024),
        ]
        merges = policy.find_merges(segs)
        assert len(merges) == 1
        assert len(merges[0]) == 2
