"""Unit tests for the RAG building blocks: chunking, parsing, embeddings, store."""

from __future__ import annotations

import httpx
import pytest
from legate.rag.chunking import chunk_text
from legate.rag.embeddings.hashing import HashingEmbedder
from legate.rag.embeddings.openai import OpenAIEmbedder
from legate.rag.parsing import UnsupportedDocumentError, extract_text
from legate.rag.vectorstore.base import VectorRecord
from legate.rag.vectorstore.memory import InMemoryVectorStore, _cosine


def test_chunk_text_short_and_empty():
    assert chunk_text("") == []
    assert chunk_text("a short paragraph") == ["a short paragraph"]


def test_chunk_text_splits_long_input():
    para = "word " * 60  # ~300 chars
    text = "\n\n".join([para] * 10)  # ~3000 chars
    chunks = chunk_text(text, chunk_size=500, overlap=50)
    assert len(chunks) > 1
    assert all(len(c) <= 500 + 50 + 10 for c in chunks)
    assert "word" in chunks[0]


def test_extract_text_txt_md_and_unsupported():
    assert extract_text("notes.txt", b"hello world") == "hello world"
    assert extract_text("readme.md", b"# Title") == "# Title"
    with pytest.raises(UnsupportedDocumentError):
        extract_text("data.xyz", b"...")


def test_hashing_embedder_is_deterministic_and_relevant():
    emb = HashingEmbedder(dim=256)
    a1 = emb.embed(["refund policy for customers"])[0]
    a2 = emb.embed(["refund policy for customers"])[0]
    assert a1 == a2  # deterministic

    ref = emb.embed(["the refund policy explains customer refunds"])[0]
    close = emb.embed(["how do refunds work"])[0]
    far = emb.embed(["airplane turbine maintenance schedule"])[0]
    assert _cosine(ref, close) > _cosine(ref, far)


def test_openai_embedder_parses_and_orders_response():
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={
                "data": [
                    {"index": 1, "embedding": [0.0, 1.0]},
                    {"index": 0, "embedding": [9.0, 9.0]},
                ]
            },
        )

    with httpx.Client(transport=httpx.MockTransport(handler)) as http:
        emb = OpenAIEmbedder(api_key="k", http_client=http)
        vecs = emb.embed(["first", "second"])
    assert vecs == [[9.0, 9.0], [0.0, 1.0]]  # reordered by index
    assert emb.embed([]) == []


def test_in_memory_vector_store_search_and_delete():
    store = InMemoryVectorStore()
    store.upsert(
        "col",
        [
            VectorRecord(id="1", vector=[1.0, 0.0], payload={"document_id": "d1", "text": "a"}),
            VectorRecord(id="2", vector=[0.0, 1.0], payload={"document_id": "d2", "text": "b"}),
        ],
    )
    hits = store.search("col", [1.0, 0.0], top_k=2)
    assert hits[0].id == "1" and hits[0].score > hits[1].score

    assert store.delete_by_document("col", "d1") == 1
    assert [h.id for h in store.search("col", [1.0, 0.0], top_k=5)] == ["2"]

    store.delete_collection("col")
    assert store.search("col", [1.0, 0.0]) == []
