"""Tests for index building."""

from __future__ import annotations

import numpy as np

from wikikb import config, embed, indexer

_DOC_A = """---
title: "A"
url: https://x/A
revid: 1
categories: []
---

# A

Genius trait gives plus five diplomacy.
"""

_DOC_B = """---
title: "B"
url: https://x/B
revid: 2
categories: []
---

# B

Herculean physique gives health.
"""


def _embed_fn(fake_model):
    return lambda texts: embed.embed_passages(texts, model=fake_model)


def test_collect_chunks_empty(tmp_kb) -> None:
    assert indexer.collect_chunks(config.CORPUS_DIR) == []


def test_build_index_creates_db_and_embeddings(tmp_kb, fake_model) -> None:
    (config.CORPUS_DIR / "a.md").write_text(_DOC_A, encoding="utf-8")
    (config.CORPUS_DIR / "b.md").write_text(_DOC_B, encoding="utf-8")
    count = indexer.build_index(embed_fn=_embed_fn(fake_model), log=lambda _: None)
    assert count >= 2
    assert config.DB_PATH.exists()
    arr = np.load(config.EMBEDDINGS_PATH)
    assert arr.shape == (count, config.EMBED_DIM)


def test_build_index_empty_corpus(tmp_kb, fake_model) -> None:
    count = indexer.build_index(embed_fn=_embed_fn(fake_model), log=lambda _: None)
    assert count == 0
    arr = np.load(config.EMBEDDINGS_PATH)
    assert arr.shape == (0, config.EMBED_DIM)


def test_build_index_embeds_context_augmented_text(tmp_kb, fake_model) -> None:
    (config.CORPUS_DIR / "a.md").write_text(_DOC_A, encoding="utf-8")
    captured: list[str] = []

    def capture_embed(texts):
        captured.extend(texts)
        return embed.embed_passages(texts, model=fake_model)

    indexer.build_index(embed_fn=capture_embed, log=lambda _: None)
    assert captured
    # embedded text is prefixed with the page title (context), body stays clean in DB
    assert all(t.startswith("A. ") for t in captured)


def test_embed_efficiently_returns_vectors_in_original_order(fake_model) -> None:
    # Worthless if it corrupts alignment: embedding row i must still be rowid i+1
    # after the internal length sort and dedupe.
    texts = ["a" * 500, "bb", "a" * 500, "ccc", "d" * 900]
    got = indexer.embed_efficiently(texts, _embed_fn(fake_model))
    want = embed.embed_passages(texts, model=fake_model)
    np.testing.assert_allclose(got, want)


def test_embed_efficiently_embeds_each_distinct_text_once(fake_model) -> None:
    seen: list[list[str]] = []

    def counting_embed(texts):
        seen.append(list(texts))
        return embed.embed_passages(texts, model=fake_model)

    got = indexer.embed_efficiently(["dup", "dup", "dup", "other"], counting_embed)
    assert seen == [["dup", "other"]]  # deduped, then sorted by length
    assert got.shape[0] == 4
    np.testing.assert_allclose(got[0], got[2])


def test_embed_efficiently_empty(fake_model) -> None:
    assert indexer.embed_efficiently([], _embed_fn(fake_model)).shape == (0, config.EMBED_DIM)


def test_noise_sections_are_not_indexed(tmp_kb, fake_model, monkeypatch) -> None:
    monkeypatch.setattr(config, "NOISE_SECTIONS", frozenset({"other languages"}))
    doc = _DOC_A + "\n## Other languages\n\nlocalized UI names table\n"
    (config.CORPUS_DIR / "a.md").write_text(doc, encoding="utf-8")
    chunks = indexer.collect_chunks(config.CORPUS_DIR)
    assert chunks
    assert not any("localized UI names" in c.text for c in chunks)


def test_main_delegates(monkeypatch) -> None:
    monkeypatch.setattr(indexer, "build_index", lambda: 3)
    assert indexer.main() == 0