from __future__ import annotations

import math
from typing import Any

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

EMBED_MODULE = "cerebellum"


def _cosine(a: list[float], b: list[float]) -> float:
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0.0 or nb == 0.0:
        return 0.0
    return dot / (na * nb)


class CerebellumEmbedRecall:
    """小织 · embedding 粗召回：query 与 atom snippet cosine 打分。"""

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

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

    def blend_scored_atoms(
        self,
        query: str,
        scored: list[tuple[Atom, float]],
        *,
        max_coarse: int,
        w_embed: float,
        w_rule: float,
        trace: list[str] | None = None,
    ) -> list[tuple[Atom, float]]:
        if not scored:
            return scored
        if not self.enabled():
            if trace is not None:
                trace.append(f"[{CEREBELLUM}] embed skipped: slot 未配置")
            return scored

        pool = scored[:max_coarse]
        texts = [query] + [(a.raw_text or "")[:512] for a, _ in pool]
        router_trace = RouterTrace()
        try:
            result = self._router.embed(EMBED_MODULE, texts, trace=router_trace)
        except LLMUnavailable as exc:
            if trace is not None:
                trace.append(f"[{CEREBELLUM}] embed skipped: {exc}")
            return scored

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

        if len(result.embeddings) < len(texts):
            if trace is not None:
                trace.append(f"[{CEREBELLUM}] embed skipped: 向量数不足")
            return scored

        query_vec = result.embeddings[0]
        embed_by_id: dict[str, float] = {}
        for (atom, _), vec in zip(pool, result.embeddings[1:]):
            embed_by_id[atom.id] = _cosine(query_vec, vec)

        max_embed = max(embed_by_id.values()) if embed_by_id else 1.0
        max_rule = max(rule for _, rule in pool) or 1.0

        def _blended(atom: Atom, rule_score: float) -> float:
            embed = embed_by_id.get(atom.id, 0.0)
            embed_norm = embed / max_embed if max_embed else 0.0
            rule_norm = rule_score / max_rule if max_rule else 0.0
            return w_embed * embed_norm + w_rule * rule_norm

        blended_pool = [(a, _blended(a, rule)) for a, rule in pool]
        blended_pool.sort(key=lambda x: x[1], reverse=True)
        pool_ids = {a.id for a, _ in pool}
        tail = [(a, s) for a, s in scored if a.id not in pool_ids]
        if trace is not None:
            trace.append(f"[{CEREBELLUM}] embed coarse {len(pool)}")
        return blended_pool + tail

    def rank_vault_chunks(
        self,
        query: str,
        chunks: list[dict[str, Any]],
        *,
        vault_id: str,
        index: Any,
        trace: list[str] | None = None,
    ) -> list[tuple[dict[str, Any], float]]:
        """vault chunk cosine 排序；缺失向量时 lazy embed 写入 sidecar。"""
        from niannian.vault.embed_index import chunk_embed_key

        if not chunks:
            return []
        if not self.enabled():
            if trace is not None:
                trace.append(f"[{CEREBELLUM}] vault embed skipped: slot 未配置")
            return [(c, 0.0) for c in chunks]

        pending: list[tuple[dict[str, Any], str, str]] = []
        ready: list[tuple[dict[str, Any], str, list[float]]] = []
        for chunk in chunks:
            idx = int(chunk.get("index", 0))
            key = chunk_embed_key(vault_id, idx)
            vec = index.get_vector(key)
            if vec:
                ready.append((chunk, key, vec))
                continue
            heading = str(chunk.get("heading") or "")
            body = str(chunk.get("text") or "")
            pending.append((chunk, key, f"{heading} {body}".strip()[:512]))

        query_vec: list[float] = []
        if pending:
            router_trace = RouterTrace()
            try:
                result = self._router.embed(
                    EMBED_MODULE,
                    [query] + [text for _, _, text in pending],
                    trace=router_trace,
                )
            except LLMUnavailable as exc:
                if trace is not None:
                    trace.append(f"[{CEREBELLUM}] vault embed skipped: {exc}")
                return [(c, 0.0) for c in chunks]
            if trace is not None:
                trace.extend(router_trace.lines)
            if not result.embeddings:
                return [(c, 0.0) for c in chunks]
            query_vec = result.embeddings[0]
            for (chunk, key, preview), vec in zip(pending, result.embeddings[1:]):
                if vec:
                    index.upsert_lazy(
                        key=key,
                        vault_id=vault_id,
                        chunk_index=int(chunk.get("index", 0)),
                        vector=vec,
                        text_preview=preview,
                    )
                ready.append((chunk, key, vec))
        elif ready:
            router_trace = RouterTrace()
            try:
                result = self._router.embed(EMBED_MODULE, [query], trace=router_trace)
                query_vec = result.embeddings[0] if result.embeddings else []
            except LLMUnavailable:
                query_vec = []

        scored = [(chunk, _cosine(query_vec, vec)) for chunk, _key, vec in ready if vec]
        scored.sort(key=lambda x: x[1], reverse=True)
        if trace is not None:
            trace.append(f"[{CEREBELLUM}] vault embed rank {len(scored)}")
        return scored or [(c, 0.0) for c in chunks]
