"""Unit tests for coalex.ext.retrieval."""

from __future__ import annotations

import json

import coalex
import pytest
from coalex.ext.retrieval import Document, RetrievalSpan, encode_documents, retrieval_span
from opentelemetry.trace import StatusCode

# ---------------------------------------------------------------------------
# Document dataclass
# ---------------------------------------------------------------------------


def test_document_to_openinference_dict_omits_none_fields():
    doc = Document(content="hello")
    d = doc.to_openinference_dict()
    assert d == {"document.content": "hello"}
    assert "document.id" not in d
    assert "document.score" not in d
    assert "document.metadata" not in d


def test_document_to_openinference_dict_includes_score():
    doc = Document(content="hello", id="doc1", score=0.95)
    d = doc.to_openinference_dict()
    assert d["document.id"] == "doc1"
    assert d["document.score"] == 0.95


def test_encode_documents_json_roundtrip():
    docs = [Document(content="a", id="1", score=0.8), Document(content="b")]
    encoded = encode_documents(docs)
    decoded = json.loads(encoded)
    assert len(decoded) == 2
    assert decoded[0]["document.content"] == "a"
    assert decoded[0]["document.id"] == "1"
    assert decoded[0]["document.score"] == 0.8
    assert decoded[1]["document.content"] == "b"
    assert "document.id" not in decoded[1]


# ---------------------------------------------------------------------------
# @retrieval_span decorator
# ---------------------------------------------------------------------------


def test_retrieval_span_decorator_sets_kind(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return [Document(content="result")]

    fetch(query="test query")
    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 1
    assert spans[0].attributes["openinference.span.kind"] == "RETRIEVER"


def test_retrieval_span_decorator_encodes_documents(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return [Document(content="abc", id="pmid123", score=0.9)]

    fetch(query="q")
    spans = memory_exporter.get_finished_spans()
    docs = json.loads(spans[0].attributes["retrieval.documents"])
    assert docs[0]["document.content"] == "abc"
    assert docs[0]["document.id"] == "pmid123"
    assert docs[0]["document.score"] == 0.9


def test_retrieval_span_decorator_extracts_query_kwarg(memory_exporter):
    @retrieval_span(query_arg="query")
    def fetch(query: str) -> list[Document]:
        return []

    fetch(query="what is diabetes?")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["input.value"] == "what is diabetes?"


def test_retrieval_span_decorator_extracts_query_positional(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return []

    fetch("positional query")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["input.value"] == "positional query"


def test_retrieval_span_decorator_handles_empty_list(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return []

    fetch(query="q")
    spans = memory_exporter.get_finished_spans()
    assert "retrieval.documents" not in spans[0].attributes


def test_retrieval_span_custom_name(memory_exporter):
    @retrieval_span(name="pubmed_retrieval")
    def fetch(query: str) -> list[Document]:
        return []

    fetch(query="q")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].name == "pubmed_retrieval"


def test_retrieval_span_is_child_of_coalex_context(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return [Document(content="x")]

    with coalex.coalex_context(agent_id="test-agent", request_id="req-1"):
        fetch(query="q")

    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 2
    parent = next(s for s in spans if s.name == "coalex.invocation")
    child = next(s for s in spans if s.name == "fetch")
    assert child.parent is not None
    assert child.parent.span_id == parent.context.span_id


# ---------------------------------------------------------------------------
# output.value on RETRIEVER spans
# ---------------------------------------------------------------------------


def test_retrieval_span_decorator_sets_output_value(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return [Document(content="result")]

    fetch(query="test")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["output.value"] == "Retrieved 1 document(s)"


def test_retrieval_span_decorator_output_value_with_ids(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return [Document(content="a", id="pmid1"), Document(content="b", id="pmid2")]

    fetch(query="test")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["output.value"] == "Retrieved 2 document(s): [pmid1, pmid2]"


def test_retrieval_span_decorator_no_output_value_on_empty(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return []

    fetch(query="test")
    spans = memory_exporter.get_finished_spans()
    assert "output.value" not in spans[0].attributes


async def test_retrieval_span_async_sets_output_value(memory_exporter):
    @retrieval_span(name="async_retrieval")
    async def fetch(query: str) -> list[Document]:
        return [Document(content="async", id="a1")]

    await fetch(query="test")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["output.value"] == "Retrieved 1 document(s): [a1]"


# ---------------------------------------------------------------------------
# RetrievalSpan context manager
# ---------------------------------------------------------------------------


def test_retrieval_span_context_manager_set_documents(memory_exporter):
    with RetrievalSpan("my_retrieval") as r:
        r.set_documents([Document(content="doc1", id="id1")])

    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["openinference.span.kind"] == "RETRIEVER"
    docs = json.loads(spans[0].attributes["retrieval.documents"])
    assert docs[0]["document.id"] == "id1"


def test_retrieval_span_context_manager_sets_output_value(memory_exporter):
    with RetrievalSpan("my_retrieval") as r:
        r.set_documents([Document(content="doc1", id="id1"), Document(content="doc2")])

    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["output.value"] == "Retrieved 2 document(s): [id1]"


def test_retrieval_span_context_manager_set_query(memory_exporter):
    with RetrievalSpan("my_retrieval") as r:
        r.set_query("my search query")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["input.value"] == "my search query"


# ---------------------------------------------------------------------------
# Async decorator
# ---------------------------------------------------------------------------


async def test_retrieval_span_async_encodes_documents(memory_exporter):
    @retrieval_span(name="async_retrieval", query_arg="query")
    async def fetch(query: str) -> list[Document]:
        return [Document(content="async result", id="a1", score=0.88)]

    await fetch(query="async query")
    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 1
    assert spans[0].attributes["openinference.span.kind"] == "RETRIEVER"
    assert spans[0].attributes["input.value"] == "async query"
    docs = json.loads(spans[0].attributes["retrieval.documents"])
    assert docs[0]["document.id"] == "a1"
    assert docs[0]["document.score"] == 0.88


async def test_retrieval_span_async_handles_empty_list(memory_exporter):
    @retrieval_span()
    async def fetch(query: str) -> list[Document]:
        return []

    await fetch(query="q")
    spans = memory_exporter.get_finished_spans()
    assert "retrieval.documents" not in spans[0].attributes


# ---------------------------------------------------------------------------
# Status OK / ERROR tests
# ---------------------------------------------------------------------------


def test_retrieval_span_decorator_status_ok_on_success(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        return [Document(content="result")]

    fetch(query="q")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.OK


def test_retrieval_span_decorator_status_error_on_exception(memory_exporter):
    @retrieval_span()
    def fetch(query: str) -> list[Document]:
        raise RuntimeError("fetch failed")

    with pytest.raises(RuntimeError, match="fetch failed"):
        fetch(query="q")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.ERROR
    assert "fetch failed" in spans[0].status.description
    exception_events = [e for e in spans[0].events if e.name == "exception"]
    assert len(exception_events) == 1


async def test_retrieval_span_async_status_ok(memory_exporter):
    @retrieval_span()
    async def fetch(query: str) -> list[Document]:
        return [Document(content="result")]

    await fetch(query="q")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.OK


async def test_retrieval_span_async_status_error(memory_exporter):
    @retrieval_span()
    async def fetch(query: str) -> list[Document]:
        raise RuntimeError("async fetch failed")

    with pytest.raises(RuntimeError, match="async fetch failed"):
        await fetch(query="q")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.ERROR
    exception_events = [e for e in spans[0].events if e.name == "exception"]
    assert len(exception_events) == 1


def test_retrieval_span_context_manager_status_ok(memory_exporter):
    with RetrievalSpan("my_retrieval") as r:
        r.set_query("q")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.OK


def test_retrieval_span_context_manager_status_error(memory_exporter):
    with pytest.raises(ValueError, match="ctx error"), RetrievalSpan("my_retrieval") as r:
        r.set_query("q")
        raise ValueError("ctx error")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.ERROR
    exception_events = [e for e in spans[0].events if e.name == "exception"]
    assert len(exception_events) == 1
