"""Tests for EPUB search core helpers."""

from __future__ import annotations

from dataclasses import replace
from datetime import UTC, datetime
from os import environ
from pathlib import Path
from threading import Event
from types import ModuleType

import pytest
from sqlalchemy import create_engine, select
from sqlalchemy.orm import sessionmaker

from python.ebook_search.answer import answer_query
from python.ebook_search.bm25_corpus import (
    BM25Corpus,
    BM25CorpusUnavailableError,
    BM25Manifest,
    ensure_bm25_corpus,
    load_bm25_corpus,
)
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config, normalize_embedding_model
from python.ebook_search.embeddings import MODEL_DIMENSIONS, ensure_embedding_models
from python.ebook_search.ingest import chunk_text, find_existing_source
from python.ebook_search.search import (
    SearchResponse,
    SearchResult,
    bm25_candidates,
    reciprocal_rank_fusion,
    retrieval_query_from_text,
    search_ebooks,
)
from python.ebook_search.timing import RuntimeStep
from python.orm.richie import EbookEmbeddingModel, EbookSource, RichieBase


def test_chunk_text_uses_overlap() -> None:
    chunks = chunk_text(" ".join(str(index) for index in range(100)), chunk_tokens=20, overlap_tokens=5)

    assert len(chunks) > 1
    assert chunks[0].token_start == 0
    assert chunks[1].token_start == 15
    assert all(chunk.token_count <= 20 for chunk in chunks)


def test_reciprocal_rank_fusion_combines_vector_and_bm25_rankings() -> None:
    vector_results = [
        SearchResult(chunk_id=1, text="a", source_title="A", score=0.9, vector_score=0.9),
        SearchResult(chunk_id=2, text="b", source_title="B", score=0.8, vector_score=0.8),
    ]
    lexical_results = [
        SearchResult(chunk_id=2, text="b", source_title="B", score=4.2, bm25_score=4.2),
        SearchResult(chunk_id=3, text="c", source_title="C", score=2.1, bm25_score=2.1),
    ]

    fused = reciprocal_rank_fusion(vector_results, lexical_results)

    assert [result.chunk_id for result in fused] == [2, 1, 3]
    assert fused[0].rank_source == "Hybrid"
    assert fused[0].vector_score == 0.8
    assert fused[0].bm25_score == 4.2
    assert fused[0].fused_score == fused[0].score


def test_find_existing_source_matches_path_or_hash() -> None:
    engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
    RichieBase.metadata.create_all(engine)
    with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
        source = EbookSource(
            title="Book",
            author=None,
            language=None,
            publisher=None,
            identifier=None,
            file_path="/old/book.epub",
            file_sha256="a" * 64,
            file_mtime=datetime.now(tz=UTC),
            file_size=10,
        )
        session.add(source)
        session.commit()

        assert find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source
        assert find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source


def test_reciprocal_rank_fusion_marks_hybrid_source() -> None:
    vector_results = [SearchResult(chunk_id=1, text="a", source_title="A")]
    lexical_results = [SearchResult(chunk_id=2, text="b", source_title="B")]

    fused = reciprocal_rank_fusion(vector_results, lexical_results)

    assert {result.rank_source for result in fused} == {"Hybrid"}


def test_search_response_sums_runtime_steps() -> None:
    response = SearchResponse(
        query="query",
        results=[],
        rank_label="Hybrid",
        timings=(
            RuntimeStep(name="A", duration_ms=1.25),
            RuntimeStep(name="B", duration_ms=2.75),
            RuntimeStep(name="Parallel detail", duration_ms=10.0, counts_toward_total=False),
        ),
    )

    assert response.total_runtime_ms == 4.0


def test_search_ebooks_runs_vector_and_bm25_in_parallel(monkeypatch) -> None:
    engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
    vector_started = Event()
    bm25_started = Event()
    received_engines: list[object] = []

    def fake_vector_candidates(received_engine, query, _config):
        """Return vector candidates after confirming BM25 has started."""
        received_engines.append(received_engine)
        assert query == "parallel"
        vector_started.set()
        assert bm25_started.wait(timeout=2)
        return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]

    def fake_bm25_candidates(query, _config):
        """Return BM25 candidates after confirming vector search has started."""
        assert query == "parallel"
        bm25_started.set()
        assert vector_started.wait(timeout=2)
        return [SearchResult(chunk_id=2, text="bm25", source_title="BM25", bm25_score=2.0)]

    monkeypatch.setattr("python.ebook_search.search.vector_candidates", fake_vector_candidates)
    monkeypatch.setattr("python.ebook_search.search.bm25_candidates", fake_bm25_candidates)
    config = EbookSearchConfig(rerank=RerankConfig(enabled=False))

    response = search_ebooks(engine, "parallel", config)

    timings = {step.name: step for step in response.timings}
    assert [result.chunk_id for result in response.results] == [1, 2]
    assert timings["Embedding + vector search"].counts_toward_total is False
    assert timings["BM25 search"].counts_toward_total is False
    assert timings["Hybrid retrieval"].counts_toward_total is True
    assert received_engines == [engine]


def test_retrieval_query_keeps_entity_and_series_terms() -> None:
    assert retrieval_query_from_text("what does Damien Montgomery stand for in starship mage") == (
        "damien montgomery stand starship mage"
    )


def test_bm25_candidates_scores_whole_corpus(monkeypatch) -> None:
    record = {
        "chunk_id": 2,
        "text": "high",
        "source_title": "B",
        "source_author": None,
        "chapter_title": None,
        "page_label": None,
        "bm25_text": "high",
    }
    manifest = BM25Manifest(created_at=datetime.now(tz=UTC), db_updated_at=None, chunk_count=1)
    corpus = BM25Corpus(retriever=object(), records=(record,), manifest=manifest)
    captured: dict[str, object] = {}

    def fake_score_bm25_corpus(query, saved_corpus, *, limit):
        captured["query"] = query
        captured["corpus"] = saved_corpus
        captured["limit"] = limit
        return [(record, 1.5)]

    monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", lambda _config: corpus)
    monkeypatch.setattr("python.ebook_search.search.score_bm25_corpus", fake_score_bm25_corpus)
    config = EbookSearchConfig(rerank=RerankConfig(enabled=False))

    results = bm25_candidates("high", config)

    assert captured["query"] == "high"
    assert captured["corpus"] == corpus
    assert captured["limit"] == 120
    assert [result.chunk_id for result in results] == [2]
    assert [result.bm25_score for result in results] == [1.5]


def test_bm25_candidates_raises_when_corpus_is_unavailable(monkeypatch) -> None:
    def fake_load_bm25_corpus(_config):
        raise BM25CorpusUnavailableError

    monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", fake_load_bm25_corpus)
    config = EbookSearchConfig(rerank=RerankConfig(enabled=False))

    with pytest.raises(BM25CorpusUnavailableError):
        bm25_candidates("high", config)


def test_load_bm25_corpus_caches_disk_load(monkeypatch, tmp_path) -> None:
    load_bm25_corpus.cache_clear()
    manifest = BM25Manifest(created_at=datetime.now(tz=UTC), db_updated_at=None, chunk_count=1)
    record = {
        "chunk_id": 2,
        "text": "cached",
        "source_title": "B",
        "source_author": None,
        "chapter_title": None,
        "page_label": None,
        "bm25_text": "cached",
    }
    load_count = 0

    class FakeRetriever:
        """Fake persisted BM25 retriever."""

        corpus = (record,)

    class FakeBM25:
        """Fake BM25 class with observable load count."""

        @staticmethod
        def load(index_path, *, load_corpus, mmap):
            nonlocal load_count
            load_count += 1
            assert index_path == tmp_path
            assert load_corpus is True
            assert mmap is True
            return FakeRetriever()

    fake_bm25s = ModuleType("bm25s")
    fake_bm25s.BM25 = FakeBM25
    monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: manifest)
    monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: True)
    monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25s", fake_bm25s)
    config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(tmp_path))

    try:
        first = load_bm25_corpus(config)
        second = load_bm25_corpus(config)
    finally:
        load_bm25_corpus.cache_clear()

    assert first is second
    assert first is not None
    assert first.records == (record,)
    assert load_count == 1


def test_load_bm25_corpus_raises_when_index_is_missing(monkeypatch, tmp_path) -> None:
    load_bm25_corpus.cache_clear()
    monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: None)
    monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: False)
    config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(tmp_path))

    try:
        with pytest.raises(BM25CorpusUnavailableError, match="BM25 corpus is not available"):
            load_bm25_corpus(config)
    finally:
        load_bm25_corpus.cache_clear()


def test_ensure_bm25_corpus_refreshes_missing_index(monkeypatch) -> None:
    refreshed: list[object] = []
    db_updated_at = datetime.now(tz=UTC)

    monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: None)
    monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: False)
    monkeypatch.setattr("python.ebook_search.bm25_corpus.corpus_last_updated_at", lambda _session: db_updated_at)
    monkeypatch.setattr(
        "python.ebook_search.bm25_corpus.refresh_bm25_corpus",
        lambda session, config, *, db_updated_at: refreshed.append((session, config, db_updated_at)),
    )

    config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
    session = object()

    ensure_bm25_corpus(session, config)

    assert refreshed == [(session, config, db_updated_at)]


def test_ensure_bm25_corpus_refreshes_stale_index(monkeypatch) -> None:
    refreshed: list[object] = []
    created_at = datetime(2026, 1, 1, tzinfo=UTC)
    db_updated_at = datetime(2026, 1, 2, tzinfo=UTC)
    manifest = BM25Manifest(created_at=created_at, db_updated_at=created_at, chunk_count=10)

    monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: manifest)
    monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: True)
    monkeypatch.setattr("python.ebook_search.bm25_corpus.corpus_last_updated_at", lambda _session: db_updated_at)
    monkeypatch.setattr(
        "python.ebook_search.bm25_corpus.refresh_bm25_corpus",
        lambda session, config, *, db_updated_at: refreshed.append((session, config, db_updated_at)),
    )

    config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
    session = object()

    ensure_bm25_corpus(session, config)

    assert refreshed == [(session, config, db_updated_at)]


def test_supported_embedding_models_match_service_names() -> None:
    assert MODEL_DIMENSIONS == {
        "qwen3-embedding-0.6b": 1024,
        "qwen3-embedding-4b": 2560,
        "qwen3-embedding-8b": 4096,
    }


def test_ensure_embedding_models_registers_service_names() -> None:
    engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
    RichieBase.metadata.create_all(engine)
    with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
        ensure_embedding_models(session)
        session.commit()

        models = list(session.scalars(select(EbookEmbeddingModel).order_by(EbookEmbeddingModel.name)))

    assert [(model.name, model.dimension) for model in models] == [
        ("qwen3-embedding-0.6b", 1024),
        ("qwen3-embedding-4b", 2560),
        ("qwen3-embedding-8b", 4096),
    ]


def test_embedding_model_aliases_normalize_to_provider_names() -> None:
    assert normalize_embedding_model() == "qwen3-embedding-0.6b"

    environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-0.6b"
    assert normalize_embedding_model() == "qwen3-embedding-0.6b"

    environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen3-Embedding-0.6B"
    assert normalize_embedding_model() == "qwen3-embedding-0.6b"

    environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen/Qwen3-Embedding-4B"

    assert normalize_embedding_model() == "qwen3-embedding-4b"

    environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding:8b"
    assert normalize_embedding_model() == "qwen3-embedding-8b"

    environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-8b"
    assert normalize_embedding_model() == "qwen3-embedding-8b"


def test_answer_generation_is_enabled_by_default(monkeypatch) -> None:
    monkeypatch.delenv("EBOOK_SEARCH_ANSWER_ENABLED", raising=False)

    config = load_config()

    assert config.answer_enabled is True


def test_chat_defaults_use_ollama_cloud(monkeypatch) -> None:
    monkeypatch.delenv("EBOOK_SEARCH_VLLM_BASE_URL", raising=False)
    monkeypatch.delenv("EBOOK_SEARCH_CHAT_MODEL", raising=False)

    config = load_config()

    assert config.vllm_base_url == "https://ollama.com/v1"
    assert config.chat_model == "deepseek-v4-flash"


def test_chat_api_key_falls_back_to_ollama_api_key(monkeypatch) -> None:
    monkeypatch.delenv("EBOOK_SEARCH_VLLM_API_KEY", raising=False)
    monkeypatch.setenv("OLLAMA_API_KEY", "ollama-key")

    config = load_config()

    assert config.vllm_api_key == "ollama-key"


def test_answer_query_does_not_call_model_when_disabled() -> None:
    config = replace(load_config(), answer_enabled=False)
    result = SearchResult(chunk_id=1, text="source text", source_title="Book")

    answer = answer_query("question", [result], config)

    assert "Answer generation is disabled" in answer