"""Tests for rbsp.indexer.warmup -- index preloading and cache warming."""

import os
import sys
import time
import unittest

sys.path.insert(0, os.path.dirname(__file__))
from conftest import RBSPTestCase, SAMPLE_TEXTS

from rbsp.indexer.warmup import IndexWarmer, WarmupConfig, WarmupStats


# ---------------------------------------------------------------------------
# Fake engine for unit-testing warmup without a full RBSPEngine
# ---------------------------------------------------------------------------

class _FakeInverted:
    """Minimal inverted index stub with posting lists."""

    def __init__(self):
        self._index = {
            "python": {"doc1", "doc2"},
            "java": {"doc3"},
            "class": {"doc1", "doc2", "doc3"},
        }

    def term_count(self):
        return len(self._index)

    def doc_count(self):
        all_docs = set()
        for docs in self._index.values():
            all_docs.update(docs)
        return len(all_docs)


class _FakeLSH:
    """Minimal LSH stub with signature storage."""

    def __init__(self):
        self._signatures = {
            "doc1": [1, 2, 3, 4],
            "doc2": [5, 6, 7, 8],
        }

    def __len__(self):
        return len(self._signatures)


class _FakeBM25:
    """Minimal BM25 stub with doc_freqs and idf()."""

    def __init__(self):
        self._doc_freqs = {
            "python": 2,
            "java": 1,
            "class": 3,
        }
        self._num_docs = 3
        self._avg_doc_len = 10.0

    def idf(self, term):
        import math
        df = self._doc_freqs.get(term, 0)
        n = self._num_docs
        return math.log((n - df + 0.5) / (df + 0.5) + 1.0)


class _FakeEngine:
    """Fake engine exposing the attributes that IndexWarmer expects."""

    def __init__(self, with_resolve=True):
        self._inverted = _FakeInverted()
        self._lsh = _FakeLSH()
        self._bm25 = _FakeBM25()
        self._resolve_count = 0
        self._with_resolve = with_resolve

    def resolve(self, query):
        self._resolve_count += 1
        return []

    def __getattr__(self, name):
        if name == "resolve" and not self._with_resolve:
            raise AttributeError(name)
        raise AttributeError(name)


class _MinimalEngine:
    """Engine with no internal structures (all None)."""
    _inverted = None
    _lsh = None
    _bm25 = None


# ---------------------------------------------------------------------------
# WarmupConfig tests
# ---------------------------------------------------------------------------

class TestWarmupConfig(unittest.TestCase):
    """Tests for WarmupConfig dataclass."""

    def test_defaults(self):
        cfg = WarmupConfig()
        self.assertTrue(cfg.preload_postings)
        self.assertTrue(cfg.preload_signatures)
        self.assertTrue(cfg.precompute_idf)
        self.assertEqual(cfg.warmup_queries, [])
        self.assertEqual(cfg.max_warmup_time, 5.0)

    def test_custom_values(self):
        cfg = WarmupConfig(
            preload_postings=False,
            preload_signatures=False,
            precompute_idf=False,
            warmup_queries=["hello", "world"],
            max_warmup_time=2.0,
        )
        self.assertFalse(cfg.preload_postings)
        self.assertFalse(cfg.preload_signatures)
        self.assertFalse(cfg.precompute_idf)
        self.assertEqual(cfg.warmup_queries, ["hello", "world"])
        self.assertEqual(cfg.max_warmup_time, 2.0)

    def test_warmup_queries_default_is_independent(self):
        cfg1 = WarmupConfig()
        cfg2 = WarmupConfig()
        cfg1.warmup_queries.append("test")
        self.assertEqual(cfg2.warmup_queries, [])


# ---------------------------------------------------------------------------
# WarmupStats tests
# ---------------------------------------------------------------------------

class TestWarmupStats(unittest.TestCase):
    """Tests for WarmupStats dataclass."""

    def test_defaults(self):
        stats = WarmupStats()
        self.assertEqual(stats.duration, 0.0)
        self.assertEqual(stats.postings_loaded, 0)
        self.assertEqual(stats.terms_precomputed, 0)
        self.assertEqual(stats.queries_run, 0)

    def test_custom_values(self):
        stats = WarmupStats(
            duration=1.5,
            postings_loaded=100,
            terms_precomputed=50,
            queries_run=3,
        )
        self.assertEqual(stats.duration, 1.5)
        self.assertEqual(stats.postings_loaded, 100)
        self.assertEqual(stats.terms_precomputed, 50)
        self.assertEqual(stats.queries_run, 3)


# ---------------------------------------------------------------------------
# IndexWarmer tests
# ---------------------------------------------------------------------------

class TestIndexWarmer(unittest.TestCase):
    """Tests for IndexWarmer."""

    def test_default_config(self):
        warmer = IndexWarmer()
        self.assertIsNotNone(warmer._config)
        self.assertTrue(warmer._config.preload_postings)

    def test_custom_config(self):
        cfg = WarmupConfig(preload_postings=False)
        warmer = IndexWarmer(cfg)
        self.assertFalse(warmer._config.preload_postings)

    def test_warmup_all_steps(self):
        engine = _FakeEngine()
        cfg = WarmupConfig(warmup_queries=["test query"])
        warmer = IndexWarmer(cfg)
        stats = warmer.warmup(engine)

        self.assertGreater(stats.duration, 0.0)
        self.assertEqual(stats.postings_loaded, 3)  # 3 terms in fake inverted
        self.assertEqual(stats.terms_precomputed, 3)  # 3 terms in fake bm25
        self.assertEqual(stats.queries_run, 1)
        self.assertEqual(engine._resolve_count, 1)

    def test_warmup_postings_only(self):
        engine = _FakeEngine()
        cfg = WarmupConfig(
            preload_postings=True,
            preload_signatures=False,
            precompute_idf=False,
        )
        warmer = IndexWarmer(cfg)
        stats = warmer.warmup(engine)

        self.assertEqual(stats.postings_loaded, 3)
        self.assertEqual(stats.terms_precomputed, 0)
        self.assertEqual(stats.queries_run, 0)

    def test_warmup_signatures_only(self):
        engine = _FakeEngine()
        cfg = WarmupConfig(
            preload_postings=False,
            preload_signatures=True,
            precompute_idf=False,
        )
        warmer = IndexWarmer(cfg)
        stats = warmer.warmup(engine)

        self.assertEqual(stats.postings_loaded, 0)
        self.assertEqual(stats.terms_precomputed, 0)

    def test_warmup_idf_only(self):
        engine = _FakeEngine()
        cfg = WarmupConfig(
            preload_postings=False,
            preload_signatures=False,
            precompute_idf=True,
        )
        warmer = IndexWarmer(cfg)
        stats = warmer.warmup(engine)

        self.assertEqual(stats.postings_loaded, 0)
        self.assertEqual(stats.terms_precomputed, 3)

    def test_warmup_no_steps(self):
        engine = _FakeEngine()
        cfg = WarmupConfig(
            preload_postings=False,
            preload_signatures=False,
            precompute_idf=False,
        )
        warmer = IndexWarmer(cfg)
        stats = warmer.warmup(engine)

        self.assertEqual(stats.postings_loaded, 0)
        self.assertEqual(stats.terms_precomputed, 0)
        self.assertEqual(stats.queries_run, 0)
        self.assertGreaterEqual(stats.duration, 0.0)

    def test_warmup_multiple_queries(self):
        engine = _FakeEngine()
        cfg = WarmupConfig(
            preload_postings=False,
            preload_signatures=False,
            precompute_idf=False,
            warmup_queries=["q1", "q2", "q3"],
        )
        warmer = IndexWarmer(cfg)
        stats = warmer.warmup(engine)

        self.assertEqual(stats.queries_run, 3)
        self.assertEqual(engine._resolve_count, 3)

    def test_warmup_with_minimal_engine(self):
        """Warmup on an engine with no internal structures should not crash."""
        engine = _MinimalEngine()
        warmer = IndexWarmer()
        stats = warmer.warmup(engine)

        self.assertEqual(stats.postings_loaded, 0)
        self.assertEqual(stats.terms_precomputed, 0)
        self.assertEqual(stats.queries_run, 0)

    def test_warmup_respects_time_budget(self):
        """Warmup with a near-zero budget should complete quickly."""
        engine = _FakeEngine()
        cfg = WarmupConfig(max_warmup_time=0.0)
        warmer = IndexWarmer(cfg)
        stats = warmer.warmup(engine)

        # With zero budget some steps may be skipped, but it should not hang
        self.assertLess(stats.duration, 1.0)

    def test_warmup_query_exception_handled(self):
        """If a warmup query raises, it should be caught and not stop warmup."""

        class _FailingEngine:
            _inverted = None
            _lsh = None
            _bm25 = None

            def resolve(self, query):
                raise RuntimeError("query failed")

        engine = _FailingEngine()
        cfg = WarmupConfig(
            preload_postings=False,
            preload_signatures=False,
            precompute_idf=False,
            warmup_queries=["bad query", "also bad"],
        )
        warmer = IndexWarmer(cfg)
        stats = warmer.warmup(engine)

        # Both queries attempted despite exceptions
        self.assertEqual(stats.queries_run, 2)


# ---------------------------------------------------------------------------
# estimate_warmup_time tests
# ---------------------------------------------------------------------------

class TestEstimateWarmupTime(unittest.TestCase):
    """Tests for IndexWarmer.estimate_warmup_time()."""

    def test_estimate_with_fake_engine(self):
        engine = _FakeEngine()
        warmer = IndexWarmer()
        est = warmer.estimate_warmup_time(engine)
        self.assertGreater(est, 0.0)

    def test_estimate_with_empty_engine(self):
        engine = _MinimalEngine()
        warmer = IndexWarmer()
        est = warmer.estimate_warmup_time(engine)
        self.assertEqual(est, 0.0)

    def test_estimate_includes_query_time(self):
        engine = _FakeEngine()
        cfg = WarmupConfig(
            preload_postings=False,
            preload_signatures=False,
            precompute_idf=False,
            warmup_queries=["q1", "q2"],
        )
        warmer = IndexWarmer(cfg)
        est = warmer.estimate_warmup_time(engine)
        # Should include 2 * 0.01 = 0.02 from warmup queries
        self.assertAlmostEqual(est, 0.02, places=3)

    def test_estimate_all_disabled(self):
        engine = _FakeEngine()
        cfg = WarmupConfig(
            preload_postings=False,
            preload_signatures=False,
            precompute_idf=False,
        )
        warmer = IndexWarmer(cfg)
        est = warmer.estimate_warmup_time(engine)
        self.assertEqual(est, 0.0)


# ---------------------------------------------------------------------------
# Internal method tests
# ---------------------------------------------------------------------------

class TestPreloadPostings(unittest.TestCase):
    """Tests for _preload_postings static method."""

    def test_preloads_all_terms(self):
        engine = _FakeEngine()
        count = IndexWarmer._preload_postings(engine)
        self.assertEqual(count, 3)

    def test_no_inverted_index(self):
        engine = _MinimalEngine()
        count = IndexWarmer._preload_postings(engine)
        self.assertEqual(count, 0)


class TestPreloadSignatures(unittest.TestCase):
    """Tests for _preload_signatures static method."""

    def test_preloads_all_sigs(self):
        engine = _FakeEngine()
        count = IndexWarmer._preload_signatures(engine)
        self.assertEqual(count, 2)  # 2 items in fake LSH

    def test_no_lsh(self):
        engine = _MinimalEngine()
        count = IndexWarmer._preload_signatures(engine)
        self.assertEqual(count, 0)


class TestPrecomputeIdf(unittest.TestCase):
    """Tests for _precompute_idf static method."""

    def test_computes_all_terms(self):
        engine = _FakeEngine()
        count = IndexWarmer._precompute_idf(engine)
        self.assertEqual(count, 3)

    def test_no_bm25(self):
        engine = _MinimalEngine()
        count = IndexWarmer._precompute_idf(engine)
        self.assertEqual(count, 0)

    def test_empty_doc_freqs(self):
        engine = _FakeEngine()
        engine._bm25._doc_freqs = {}
        count = IndexWarmer._precompute_idf(engine)
        self.assertEqual(count, 0)


class TestRunWarmupQueries(unittest.TestCase):
    """Tests for _run_warmup_queries static method."""

    def test_runs_all_queries(self):
        engine = _FakeEngine()
        deadline = time.monotonic() + 10.0
        count = IndexWarmer._run_warmup_queries(engine, ["a", "b"], deadline)
        self.assertEqual(count, 2)
        self.assertEqual(engine._resolve_count, 2)

    def test_no_resolve_method(self):
        engine = _MinimalEngine()
        deadline = time.monotonic() + 10.0
        count = IndexWarmer._run_warmup_queries(engine, ["a"], deadline)
        self.assertEqual(count, 0)

    def test_expired_deadline(self):
        engine = _FakeEngine()
        deadline = time.monotonic() - 1.0  # already past
        count = IndexWarmer._run_warmup_queries(engine, ["a", "b"], deadline)
        self.assertEqual(count, 0)


# ---------------------------------------------------------------------------
# Package-level import test
# ---------------------------------------------------------------------------

class TestPackageImport(unittest.TestCase):
    """Verify that warmup classes are exported from rbsp.indexer."""

    def test_import_from_package(self):
        from rbsp.indexer import IndexWarmer, WarmupConfig, WarmupStats
        self.assertIsNotNone(IndexWarmer)
        self.assertIsNotNone(WarmupConfig)
        self.assertIsNotNone(WarmupStats)


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