"""Tests for compiler/stages/index.py — offline using a stub embedder."""

from __future__ import annotations

import hashlib

import numpy as np
from compiler.stages.index import _sub_chunk, build_index, query_index

DIM = 32  # tiny embedding dim for tests


class _HashEmbedder:
    """Deterministic stub: text → fixed-dim unit vector derived from content hash."""

    def encode(
        self,
        sentences: list[str],
        *,
        show_progress_bar: bool = False,
        normalize_embeddings: bool = False,
    ) -> np.ndarray:
        vecs = []
        for s in sentences:
            h = hashlib.sha256(s.encode()).digest()
            # Repeat hash bytes to fill DIM floats
            raw = (h * (DIM // len(h) + 1))[:DIM]
            vec = np.frombuffer(raw, dtype=np.uint8).astype(np.float32) - 128.0
            if normalize_embeddings:
                norm = np.linalg.norm(vec)
                if norm > 0:
                    vec /= norm
            vecs.append(vec)
        return np.stack(vecs)


FAKE_EMBEDDER = _HashEmbedder()

_PASSAGES = [
    "You have power over your mind, not outside events."
    " Realize this, and you will find strength.",
    "The impediment to action advances action. What stands in the way becomes the way.",
    "Very little is needed to make a happy life;"
    " it is all within yourself in your way of thinking.",
    "Accept the things to which fate binds you"
    " and love the people with whom fate brings you together.",
    "Dwell on the beauty of life. Watch the stars and see yourself running with them.",
]
INDEXABLE = {"stoic-texts": _PASSAGES}


# ── _sub_chunk ───────────────────────────────────────────────────────────────


def test_sub_chunk_short_passthrough():
    text = "Short paragraph."
    assert _sub_chunk(text) == [text]


def test_sub_chunk_splits_long_text():
    long_text = "Sentence one is here. " * 100
    chunks = _sub_chunk(long_text, max_chars=200)
    assert len(chunks) > 1
    for chunk in chunks:
        assert len(chunk) <= 300


# ── build_index + query_index ─────────────────────────────────────────────────


def test_build_index_creates_table(tmp_path):
    build_index(INDEXABLE, tmp_path / "index", embedder=FAKE_EMBEDDER, overwrite=True)
    import lancedb

    db = lancedb.connect(str(tmp_path / "index"))
    assert "chunks" in db.list_tables().tables
    table = db.open_table("chunks")
    assert table.count_rows() == len(INDEXABLE["stoic-texts"])


def test_query_index_returns_results(tmp_path):
    index_dir = tmp_path / "index"
    build_index(INDEXABLE, index_dir, embedder=FAKE_EMBEDDER, overwrite=True)
    results = query_index("mind and power", index_dir, top_k=3, embedder=FAKE_EMBEDDER)
    assert len(results) > 0
    assert "chunk_id" in results[0]
    assert "source_id" in results[0]
    assert "text" in results[0]


def test_query_index_empty_returns_empty(tmp_path):
    results = query_index(
        "anything", tmp_path / "index", embedder=FAKE_EMBEDDER
    )
    assert results == []


def test_build_index_overwrite(tmp_path):
    index_dir = tmp_path / "index"
    build_index(INDEXABLE, index_dir, embedder=FAKE_EMBEDDER, overwrite=True)
    build_index(INDEXABLE, index_dir, embedder=FAKE_EMBEDDER, overwrite=True)
    results = query_index("virtue", index_dir, top_k=2, embedder=FAKE_EMBEDDER)
    assert len(results) > 0


def test_build_index_empty_corpus_no_error(tmp_path):
    build_index({}, tmp_path / "index", embedder=FAKE_EMBEDDER)


def test_chunk_ids_are_unique(tmp_path):
    index_dir = tmp_path / "index"
    build_index(INDEXABLE, index_dir, embedder=FAKE_EMBEDDER, overwrite=True)
    import lancedb

    db = lancedb.connect(str(index_dir))
    table = db.open_table("chunks")
    rows = table.search([0.0] * DIM).limit(1000).to_list()
    chunk_ids = [r["chunk_id"] for r in rows]
    assert len(set(chunk_ids)) == len(chunk_ids)