"""Unit tests for coalex.ext.reranker."""

from __future__ import annotations

import json

from coalex.ext.reranker import RerankerSpan, reranker_span
from coalex.ext.retrieval import Document


def test_reranker_span_decorator_sets_kind(memory_exporter):
    @reranker_span(model_name="rerank-v3", top_k=2)
    def rerank(query: str, docs: list[Document]) -> list[Document]:
        return docs[:2]

    docs = [Document(content="a", score=0.9), Document(content="b", score=0.7)]
    rerank(query="q", docs=docs)
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["openinference.span.kind"] == "RERANKER"


def test_reranker_span_decorator_records_model_and_top_k(memory_exporter):
    @reranker_span(model_name="rerank-english-v3", top_k=3)
    def rerank(query: str, docs: list[Document]) -> list[Document]:
        return docs

    rerank(query="q", docs=[Document(content="x")])
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["reranker.model_name"] == "rerank-english-v3"
    assert spans[0].attributes["reranker.top_k"] == 3


def test_reranker_span_decorator_encodes_input_output(memory_exporter):
    @reranker_span()
    def rerank(query: str, docs: list[Document]) -> list[Document]:
        return [docs[0]]

    input_docs = [Document(content="a", id="1"), Document(content="b", id="2")]
    rerank("q", input_docs)
    spans = memory_exporter.get_finished_spans()
    in_docs = json.loads(spans[0].attributes["reranker.input_documents"])
    out_docs = json.loads(spans[0].attributes["reranker.output_documents"])
    assert len(in_docs) == 2
    assert len(out_docs) == 1
    assert out_docs[0]["document.id"] == "1"


def test_reranker_span_context_manager(memory_exporter):
    input_docs = [Document(content="doc1", id="d1", score=0.5)]
    output_docs = [Document(content="doc1", id="d1", score=0.95)]

    with RerankerSpan("rerank_step", model_name="rerank-v3", top_k=1) as r:
        r.set_query("test query")
        r.set_input_documents(input_docs)
        r.set_output_documents(output_docs)

    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["openinference.span.kind"] == "RERANKER"
    assert spans[0].attributes["reranker.query"] == "test query"
    assert spans[0].attributes["reranker.model_name"] == "rerank-v3"
    assert spans[0].attributes["reranker.top_k"] == 1
    out = json.loads(spans[0].attributes["reranker.output_documents"])
    assert out[0]["document.score"] == 0.95


def test_reranker_span_docs_arg_kwarg_name(memory_exporter):
    @reranker_span(docs_arg="candidates")
    def rerank(query: str, candidates: list[Document]) -> list[Document]:
        return candidates[:1]

    docs = [Document(content="a", id="1"), Document(content="b", id="2")]
    rerank(query="q", candidates=docs)
    spans = memory_exporter.get_finished_spans()
    in_docs = json.loads(spans[0].attributes["reranker.input_documents"])
    assert len(in_docs) == 2


async def test_reranker_span_async(memory_exporter):
    @reranker_span(model_name="rerank-v3", top_k=1)
    async def rerank(query: str, documents: list[Document]) -> list[Document]:
        return documents[:1]

    docs = [Document(content="a", id="1"), Document(content="b", id="2")]
    result = await rerank(query="test", documents=docs)
    assert len(result) == 1
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["openinference.span.kind"] == "RERANKER"
    assert spans[0].attributes["reranker.model_name"] == "rerank-v3"
    in_docs = json.loads(spans[0].attributes["reranker.input_documents"])
    assert len(in_docs) == 2
    out_docs = json.loads(spans[0].attributes["reranker.output_documents"])
    assert len(out_docs) == 1
