"""Implement cross-encoder reranking via an OpenAI-compatible reranker API."""

from __future__ import annotations

import time

import requests
from loguru import logger

from app.config.settings import settings
from app.domain.retrieval import Reranker, RetrievedChunk
from app.shared.model_usage_tracker import get_model_usage_tracker


class OpenAICompatibleReranker(Reranker):
    """Call a TEI / Cohere-style reranker endpoint to re-score retrieved chunks."""

    def __init__(
        self,
        base_url: str | None = None,
        model: str | None = None,
        api_key: str | None = None,
        timeout: int = 30,
    ) -> None:
        self._base_url = (base_url or settings.reranker_base_url).rstrip("/")
        self._model = model or settings.reranker_model
        self._api_key = api_key or settings.reranker_api_key
        self._timeout = timeout

    def rerank(self, query: str, chunks: list[RetrievedChunk], top_k: int) -> list[RetrievedChunk]:
        """Return up to top_k chunks re-sorted by cross-encoder score."""
        if not chunks:
            return []

        texts = [chunk.text for chunk in chunks]
        start = time.time()
        try:
            scores = self._call_reranker(query, texts)
        except Exception as exc:
            logger.warning("Reranker call failed ({}), falling back to original order: {}", type(exc).__name__, exc)
            # Record the failure so the Status page reflects real reranker health.
            get_model_usage_tracker().record(
                provider="reranker",
                model=self._model,
                success=False,
                latency_ms=int((time.time() - start) * 1000),
                error=str(exc),
            )
            return chunks[:top_k]

        elapsed_ms = int((time.time() - start) * 1000)
        logger.debug("Reranker scored {} chunks in {}ms", len(chunks), elapsed_ms)
        # TEI/Cohere-style rerank responses carry no token usage field —
        # only call success/latency is meaningful for this role.
        get_model_usage_tracker().record(
            provider="reranker",
            model=self._model,
            success=True,
            latency_ms=elapsed_ms,
        )

        ranked = sorted(
            [(score, chunk) for score, chunk in zip(scores, chunks)],
            key=lambda x: x[0],
            reverse=True,
        )
        result = []
        for score, chunk in ranked[:top_k]:
            chunk.score = float(score)
            result.append(chunk)
        return result

    def _call_reranker(self, query: str, texts: list[str]) -> list[float]:
        """Call the reranker API and return a score per text.

        Tries TEI format first (POST /rerank with model+texts), then falls back
        to Cohere/OpenAI format (POST /v1/rerank with model+documents).
        Both formats now include the model name, which most gateways require.
        """
        headers = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"

        # TEI format: POST /rerank — include model name (required by gateway proxies)
        payload = {
            "model": self._model,
            "query": query,
            "texts": texts,
            "raw_scores": False,
            "return_text": False,
        }
        url = f"{self._base_url}/rerank"
        resp = requests.post(url, json=payload, headers=headers, timeout=self._timeout)

        if resp.status_code in (404, 400):
            # Gateway returned an error — try Cohere/OpenAI-style format as fallback.
            logger.debug(
                "TEI rerank returned {} — trying Cohere format. Body: {}",
                resp.status_code,
                resp.text[:200],
            )
            payload_v1 = {"model": self._model, "query": query, "documents": texts}
            url = f"{self._base_url}/v1/rerank"
            resp = requests.post(url, json=payload_v1, headers=headers, timeout=self._timeout)

        if not resp.ok:
            # Surface a clear error message so callers can log it meaningfully.
            try:
                err_body = resp.json()
                err_msg = err_body.get("error", {}).get("message", resp.text[:200])
            except Exception:
                err_msg = resp.text[:200]
            resp.raise_for_status()  # raises HTTPError with status code
            raise ValueError(err_msg)  # unreachable but satisfies type checker

        resp.raise_for_status()
        data = resp.json()

        # TEI response: list of {"index": N, "score": F}
        if isinstance(data, list):
            ordered = sorted(data, key=lambda x: x["index"])
            return [float(item["score"]) for item in ordered]

        # Cohere/OpenAI response: {"results": [{"index": N, "relevance_score": F}]}
        results = data.get("results", [])
        ordered = sorted(results, key=lambda x: x["index"])
        return [float(item.get("relevance_score", item.get("score", 0))) for item in ordered]
