from __future__ import annotations

import json

import httpx
import pytest

from app.adapters.bailian.rerank import BailianRerankerAdapter
from app.ports.model_providers import ModelProviderError, ProviderErrorKind

BASE_URL = "https://workspace.cn-beijing.maas.aliyuncs.com/compatible-api/v1"
API_KEY = "sk-test-rerank-secret"


@pytest.mark.asyncio
async def test_rerank_uses_separate_endpoint_and_maps_result_indices() -> None:
    documents = ["copper evidence", "gold evidence", "irrelevant"]

    def handler(request: httpx.Request) -> httpx.Response:
        assert request.url.path == "/compatible-api/v1/reranks"
        assert request.headers["Authorization"] == f"Bearer {API_KEY}"
        payload = json.loads(request.content)
        assert payload == {
            "model": "qwen3-rerank",
            "query": "where is gold",
            "documents": documents,
            "top_n": 2,
            "instruct": "rank geological evidence",
        }
        return httpx.Response(
            200,
            json={
                "id": "req_rerank_1",
                "model": "qwen3-rerank",
                "results": [
                    {"index": 1, "relevance_score": 0.91},
                    {"index": 0, "relevance_score": 0.72},
                ],
                "usage": {"input_tokens": 42, "total_tokens": 42},
            },
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianRerankerAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
        )
        result = await adapter.rerank(
            "where is gold",
            documents,
            top_n=2,
            instruct="rank geological evidence",
        )

    assert [item.index for item in result.items] == [1, 0]
    assert [item.document for item in result.items] == [documents[1], documents[0]]
    assert [item.relevance_score for item in result.items] == [0.91, 0.72]
    assert result.request_id == "req_rerank_1"


@pytest.mark.asyncio
async def test_rerank_enforces_500_4000_and_repeated_query_120000_formula() -> None:
    calls = 0

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        return httpx.Response(
            200,
            json={"results": [{"index": 0, "relevance_score": 1.0}]},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianRerankerAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
        )

        await adapter.rerank("q" * 4_000, ["d" * 4_000], top_n=1)
        exactly_120000 = ["d" * 40] * 500
        await adapter.rerank("q" * 200, exactly_120000, top_n=1)
        assert calls == 2

        invalid_requests = (
            ("q" * 4_001, ["d"]),
            ("q", ["d" * 4_001]),
            ("q", ["d"] * 501),
            ("q" * 200, [*(["d" * 40] * 499), "d" * 41]),
        )
        for query, documents in invalid_requests:
            with pytest.raises(ModelProviderError) as exc_info:
                await adapter.rerank(query, documents, top_n=1)
            assert exc_info.value.kind is ProviderErrorKind.INVALID_REQUEST
        assert calls == 2


@pytest.mark.asyncio
async def test_rerank_rejects_invalid_provider_index() -> None:
    def handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={"results": [{"index": 2, "relevance_score": 0.9}]},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianRerankerAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
        )
        with pytest.raises(ModelProviderError) as exc_info:
            await adapter.rerank("query", ["only document"], top_n=1)

    assert exc_info.value.kind is ProviderErrorKind.INVALID_RESPONSE
    assert exc_info.value.provider_code == "invalid_rerank_index"


@pytest.mark.asyncio
async def test_rerank_retries_5xx_with_a_fixed_bound() -> None:
    calls = 0

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        if calls == 1:
            return httpx.Response(503, json={"error": {"code": "temporarily_unavailable"}})
        return httpx.Response(
            200,
            json={"results": [{"index": 0, "relevance_score": 0.9}]},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianRerankerAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
            max_retries=1,
            retry_base_seconds=0,
        )
        result = await adapter.rerank("query", ["document"], top_n=1)

    assert result.items[0].index == 0
    assert calls == 2