"""Index preloading and cache warming for fast first-query performance.

When an index is loaded from disk, the first query can be slow because
caches are cold.  This module pre-populates hot data structures so the
first user-visible query runs at full speed.

Usage::

    from rbsp.indexer.warmup import IndexWarmer, WarmupConfig

    warmer = IndexWarmer(WarmupConfig(warmup_queries=["auth login"]))
    stats = warmer.warmup(engine)
    print(f"Warmup done in {stats.duration:.2f}s")
"""

from __future__ import annotations

import math
import time
from dataclasses import dataclass, field
from typing import Any, List


@dataclass
class WarmupConfig:
    """Configuration for index warmup behaviour.

    Attributes:
        preload_postings:  Iterate all inverted-index posting lists to
            force them into memory/CPU cache.
        preload_signatures:  Touch all LSH signatures so they are
            resident in memory.
        precompute_idf:  Pre-compute IDF values for the most frequent
            terms in the corpus.
        warmup_queries:  A list of representative queries to execute
            during warmup so that all query-path caches are hot.
        max_warmup_time:  Hard time budget (seconds).  Warmup steps
            that exceed this budget are skipped.
    """

    preload_postings: bool = True
    preload_signatures: bool = True
    precompute_idf: bool = True
    warmup_queries: List[str] = field(default_factory=list)
    max_warmup_time: float = 5.0


@dataclass
class WarmupStats:
    """Statistics collected during a warmup run.

    Attributes:
        duration:           Total wall-clock seconds spent on warmup.
        postings_loaded:    Number of posting lists touched.
        terms_precomputed:  Number of terms whose IDF was pre-computed.
        queries_run:        Number of warmup queries executed.
    """

    duration: float = 0.0
    postings_loaded: int = 0
    terms_precomputed: int = 0
    queries_run: int = 0


class IndexWarmer:
    """Pre-populates hot data structures on an :class:`RBSPEngine`.

    Args:
        config: Optional warmup configuration.  Uses defaults if *None*.
    """

    def __init__(self, config: WarmupConfig | None = None) -> None:
        self._config = config or WarmupConfig()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def warmup(self, engine: Any) -> WarmupStats:
        """Run the full warmup sequence on *engine*.

        Executes each enabled warmup step in order, respecting
        :attr:`WarmupConfig.max_warmup_time`.

        Args:
            engine: An :class:`~rbsp.runtime.engine.RBSPEngine` instance.

        Returns:
            Statistics about what was warmed up.
        """
        stats = WarmupStats()
        start = time.monotonic()
        deadline = start + self._config.max_warmup_time

        if self._config.preload_postings and time.monotonic() < deadline:
            stats.postings_loaded = self._preload_postings(engine)

        if self._config.preload_signatures and time.monotonic() < deadline:
            self._preload_signatures(engine)

        if self._config.precompute_idf and time.monotonic() < deadline:
            stats.terms_precomputed = self._precompute_idf(engine)

        if self._config.warmup_queries and time.monotonic() < deadline:
            stats.queries_run = self._run_warmup_queries(
                engine, self._config.warmup_queries, deadline,
            )

        stats.duration = time.monotonic() - start
        return stats

    def estimate_warmup_time(self, engine: Any) -> float:
        """Estimate how long warmup will take (seconds).

        Provides a rough heuristic based on the number of indexed
        documents and terms.  Useful for showing progress indicators.

        Args:
            engine: An :class:`~rbsp.runtime.engine.RBSPEngine` instance.

        Returns:
            Estimated warmup duration in seconds.
        """
        estimate = 0.0

        inverted = getattr(engine, "_inverted", None)
        lsh = getattr(engine, "_lsh", None)

        if self._config.preload_postings and inverted is not None:
            # ~1 us per posting list on modern hardware
            estimate += inverted.term_count() * 1e-6

        if self._config.preload_signatures and lsh is not None:
            # ~0.5 us per signature read
            estimate += len(lsh) * 0.5e-6

        if self._config.precompute_idf and inverted is not None:
            # IDF is a log computation per term
            estimate += inverted.term_count() * 2e-6

        # Each warmup query is roughly equivalent to a normal resolve
        estimate += len(self._config.warmup_queries) * 0.01

        return estimate

    # ------------------------------------------------------------------
    # Internal warmup steps
    # ------------------------------------------------------------------

    @staticmethod
    def _preload_postings(engine: Any) -> int:
        """Iterate all posting lists to force them into memory/cache.

        Touches every term's posting set in the inverted index so that
        subsequent lookups hit warm memory pages.

        Returns:
            Number of posting lists loaded.
        """
        inverted = getattr(engine, "_inverted", None)
        if inverted is None:
            return 0

        count = 0
        # Access the internal _index dict to touch every posting set.
        index_data = getattr(inverted, "_index", None)
        if index_data is not None:
            for _term, doc_set in index_data.items():
                # Force the set into cache by reading its length.
                len(doc_set)
                count += 1
        return count

    @staticmethod
    def _preload_signatures(engine: Any) -> int:
        """Touch all LSH signatures to pull them into memory.

        Returns:
            Number of signatures loaded.
        """
        lsh = getattr(engine, "_lsh", None)
        if lsh is None:
            return 0

        sigs = getattr(lsh, "_signatures", None)
        if sigs is None:
            return 0

        count = 0
        for _item_id, sig in sigs.items():
            # Touch the first and last element to fault in the list.
            if sig:
                _ = sig[0]
                _ = sig[-1]
            count += 1
        return count

    @staticmethod
    def _precompute_idf(engine: Any) -> int:
        """Pre-compute IDF values for all terms in the inverted index.

        This forces the BM25 scorer's IDF code path through the CPU
        cache and validates that corpus statistics are ready.

        Returns:
            Number of terms whose IDF was computed.
        """
        bm25 = getattr(engine, "_bm25", None)
        inverted = getattr(engine, "_inverted", None)
        if bm25 is None or inverted is None:
            return 0

        doc_freqs = getattr(bm25, "_doc_freqs", None)
        if not doc_freqs:
            return 0

        count = 0
        for term in doc_freqs:
            # Call idf() to force computation and warm any internal
            # caches in the math library / CPU branch predictor.
            bm25.idf(term)
            count += 1
        return count

    @staticmethod
    def _run_warmup_queries(
        engine: Any,
        queries: List[str],
        deadline: float,
    ) -> int:
        """Execute warmup queries to exercise the full query pipeline.

        Stops early if the deadline is reached.

        Returns:
            Number of queries actually executed.
        """
        resolve = getattr(engine, "resolve", None)
        if resolve is None:
            return 0

        count = 0
        for query in queries:
            if time.monotonic() >= deadline:
                break
            try:
                resolve(query)
            except Exception:
                pass
            count += 1
        return count
