"""Tiered merge policy for RBSP segment management.

Decides which index segments to merge and when, balancing query performance
(fewer segments = faster) against write amplification (fewer merges = less I/O).

Two policies are provided:

- **TieredMergePolicy**: Lucene-style tiered approach that groups segments by
  size and merges within tiers when the count exceeds ``segments_per_tier``.
- **LogMergePolicy**: Simpler logarithmic approach that merges groups of
  ``merge_factor`` same-tier segments.

Merge policies accept any object with ``id``, ``doc_count``, and ``size_bytes``
attributes — the :class:`~rbsp.indexer.segments.Segment` dataclass satisfies
this directly, or use :class:`MergeSegment` for lightweight standalone usage.

Zero external dependencies.
"""

from __future__ import annotations

import math
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Dict, List


# ---------------------------------------------------------------------------
# Lightweight segment descriptor (for standalone usage)
# ---------------------------------------------------------------------------

@dataclass(slots=True)
class MergeSegment:
    """Lightweight segment descriptor for merge policy evaluation.

    The merge policies accept any object with ``id``, ``doc_count``, and
    ``size_bytes`` attributes (duck typing).  Use this dataclass when you
    don't need the full :class:`~rbsp.indexer.segments.Segment`.
    """

    id: str
    doc_count: int
    size_bytes: int


# ---------------------------------------------------------------------------
# Base class
# ---------------------------------------------------------------------------

class MergePolicy(ABC):
    """Abstract base for merge policies."""

    @abstractmethod
    def find_merges(self, segments: List[Any]) -> List[List[Any]]:
        """Determine which segments should be merged.

        Args:
            segments: All live segments in the index.  Each element must
                have ``id: str``, ``doc_count: int``, and
                ``size_bytes: int`` attributes.

        Returns:
            A list of merge groups.  Each group is a list of segments that
            should be merged into a single new segment.  An empty list means
            no merges are needed right now.
        """


# ---------------------------------------------------------------------------
# Tiered merge policy (Lucene-style)
# ---------------------------------------------------------------------------

class TieredMergePolicy(MergePolicy):
    """Lucene-style tiered merge policy.

    Segments are sorted by size and grouped into exponentially-growing tiers.
    When any tier accumulates more than *segments_per_tier* segments a merge
    is scheduled for the smallest ``max_merge_at_once`` segments in that tier.
    """

    def __init__(
        self,
        *,
        max_merge_at_once: int = 10,
        segments_per_tier: int = 10,
        floor_segment_size: int = 1024,
        max_merged_segment_size: int = 5_000_000,
    ) -> None:
        if max_merge_at_once < 2:
            raise ValueError("max_merge_at_once must be >= 2")
        if segments_per_tier < 2:
            raise ValueError("segments_per_tier must be >= 2")

        self.max_merge_at_once = max_merge_at_once
        self.segments_per_tier = segments_per_tier
        self.floor_segment_size = floor_segment_size
        self.max_merged_segment_size = max_merged_segment_size

    def find_merges(self, segments: List[Any]) -> List[List[Any]]:
        if len(segments) <= 1:
            return []

        # Apply floor — treat tiny segments as at least floor_segment_size
        def _effective_size(seg: Any) -> int:
            return max(seg.size_bytes, self.floor_segment_size)

        # Sort ascending by effective size
        sorted_segs = sorted(segments, key=_effective_size)

        # Assign segments to tiers.  Tier boundaries grow exponentially:
        # tier 0: [floor, floor * spt)
        # tier 1: [floor * spt, floor * spt^2)
        # ...
        tiers: Dict[int, List[Any]] = {}
        for seg in sorted_segs:
            size = _effective_size(seg)
            if size <= self.floor_segment_size:
                tier = 0
            else:
                tier = int(math.log(size / self.floor_segment_size,
                                    self.segments_per_tier))
            tiers.setdefault(tier, []).append(seg)

        merges: List[List[Any]] = []

        # Walk tiers from smallest to largest
        for tier_idx in sorted(tiers):
            tier_segs = tiers[tier_idx]
            if len(tier_segs) <= self.segments_per_tier:
                continue

            # Pick the smallest segments first for merging (they are already
            # sorted ascending by size within the tier).
            remaining = list(tier_segs)
            while len(remaining) > self.segments_per_tier:
                excess = len(remaining) - self.segments_per_tier
                # Need at least 2 segments for a meaningful merge
                batch_size = min(self.max_merge_at_once, max(excess, 2))
                if batch_size > len(remaining):
                    break

                batch = remaining[:batch_size]

                # Skip if the merged result would exceed the cap
                total = sum(_effective_size(s) for s in batch)
                if total > self.max_merged_segment_size:
                    break

                merges.append(batch)
                remaining = remaining[batch_size:]

        return merges


# ---------------------------------------------------------------------------
# Log merge policy
# ---------------------------------------------------------------------------

class LogMergePolicy(MergePolicy):
    """Simple logarithmic merge policy.

    Segments are assigned to tiers based on the floor-log of their size.
    Whenever a tier contains at least ``merge_factor`` segments the oldest
    (smallest) ``merge_factor`` are scheduled for merging.
    """

    def __init__(self, *, merge_factor: int = 10) -> None:
        if merge_factor < 2:
            raise ValueError("merge_factor must be >= 2")
        self.merge_factor = merge_factor

    def find_merges(self, segments: List[Any]) -> List[List[Any]]:
        if len(segments) <= 1:
            return []

        # Assign each segment to a tier = floor(log2(size_bytes)).
        # Segments with size 0 go to tier 0.
        tiers: Dict[int, List[Any]] = {}
        for seg in segments:
            tier = int(math.log2(seg.size_bytes)) if seg.size_bytes > 0 else 0
            tiers.setdefault(tier, []).append(seg)

        merges: List[List[Any]] = []

        for tier_idx in sorted(tiers):
            tier_segs = sorted(tiers[tier_idx], key=lambda s: s.size_bytes)
            while len(tier_segs) >= self.merge_factor:
                batch = tier_segs[:self.merge_factor]
                merges.append(batch)
                tier_segs = tier_segs[self.merge_factor:]

        return merges
