"""
Re-ranking client for Synesis.

Sends query-document pairs to Synesis's /v1/rerank endpoint for
cross-attention relevance scoring. Supports content-type-aware
instruction injection and multimodal queries.
"""

import base64
import logging
import time
from typing import Optional

import requests

from library.metrics import (
    RERANK_CANDIDATES,
    RERANK_DURATION,
    RERANK_REQUESTS_TOTAL,
)

logger = logging.getLogger(__name__)


class RerankerClient:
    """
    Client for re-ranking search candidates via Synesis.

    Uses the system reranker model's API configuration to send
    query-document pairs to Synesis's ``POST /v1/rerank`` endpoint.
    """

    def __init__(self, reranker_model, user=None):
        """
        :param reranker_model: ``LLMModel`` instance for re-ranking.
        :param user: Optional Django user for usage tracking.
        """
        self.model = reranker_model
        self.api = reranker_model.api
        self.user = user
        self.base_url = self.api.base_url.rstrip("/")
        self.model_name = self.model.name
        self.timeout = self.api.timeout_seconds or 30

        logger.info(
            "RerankerClient initialized model=%s base_url=%s",
            self.model_name,
            self.base_url,
        )

    def rerank(
        self,
        query: str,
        candidates: list,
        instruction: str = "",
        top_n: Optional[int] = None,
        query_image: Optional[bytes] = None,
    ) -> list:
        """
        Re-rank search candidates via Synesis /v1/rerank.

        :param query: Query text.
        :param candidates: List of SearchCandidate instances.
        :param instruction: Content-type-aware re-ranking instruction.
        :param top_n: Return only top N results (server-side truncation).
        :param query_image: Optional image bytes for multimodal query.
        :returns: Re-ranked list of SearchCandidate instances with updated scores.
        """
        if not candidates:
            return []

        RERANK_CANDIDATES.observe(len(candidates))

        start = time.time()

        try:
            # Build Synesis rerank request
            payload = self._build_payload(
                query, candidates, instruction, top_n, query_image
            )

            # Convention shared with every other service client in Mnemosyne
            # (embedding_client, vision, concepts): ``base_url`` is the
            # OpenAI-style ``/v1`` root (e.g. ``http://pan.helu.ca:8400/v1``),
            # and each client appends a path-only segment. Prepending
            # ``/v1`` here would build ``…/v1/v1/rerank`` and 404 against
            # Synesis.
            url = f"{self.base_url}/rerank"
            headers = {"Content-Type": "application/json"}
            if self.api.api_key:
                headers["Authorization"] = f"Bearer {self.api.api_key}"

            logger.debug(
                "Rerank request candidates=%d instruction_len=%d top_n=%s",
                len(candidates),
                len(instruction),
                top_n,
            )

            resp = requests.post(
                url, json=payload, headers=headers, timeout=self.timeout
            )
            if resp.status_code != 200:
                logger.error(
                    "Rerank failed status=%d body=%s",
                    resp.status_code,
                    resp.text[:500],
                )
            resp.raise_for_status()
            data = resp.json()

            # Parse results and update candidate scores
            reranked = self._apply_scores(candidates, data)

            elapsed = time.time() - start
            RERANK_DURATION.labels(model_name=self.model_name).observe(elapsed)
            RERANK_REQUESTS_TOTAL.labels(
                model_name=self.model_name, status="success"
            ).inc()

            logger.info(
                "Rerank completed candidates=%d returned=%d elapsed=%.3fs",
                len(candidates),
                len(reranked),
                elapsed,
            )

            self._log_usage(len(candidates), elapsed)

            return reranked

        except Exception as exc:
            elapsed = time.time() - start
            RERANK_REQUESTS_TOTAL.labels(
                model_name=self.model_name, status="error"
            ).inc()
            logger.error("Rerank failed: %s", exc, exc_info=True)
            raise

    def _build_payload(
        self,
        query: str,
        candidates: list,
        instruction: str,
        top_n: Optional[int],
        query_image: Optional[bytes],
    ) -> dict:
        """
        Build the Synesis /v1/rerank request payload.

        :param query: Query text.
        :param candidates: List of SearchCandidate instances.
        :param instruction: Re-ranking instruction.
        :param top_n: Optional top-N limit.
        :param query_image: Optional query image bytes.
        :returns: Request payload dict.
        """
        # Build query object
        query_obj: dict = {"text": query}
        if query_image:
            b64 = base64.b64encode(query_image).decode("utf-8")
            query_obj["image"] = f"data:image/png;base64,{b64}"

        # Build document objects from candidate text previews
        documents = []
        for candidate in candidates:
            doc: dict = {"text": candidate.text_preview}
            documents.append(doc)

        payload: dict = {
            "query": query_obj,
            "documents": documents,
        }

        if instruction:
            payload["instruction"] = instruction

        if top_n is not None and top_n > 0:
            payload["top_n"] = top_n

        return payload

    def _apply_scores(self, candidates: list, data: dict) -> list:
        """
        Apply Synesis rerank scores back to SearchCandidate instances.

        Synesis returns results sorted by score descending with original indices.

        :param candidates: Original candidate list.
        :param data: Synesis response.
        :returns: Re-ranked candidates with updated scores.
        """
        from .fusion import SearchCandidate

        results = data.get("results", [])
        if not results:
            logger.warning("Rerank response contained no results")
            return candidates

        # Map original index → score
        reranked = []
        for result in results:
            orig_idx = result.get("index", 0)
            score = result.get("score", 0.0)

            if orig_idx < len(candidates):
                original = candidates[orig_idx]
                reranked.append(
                    SearchCandidate(
                        chunk_uid=original.chunk_uid,
                        item_uid=original.item_uid,
                        item_title=original.item_title,
                        library_type=original.library_type,
                        text_preview=original.text_preview,
                        chunk_s3_key=original.chunk_s3_key,
                        chunk_index=original.chunk_index,
                        score=score,
                        source=original.source,
                        metadata={**original.metadata, "reranked": True},
                    )
                )

        # Already sorted by Synesis, but ensure descending order
        reranked.sort(key=lambda c: c.score, reverse=True)

        return reranked

    def _log_usage(self, document_count: int, elapsed: float):
        """
        Log re-ranking usage to LLMUsage model.

        :param document_count: Number of documents re-ranked.
        :param elapsed: Time taken in seconds.
        """
        try:
            from llm_manager.models import LLMUsage

            LLMUsage.objects.create(
                model=self.model,
                user=self.user,
                input_tokens=document_count,  # Approximate — one "token" per doc
                output_tokens=0,
                cached_tokens=0,
                total_cost=0,
                purpose="reranking",
                request_metadata={
                    "document_count": document_count,
                    "elapsed_ms": round(elapsed * 1000, 2),
                },
            )
        except Exception as exc:
            logger.warning("Failed to log rerank usage: %s", exc)
