"""Alibaba Cloud Model Studio compatible rerank adapter."""

from __future__ import annotations

import math
from collections.abc import Mapping, Sequence
from typing import Any

import httpx

from app.adapters.bailian._base import (
    BailianHttpAdapter,
    conservative_utf8_token_count,
    count_tokens,
    extract_request_id,
    invalid_request,
    invalid_response,
    parse_usage,
    response_model,
)
from app.ports.model_providers import RankedItem, RerankResult, TokenCounter

RERANK_MAX_DOCUMENTS = 500
RERANK_MAX_TOKENS_PER_TEXT = 4_000
RERANK_MAX_TOKENS_PER_REQUEST = 120_000


class BailianRerankerAdapter(BailianHttpAdapter):
    """Call ``compatible-api/v1/reranks`` and map indices locally."""

    def __init__(
        self,
        *,
        api_key: str,
        base_url: str,
        model: str = "qwen3-rerank",
        token_counter: TokenCounter = conservative_utf8_token_count,
        http_client: httpx.AsyncClient | None = None,
        timeout_seconds: float = 30.0,
        max_retries: int = 0,
        retry_base_seconds: float = 0.5,
    ) -> None:
        if not model or model != model.strip():
            raise invalid_request("rerank.configuration", "invalid_model")
        if not callable(token_counter):
            raise invalid_request("rerank.configuration", "invalid_token_counter")

        super().__init__(
            api_key=api_key,
            base_url=base_url,
            expected_path="/compatible-api/v1",
            http_client=http_client,
            timeout_seconds=timeout_seconds,
            max_retries=max_retries,
            retry_base_seconds=retry_base_seconds,
        )
        self._model = model
        self._token_counter = token_counter

    async def rerank(
        self,
        query: str,
        documents: Sequence[str],
        *,
        top_n: int,
        instruct: str | None = None,
    ) -> RerankResult:
        operation = "rerank.create"
        validated_documents = self._validate_request(
            query=query,
            documents=documents,
            top_n=top_n,
            instruct=instruct,
            operation=operation,
        )
        payload: dict[str, Any] = {
            "model": self._model,
            "query": query,
            "documents": list(validated_documents),
            "top_n": top_n,
        }
        if instruct is not None:
            payload["instruct"] = instruct

        sensitive_values = (
            query,
            *validated_documents,
            *((instruct,) if instruct is not None else ()),
        )
        response = await self._post_json(
            operation=operation,
            path="reranks",
            payload=payload,
            sensitive_values=sensitive_values,
        )
        items = self._parse_results(
            response.body,
            documents=validated_documents,
            top_n=top_n,
            operation=operation,
        )
        return RerankResult(
            items=items,
            model=response_model(
                response.body,
                self._model,
                sensitive_values=(*sensitive_values, self._api_key),
            ),
            request_id=extract_request_id(
                response.body,
                sensitive_values=(*sensitive_values, self._api_key),
            ),
            usage=parse_usage(response.body.get("usage")),
            elapsed_ms=response.elapsed_ms,
        )

    def _validate_request(
        self,
        *,
        query: str,
        documents: Sequence[str],
        top_n: int,
        instruct: str | None,
        operation: str,
    ) -> tuple[str, ...]:
        if not isinstance(query, str) or not query:
            raise invalid_request(operation, "invalid_query")
        if isinstance(documents, (str, bytes)) or not isinstance(documents, Sequence):
            raise invalid_request(operation, "invalid_document_collection")
        validated_documents = tuple(documents)
        if not validated_documents:
            raise invalid_request(operation, "empty_documents")
        if len(validated_documents) > RERANK_MAX_DOCUMENTS:
            raise invalid_request(operation, "document_count_exceeded")
        if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n <= 0:
            raise invalid_request(operation, "invalid_top_n")
        if instruct is not None and (not isinstance(instruct, str) or not instruct):
            raise invalid_request(operation, "invalid_instruct")

        query_tokens = count_tokens(
            self._token_counter,
            query,
            operation=operation,
        )
        if query_tokens > RERANK_MAX_TOKENS_PER_TEXT:
            raise invalid_request(operation, "query_token_limit_exceeded")

        document_tokens_total = 0
        for document in validated_documents:
            if not isinstance(document, str) or not document:
                raise invalid_request(operation, "invalid_document")
            document_tokens = count_tokens(
                self._token_counter,
                document,
                operation=operation,
            )
            if document_tokens > RERANK_MAX_TOKENS_PER_TEXT:
                raise invalid_request(operation, "document_token_limit_exceeded")
            document_tokens_total += document_tokens

        request_tokens = query_tokens * len(validated_documents) + document_tokens_total
        if request_tokens > RERANK_MAX_TOKENS_PER_REQUEST:
            raise invalid_request(operation, "request_token_limit_exceeded")
        return validated_documents

    def _parse_results(
        self,
        body: Mapping[str, Any],
        *,
        documents: tuple[str, ...],
        top_n: int,
        operation: str,
    ) -> tuple[RankedItem, ...]:
        results = body.get("results")
        if not isinstance(results, list) or len(results) > min(top_n, len(documents)):
            raise invalid_response(operation, "invalid_rerank_count")

        seen_indices: set[int] = set()
        parsed: list[RankedItem] = []
        previous_score = math.inf
        for item in results:
            if not isinstance(item, Mapping):
                raise invalid_response(operation, "invalid_rerank_item")
            index = item.get("index")
            if (
                isinstance(index, bool)
                or not isinstance(index, int)
                or not 0 <= index < len(documents)
                or index in seen_indices
            ):
                raise invalid_response(operation, "invalid_rerank_index")
            raw_score = item.get("relevance_score")
            if isinstance(raw_score, bool) or not isinstance(raw_score, (int, float)):
                raise invalid_response(operation, "invalid_rerank_score")
            score = float(raw_score)
            if not math.isfinite(score) or not 0 <= score <= 1 or score > previous_score:
                raise invalid_response(operation, "invalid_rerank_score")

            seen_indices.add(index)
            previous_score = score
            parsed.append(
                RankedItem(
                    index=index,
                    relevance_score=score,
                    document=documents[index],
                )
            )
        return tuple(parsed)