"""Segment-based Index Architecture for RBSP.

Implements a Lucene-style segment model where the index is composed of
multiple immutable segments that are periodically merged.  This enables
concurrent reads during writes and efficient incremental indexing.

Each segment tracks its document set, byte size, creation time, and
optional lineage (which segments it was merged from).

Thread-safe for concurrent reads and writes.
"""

from __future__ import annotations

import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set


def _new_segment_id() -> str:
    """Generate a short unique segment identifier."""
    return uuid.uuid4().hex[:12]


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass(frozen=True, slots=True)
class SegmentInfo:
    """Metadata summary for a segment.

    Attributes:
        min_doc_id:  Lexicographically smallest document ID in the segment.
        max_doc_id:  Lexicographically largest document ID in the segment.
        term_count:  Number of unique terms across all documents.
        doc_count:   Number of documents in the segment.
        size_bytes:  Estimated byte size of the segment data.
    """

    min_doc_id: Optional[str] = None
    max_doc_id: Optional[str] = None
    term_count: int = 0
    doc_count: int = 0
    size_bytes: int = 0


@dataclass(slots=True)
class Segment:
    """An individual index segment.

    Attributes:
        id:          Unique segment identifier.
        doc_count:   Number of documents in this segment.
        size_bytes:  Estimated byte size of segment data.
        created_at:  Unix timestamp when the segment was created.
        merged_from: List of parent segment IDs this was merged from.
        doc_ids:     Set of document IDs contained in this segment.
        info:        Optional detailed metadata.
    """

    id: str
    doc_count: int = 0
    size_bytes: int = 0
    created_at: float = field(default_factory=time.time)
    merged_from: List[str] = field(default_factory=list)
    doc_ids: Set[str] = field(default_factory=set)
    info: Optional[SegmentInfo] = None


# ---------------------------------------------------------------------------
# Segment Manager
# ---------------------------------------------------------------------------

class SegmentManager:
    """Manages the lifecycle of index segments.

    Provides create, merge, and delete operations over a collection of
    segments.  All public methods are thread-safe.

    Usage::

        mgr = SegmentManager()
        seg = mgr.create_segment(["doc_1", "doc_2", "doc_3"])
        print(mgr.total_docs())   # 3
        print(mgr.segment_count())  # 1
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._segments: Dict[str, Segment] = {}

    # ------------------------------------------------------------------
    # Mutation
    # ------------------------------------------------------------------

    def create_segment(self, doc_ids: List[str]) -> Segment:
        """Create a new segment containing the given document IDs.

        Args:
            doc_ids: Document identifiers to include in the segment.

        Returns:
            The newly created :class:`Segment`.
        """
        seg_id = _new_segment_id()
        doc_set = set(doc_ids)
        segment = Segment(
            id=seg_id,
            doc_count=len(doc_set),
            size_bytes=0,
            doc_ids=doc_set,
        )
        with self._lock:
            self._segments[seg_id] = segment
        return segment

    def mark_merged(self, old_ids: List[str], new_id: str) -> None:
        """Mark segments as merged into a new segment.

        The old segments are removed and the new segment's ``merged_from``
        field is updated to record the lineage.

        Args:
            old_ids: Segment IDs that were merged.
            new_id:  The segment ID of the merge result.

        Raises:
            KeyError: If *new_id* does not exist in the manager.
        """
        with self._lock:
            if new_id not in self._segments:
                raise KeyError(f"Segment {new_id!r} not found")
            self._segments[new_id].merged_from = list(old_ids)
            for old_id in old_ids:
                self._segments.pop(old_id, None)

    def delete_segment(self, segment_id: str) -> bool:
        """Remove a segment from the manager.

        Args:
            segment_id: The segment to remove.

        Returns:
            True if the segment was found and removed.
        """
        with self._lock:
            return self._segments.pop(segment_id, None) is not None

    # ------------------------------------------------------------------
    # Lookup
    # ------------------------------------------------------------------

    def get_segments(self) -> List[Segment]:
        """Return a list of all active segments.

        Returns:
            Segments ordered by creation time (oldest first).
        """
        with self._lock:
            segments = list(self._segments.values())
        segments.sort(key=lambda s: s.created_at)
        return segments

    def get_segment(self, segment_id: str) -> Optional[Segment]:
        """Retrieve a segment by its ID.

        Args:
            segment_id: The segment identifier.

        Returns:
            The :class:`Segment` if found, else ``None``.
        """
        with self._lock:
            return self._segments.get(segment_id)

    # ------------------------------------------------------------------
    # Info
    # ------------------------------------------------------------------

    def total_docs(self) -> int:
        """Total number of documents across all active segments."""
        with self._lock:
            return sum(s.doc_count for s in self._segments.values())

    def segment_count(self) -> int:
        """Number of active segments."""
        with self._lock:
            return len(self._segments)

    def __len__(self) -> int:
        return self.segment_count()

    def __contains__(self, segment_id: str) -> bool:
        with self._lock:
            return segment_id in self._segments

    def clear(self) -> None:
        """Remove all segments."""
        with self._lock:
            self._segments.clear()
