import asyncio
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 []

    for attempt in range(5):
        try:
            response = await client.post(
                str(RERANKER_URL),
                **get_upstream_kwargs(),
                json={
                    "model": RERANKER_MODEL,
                    "encoding_format": "float",
                    "text_1": query,
                    "text_2": targets,
                },
            )
        except Exception as exc:
            if attempt < 4:
                await asyncio.sleep(2 ** attempt)
                continue
            raise exc

        if response.status_code == 429:
            wait = 2 ** attempt
            logger.warning("Rerank 429, retry %d/5 in %ds", attempt + 1, wait)
            await asyncio.sleep(wait)
            continue

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

    logger.error("Rerank 429 after all retries, falling back")
    return []


async def rerank_points(
    client: httpx.AsyncClient,
    query: str,
    points: list[Any],
) -> tuple[list[Any], list[Any]]:
    if not points:
        return [], []

    head = points[:RERANK_LIMIT]
    tail = points[RERANK_LIMIT:]
    targets = [extract_page_content(p) for p in head]

    try:
        scores = await _get_rerank_scores(client, query, targets)
    except Exception as exc:
        logger.warning("Rerank failed, using retrieval order: %s", exc)
        return head, tail

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

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