from __future__ import annotations

import math

from .domain import (
    CandidateScorer,
    CandidateScoringError,
    RankedCandidate,
    RerankConfig,
    RerankRequest,
    RerankResult,
)


class Reranker:
    def __init__(self, scorer: CandidateScorer, config: RerankConfig | None = None) -> None:
        self._scorer = scorer
        self._config = config or RerankConfig()

    def rerank(self, request: RerankRequest) -> RerankResult:
        if len(request.candidates) < 2:
            return self._preserve(request, reason="insufficient_candidates")
        if len(request.candidates) > self._config.max_candidates:
            return self._preserve(request, reason="unsafe_input")
        if any(
            len(candidate.surface) > self._config.max_candidate_chars
            for candidate in request.candidates
        ):
            return self._preserve(request, reason="unsafe_input")
        surfaces = [candidate.surface for candidate in request.candidates]
        if (
            any(not surface for surface in surfaces)
            or len(set(surfaces)) != len(surfaces)
            or any(not math.isfinite(candidate.prior_score) for candidate in request.candidates)
        ):
            return self._preserve(request, reason="unsafe_input")
        if (
            len(request.left_context) + len(request.right_context)
            > self._config.max_context_words
        ):
            return self._preserve(request, reason="unsafe_input")
        if len(request.reading) > self._config.max_reading_chars:
            return self._preserve(request, reason="unsafe_input")
        if (
            sum(len(word) for word in request.left_context + request.right_context)
            > self._config.max_context_chars
        ):
            return self._preserve(request, reason="unsafe_input")
        try:
            model_scores = tuple(
                float(score) for score in self._scorer.score_candidates(request)
            )
        except CandidateScoringError as error:
            return self._preserve(request, reason=error.code)
        except Exception:
            return self._preserve(request, reason="scorer_error")
        if len(model_scores) != len(request.candidates) or not all(
            math.isfinite(score) for score in model_scores
        ):
            return self._preserve(request, reason="invalid_scores")
        combined_scores = tuple(
            model_score + self._config.prior_weight * candidate.prior_score
            for candidate, model_score in zip(request.candidates, model_scores, strict=True)
        )
        model_order = sorted(
            range(len(request.candidates)),
            key=lambda index: (-combined_scores[index], index),
        )
        margin = combined_scores[model_order[0]] - combined_scores[model_order[1]]
        changed = model_order[0] != 0 and margin >= self._config.min_margin
        order = model_order if changed else list(range(len(request.candidates)))
        if changed:
            reason = "accepted"
        elif model_order[0] == 0:
            reason = "baseline_best"
        else:
            reason = "low_margin"
        ranked = tuple(
            RankedCandidate(
                surface=request.candidates[index].surface,
                original_rank=index,
                prior_score=request.candidates[index].prior_score,
                model_score=model_scores[index],
                combined_score=combined_scores[index],
            )
            for index in order
        )
        return RerankResult(
            reading=request.reading,
            ranked=ranked,
            changed=changed,
            reason=reason,
            margin=margin,
        )

    @staticmethod
    def _preserve(request: RerankRequest, *, reason: str) -> RerankResult:
        return RerankResult(
            reading=request.reading,
            ranked=tuple(
                RankedCandidate(
                    surface=candidate.surface,
                    original_rank=index,
                    prior_score=candidate.prior_score,
                    model_score=None,
                    combined_score=None,
                )
                for index, candidate in enumerate(request.candidates)
            ),
            changed=False,
            reason=reason,
            margin=None,
        )