from typing import Any

import httpx

from config import RERANK_LIMIT, RERANKER_MODEL, RERANKER_URL, get_upstream_kwargs, logger
from retrieval import extract_page_content


async def get_rerank_scores(
    client: httpx.AsyncClient,
    query: str,
    targets: list[str],
) -> list[float]:
    if not targets:
        return []

    response = await client.post(
        str(RERANKER_URL),
        **get_upstream_kwargs(),
        json={
            "model": RERANKER_MODEL,
            "encoding_format": "float",
            "text_1": query,
            "text_2": targets,
        },
    )
    response.raise_for_status()

    data = response.json().get("data") or []
    return [float(sample["score"]) for sample in data]


async def rerank_points(
    client: httpx.AsyncClient,
    query: str,
    points: list[Any],
) -> tuple[list[Any], list[Any]]:
    """Return (reranked_head, retrieval_tail) so we don't lose candidates."""
    if not points:
        return [], []

    rerank_candidates = points[:RERANK_LIMIT]
    tail = points[RERANK_LIMIT:]

    targets = [extract_page_content(p) for p in rerank_candidates]
    try:
        scores = await get_rerank_scores(client, query, targets)
    except Exception as exc:
        logger.warning("Rerank failed, using retrieval order: %s", exc)
        return rerank_candidates, tail

    if len(scores) != len(rerank_candidates):
        logger.warning("Rerank score count mismatch, using retrieval order")
        return rerank_candidates, tail

    paired = sorted(zip(scores, rerank_candidates), key=lambda x: x[0], reverse=True)
    reranked = [p for _, p in paired]
    return reranked, tail