import json
from pathlib import Path
from typing import Any

import pytest

from app.adapters.fake import FakeEmbeddingProvider, FakeReranker

PROJECT_ROOT = Path(__file__).resolve().parents[3]
SAMPLE_ROOT = PROJECT_ROOT / "data" / "samples" / "public"


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]


def cosine(left: tuple[float, ...], right: tuple[float, ...]) -> float:
    return sum(
        left_value * right_value for left_value, right_value in zip(left, right, strict=True)
    )


@pytest.mark.asyncio
async def test_synthetic_questions_retrieve_expected_document_after_rerank() -> None:
    documents = load_jsonl(SAMPLE_ROOT / "demo_documents.jsonl")
    queries = load_jsonl(SAMPLE_ROOT / "demo_queries.jsonl")
    embedder = FakeEmbeddingProvider()
    reranker = FakeReranker()
    document_texts = [f"{item['title']}\n{item['content']}" for item in documents]
    document_vectors: list[tuple[float, ...]] = []
    for offset in range(0, len(document_texts), 10):
        batch = await embedder.embed_documents(document_texts[offset : offset + 10])
        document_vectors.extend(batch.vectors)

    hits = 0
    answerable_queries = [query for query in queries if query["answerable"]]
    for query in answerable_queries:
        query_vector = (await embedder.embed_query(query["query"])).vectors[0]
        candidate_indexes = sorted(
            range(len(documents)),
            key=lambda index: cosine(query_vector, document_vectors[index]),
            reverse=True,
        )[:5]
        candidate_texts = [document_texts[index] for index in candidate_indexes]
        reranked = await reranker.rerank(query["query"], candidate_texts, top_n=3)
        result_ids = [documents[candidate_indexes[item.index]]["doc_id"] for item in reranked.items]
        if set(query["expected_doc_ids"]) & set(result_ids):
            hits += 1

    assert hits / len(answerable_queries) >= 0.8