from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from niannian.llm.router import LLMRouter
from niannian.llm.types import LLMUnavailable, RouterTrace
from niannian.personas import CEREBELLUM

RERANK_MODULE = "retrieval_rerank"


@dataclass
class RerankCandidate:
    id: str
    channel: str  # atom | vault | wiki
    text: str
    meta: dict[str, Any]

    @property
    def source(self) -> str:
        """召回通道标签：vault | wiki | atom（等同 channel）。"""
        return self.channel if self.channel in ("vault", "wiki", "atom") else self.channel


class CerebellumRerank:
    """小织 · 向量重排：query + 候选 snippet → relevance 分数。"""

    def __init__(self, router: LLMRouter) -> None:
        self._router = router

    def enabled(self) -> bool:
        return self._router.is_module_configured(RERANK_MODULE)

    def rerank_candidates(
        self,
        query: str,
        candidates: list[RerankCandidate],
        *,
        top_n: int | None = None,
        trace: list[str] | None = None,
    ) -> list[tuple[RerankCandidate, float]]:
        if not candidates:
            return []
        if not self.enabled():
            if trace is not None:
                trace.append(f"[{CEREBELLUM}] rerank skipped: slot 未配置")
            return [(c, 0.0) for c in candidates]

        router_trace = RouterTrace()
        documents = [c.text for c in candidates]
        limit = top_n if top_n is not None else len(candidates)
        try:
            result = self._router.rerank(
                RERANK_MODULE,
                query,
                documents,
                top_n=limit,
                trace=router_trace,
            )
        except LLMUnavailable as exc:
            if trace is not None:
                trace.append(f"[{CEREBELLUM}] rerank skipped: {exc}")
            return [(c, 0.0) for c in candidates]

        if trace is not None:
            trace.extend(router_trace.lines)

        scored: list[tuple[RerankCandidate, float]] = []
        for idx, score in zip(result.indices, result.scores):
            if 0 <= idx < len(candidates):
                scored.append((candidates[idx], score))
        if not scored:
            return [(c, 0.0) for c in candidates]
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:limit]
