"""ChainReranker — runs several rerankers in sequence."""
from __future__ import annotations

from typing import List

from ..core.types import Query, Retrieved


class ChainReranker:
    """Apply multiple rerankers one after another."""

    name = "chain"

    def __init__(self, stages):
        self.stages = list(stages)

    def rerank(self, query: Query, candidates: List[Retrieved], k: int) -> List[Retrieved]:
        cur = candidates
        for s in self.stages[:-1]:
            cur = s.rerank(query, cur, max(k, len(cur)))
        if self.stages:
            cur = self.stages[-1].rerank(query, cur, k)
        return cur[:k]