#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import asyncio
from dataclasses import dataclass, field
import logging
import re

import numpy as np
from sklearn.cluster import AgglomerativeClustering
from sklearn.mixture import GaussianMixture

from api.db.services.task_service import has_canceled
from common.connection_utils import timeout
from common.exceptions import TaskCanceledException
from common.token_utils import truncate
from rag.graphrag.utils import (
    chat_limiter,
    get_embed_cache,
    get_llm_cache,
    set_embed_cache,
    set_llm_cache,
)
from common.misc_utils import thread_pool_exec
from rag.utils.raptor_utils import (
    AHC_CLUSTERING_METHOD,
    GMM_CLUSTERING_METHOD,
    PSI_TREE_BUILDER,
    RAPTOR_TREE_BUILDER,
    SUPPORTED_CLUSTERING_METHODS,
    SUPPORTED_TREE_BUILDERS,
)

# Regularization added to GMM covariance diagonals; keeps components
# from collapsing on singleton/near-identical reduced points.
_GMM_REG_COVAR = 1e-4


@dataclass
class _PsiTreeNode:
    """Node used to represent the in-memory Psi merge tree."""

    index: int
    text: str = ""
    embedding: np.ndarray | None = None
    children: list["_PsiTreeNode"] = field(default_factory=list)
    parent: "_PsiTreeNode | None" = None
    # Original (leaf-level) chunk ids that contributed to this node. On
    # a leaf this is a single-element list with the leaf's own id; on an
    # internal node it's the order-preserving deduped union of its
    # children's lists. Carried up through the merge tree so each
    # produced summary knows which source chunks it covers.
    source_chunk_ids: list[str] = field(default_factory=list)


class _PsiUnionFind:
    """Build parent links for the Psi merge tree from ranked leaf pairs."""

    def __init__(self, n: int):
        """Initialize the union-find state for n leaf nodes."""
        self._rank = [0 for _ in range(n)]
        self._parent_chains = [[] for _ in range(n)]
        self._node_ids = [[i] for i in range(n)]
        self._tree = [-1 for _ in range(max(1, 2 * n - 1))]
        self._next_id = n

    @staticmethod
    def _ordered_extend(target: list[int], values: list[int]):
        """Append unseen values while preserving their original order."""
        for value in values:
            if value not in target:
                target.append(value)

    def _find(self, i: int) -> list[int]:
        """Return the parent chain for a leaf, extending it lazily."""
        chain = self._parent_chains[i]
        if not chain or (len(chain) == 1 and chain[0] == i):
            return [i]
        if chain[0] == i:
            self._ordered_extend(chain, self._find(chain[1]))
        else:
            self._ordered_extend(chain, self._find(chain[0]))
        return chain

    def _rank_bisect_right(self, chain: list[int], rank: int) -> int:
        """Return the first chain index whose rank is greater than rank."""
        idx = 0
        while idx < len(chain) and self._rank[chain[idx]] <= rank:
            idx += 1
        return idx

    def _build(self, i: int, j: int, insert_point: int | None = None):
        """Record a merge edge in the compact parent array."""
        if insert_point is not None:
            parent_ids = self._node_ids[insert_point]
            parent_rank_idx = self._rank[i] + 1
            if parent_rank_idx >= len(parent_ids):
                logging.warning(
                    "RAPTOR Psi union fallback: rank index %d is out of bounds for node %d with %d parent ids",
                    parent_rank_idx,
                    insert_point,
                    len(parent_ids),
                )
                parent_rank_idx = len(parent_ids) - 1
            self._tree[self._node_ids[i][-1]] = parent_ids[parent_rank_idx]
            return
        self._tree[self._node_ids[i][-1]] = self._next_id
        self._tree[self._node_ids[j][-1]] = self._next_id
        self._node_ids[i].append(self._next_id)
        self._next_id += 1

    def union(self, i: int, j: int) -> bool:
        """Merge two ranked leaves and return whether a new edge was added."""
        root_i = self._find(i)[-1]
        root_j = self._find(j)[-1]
        if root_i == root_j:
            return False

        if self._rank[root_i] < self._rank[root_j]:
            if not self._parent_chains[root_j]:
                self._parent_chains[root_j].append(root_j)
            chain = self._parent_chains[j]
            higher_rank_idx = self._rank_bisect_right(chain, self._rank[root_i])
            if higher_rank_idx >= len(chain):
                higher_rank_idx = len(chain) - 1
            insert_point = chain[higher_rank_idx]
            self._ordered_extend(self._parent_chains[root_i], chain[higher_rank_idx:])
            self._build(root_i, root_j, insert_point=insert_point)
        elif self._rank[root_i] > self._rank[root_j]:
            if not self._parent_chains[root_i]:
                self._parent_chains[root_i].append(root_i)
            chain = self._parent_chains[i]
            higher_rank_idx = self._rank_bisect_right(chain, self._rank[root_j])
            if higher_rank_idx >= len(chain):
                higher_rank_idx = len(chain) - 1
            insert_point = chain[higher_rank_idx]
            self._ordered_extend(self._parent_chains[root_j], chain[higher_rank_idx:])
            self._build(root_j, root_i, insert_point=insert_point)
        else:
            if not self._parent_chains[root_i]:
                self._parent_chains[root_i].append(root_i)
            self._ordered_extend(self._parent_chains[root_j], self._parent_chains[i][-1:])
            self._rank[root_i] += 1
            self._build(root_i, root_j)
        return True

    @property
    def tree(self) -> list[int]:
        """Return the compact child-to-parent array for constructed nodes."""
        return self._tree[: self._next_id]


class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
    """Build RAPTOR summary layers with the classic or Psi tree strategy."""

    def __init__(
        self,
        max_cluster,
        llm_model,
        embd_model,
        prompt,
        max_token=512,
        threshold=0.1,
        small_layer_collapse=8,
        max_errors=3,
        tree_builder=RAPTOR_TREE_BUILDER,
        clustering_method=GMM_CLUSTERING_METHOD,
        psi_exact_max_leaves=4096,
        psi_bucket_size=1024,
    ):
        """Configure RAPTOR summarization, clustering, and Psi limits."""
        self._max_cluster = max_cluster
        self._small_layer_collapse = small_layer_collapse
        self._llm_model = llm_model
        self._embd_model = embd_model
        self._threshold = threshold
        self._prompt = prompt
        self._max_token = max_token
        self._max_errors = max(1, max_errors)
        self._error_count = 0
        self._tree_builder = tree_builder or RAPTOR_TREE_BUILDER
        if self._tree_builder not in SUPPORTED_TREE_BUILDERS:
            raise ValueError(f"Unsupported RAPTOR tree builder: {self._tree_builder}")
        self._clustering_method = clustering_method or GMM_CLUSTERING_METHOD
        if self._clustering_method not in SUPPORTED_CLUSTERING_METHODS:
            raise ValueError(f"Unsupported RAPTOR clustering method: {self._clustering_method}")
        self._psi_exact_max_leaves = max(2, int(psi_exact_max_leaves or 4096))
        self._psi_bucket_size = min(max(2, int(psi_bucket_size or 1024)), self._psi_exact_max_leaves)

    def _check_task_canceled(self, task_id: str, message: str = ""):
        """Raise if the current document task was canceled."""
        if task_id and has_canceled(task_id):
            log_msg = f"Task {task_id} cancelled during RAPTOR {message}."
            logging.info(log_msg)
            raise TaskCanceledException(f"Task {task_id} was cancelled")

    @timeout(60 * 20)
    async def _chat(self, system, history, gen_conf):
        """Call the configured LLM with caching and short retries."""
        cached = await thread_pool_exec(get_llm_cache, self._llm_model.llm_name, system, history, gen_conf)
        if cached:
            return cached

        last_exc = None
        for attempt in range(3):
            try:
                response = await self._llm_model.async_chat(system, history, gen_conf)
                response = re.sub(r"^.*</think>", "", response, flags=re.DOTALL)
                if response.find("**ERROR**") >= 0:
                    raise Exception(response)
                await thread_pool_exec(set_llm_cache, self._llm_model.llm_name, system, response, history, gen_conf)
                return response
            except Exception as exc:
                last_exc = exc
                logging.warning("RAPTOR LLM call failed on attempt %d/3: %s", attempt + 1, exc)
                if attempt < 2:
                    await asyncio.sleep(1 + attempt)

        raise last_exc if last_exc else Exception("LLM chat failed without exception")

    @timeout(20)
    async def _embedding_encode(self, txt):
        """Encode text with the configured embedding model and cache result."""
        response = await thread_pool_exec(get_embed_cache, self._embd_model.llm_name, txt)
        if response is not None:
            return response
        embds, _ = await thread_pool_exec(self._embd_model.encode, [txt])
        if len(embds) < 1 or len(embds[0]) < 1:
            raise Exception("Embedding error: empty embeddings returned")
        embds = embds[0]
        await thread_pool_exec(set_embed_cache, self._embd_model.llm_name, txt, embds)
        return embds

    def _get_optimal_clusters(self, embeddings: np.ndarray, random_state: int, task_id: str = ""):
        """Choose the GMM cluster count with the lowest BIC score."""
        max_clusters = min(self._max_cluster, len(embeddings))
        if max_clusters <= 1:
            logging.info(
                "RAPTOR GMM: _get_optimal_clusters returning 1 (max_clusters=%s, embeddings=%d)",
                max_clusters,
                len(embeddings),
            )
            return 1
        n_clusters = np.arange(1, max_clusters + 1)
        bics = []
        for n in n_clusters:
            self._check_task_canceled(task_id, "get optimal clusters")

            gm = GaussianMixture(n_components=n, random_state=random_state, covariance_type="diag", reg_covar=_GMM_REG_COVAR)
            gm.fit(embeddings)
            bics.append(gm.bic(embeddings))
        optimal_clusters = n_clusters[np.argmin(bics)]
        return int(optimal_clusters)

    def _get_clusters_ahc(self, embeddings: np.ndarray, task_id: str = "") -> np.ndarray:
        """Cluster embeddings with Ward-linkage AHC and a dendrogram gap heuristic."""
        n = len(embeddings)
        if n <= 1:
            return np.zeros(n, dtype=int)
        if n == 2:
            return np.arange(n)

        self._check_task_canceled(task_id, "_get_clusters_ahc dendrogram")
        full_clust = AgglomerativeClustering(
            n_clusters=None,
            distance_threshold=0,
            compute_distances=True,
            linkage="ward",
        )
        full_clust.fit(embeddings)

        distances = full_clust.distances_
        if len(distances) > 1:
            gaps = np.diff(distances)
            max_gap_idx = int(np.argmax(gaps))
            n_clusters = max(1, min(n - max_gap_idx - 1, self._max_cluster))
        else:
            n_clusters = max(1, min(n, self._max_cluster))
        if n_clusters <= 1:
            logging.info("RAPTOR AHC: _get_clusters_ahc selected one cluster for %d embeddings", n)
            return np.zeros(n, dtype=int)

        logging.info("RAPTOR AHC: _get_clusters_ahc selected n_clusters=%d for %d embeddings", n_clusters, n)
        self._check_task_canceled(task_id, "_get_clusters_ahc fit")
        clustering = AgglomerativeClustering(n_clusters=n_clusters, linkage="ward")
        return clustering.fit_predict(embeddings)

    def _adjust_tree_nodes(self, embeddings: np.ndarray, labels: np.ndarray, max_iter: int = 5) -> np.ndarray:
        """Refine AHC assignments by reassigning nodes to nearest centroids."""
        labels = labels.copy()
        for _ in range(max_iter):
            unique_labels = np.unique(labels)
            if len(unique_labels) <= 1:
                return labels
            centroids = np.stack([embeddings[labels == lbl].mean(axis=0) for lbl in unique_labels])
            diffs = embeddings[:, np.newaxis, :] - centroids[np.newaxis, :, :]
            sq_dists = (diffs**2).sum(axis=2)
            new_label_indices = np.argmin(sq_dists, axis=1)
            new_labels = unique_labels[new_label_indices]
            if np.array_equal(new_labels, labels):
                break
            unique_new = np.unique(new_labels)
            remap = {old: new for new, old in enumerate(unique_new)}
            labels = np.array([remap[int(lbl)] for lbl in new_labels])
        return labels

    def clustering(self, embeddings, random_state: int, task_id: str = "") -> tuple[int, list[int]]:
        """Cluster one RAPTOR layer and return contiguous labels."""
        reduced_embeddings = np.asarray(embeddings, dtype=np.float64)
        if len(reduced_embeddings) == 0:
            return 0, []

        # Degrade too much ??
        n_neighbors = min(int((len(embeddings) - 1) ** 0.8), 100)
        import umap

        reduced_embeddings = umap.UMAP(
            n_neighbors=max(2, n_neighbors),
            n_components=min(12, len(embeddings) - 2),
            metric="cosine",
        ).fit_transform(embeddings)
        if self._clustering_method == AHC_CLUSTERING_METHOD:
            logging.info("RAPTOR: using clustering_method=%s before _get_clusters_ahc", self._clustering_method)
            raw_labels = self._get_clusters_ahc(reduced_embeddings, task_id=task_id)
            raw_cluster_count = np.unique(raw_labels).size
            logging.info("RAPTOR AHC: _get_clusters_ahc produced n_clusters=%d", raw_cluster_count)
            if raw_cluster_count > 1:
                labels = self._adjust_tree_nodes(reduced_embeddings, raw_labels)
                adjusted_cluster_count = np.unique(labels).size
                logging.info("RAPTOR AHC: _adjust_tree_nodes adjusted n_clusters=%d", adjusted_cluster_count)
            else:
                labels = raw_labels
                logging.warning("RAPTOR AHC: _adjust_tree_nodes skipped because _get_clusters_ahc returned one cluster")
        else:
            n_clusters = int(self._get_optimal_clusters(reduced_embeddings, random_state, task_id=task_id))
            if n_clusters <= 1:
                labels = [0 for _ in range(len(reduced_embeddings))]
            else:
                gm = GaussianMixture(n_components=n_clusters, random_state=random_state, covariance_type="diag", reg_covar=_GMM_REG_COVAR)
                gm.fit(reduced_embeddings)
                probs = gm.predict_proba(reduced_embeddings)
                labels = []
                for prob in probs:
                    candidates = np.where(prob > self._threshold)[0]
                    labels.append(int(candidates[0]) if len(candidates) else int(np.argmax(prob)))

        normalized_labels: list[int] = []
        for label in labels:
            if isinstance(label, np.ndarray):
                normalized_labels.append(int(label[0]) if len(label) else 0)
            else:
                normalized_labels.append(int(label))

        if len(normalized_labels) <= 0:
            return 0, []
        unique_labels = np.unique(normalized_labels)
        if len(unique_labels) <= 1:
            return 1, [0 for _ in normalized_labels]
        label_map = {int(old): idx for idx, old in enumerate(unique_labels)}
        return len(unique_labels), [label_map[label] for label in normalized_labels]

    @timeout(60 * 20)
    async def _summarize_texts(self, texts: list[str], callback=None, task_id: str = ""):
        """Summarize a cluster and return text plus embedding when successful."""
        self._check_task_canceled(task_id, "summarization")

        len_per_chunk = int((self._llm_model.max_length - self._max_token) / len(texts))
        cluster_content = "\n".join([truncate(t, max(1, len_per_chunk)) for t in texts])
        try:
            async with chat_limiter:
                self._check_task_canceled(task_id, "before LLM call")

                cnt = await self._chat(
                    "You're a helpful assistant.\n\nHelp me with the following task.\n\n%s" % self._prompt.format(cluster_content=cluster_content),
                    [
                        {
                            "role": "user",
                            "content": "Beside the summarization, give a title at the first line of your summarization. Must be in the same language as the paragraphs.",
                        }
                    ],
                    {"max_tokens": max(self._max_token, 512)},  # fix issue:  #10235
                )
                cnt = re.sub(
                    "(······\n由于长度的原因，回答被截断了，要继续吗？|For the content length reason, it stopped, continue?)",
                    "",
                    cnt,
                )
                logging.debug(f"SUM: {cnt}")

                self._check_task_canceled(task_id, "before embedding")

                embds = await self._embedding_encode(cnt)
                return cnt.split("\n")[0], cnt, embds
        except TaskCanceledException:
            raise
        except Exception as exc:
            self._error_count += 1
            warn_msg = f"[RAPTOR] Skip cluster ({len(texts)} chunks) due to error: {exc}"
            logging.warning(warn_msg)
            if callback:
                callback(msg=warn_msg)
            if self._error_count >= self._max_errors:
                raise RuntimeError(f"RAPTOR aborted after {self._error_count} errors. Last error: {exc}") from exc
            return None

    @staticmethod
    def _root(node: _PsiTreeNode) -> _PsiTreeNode:
        """Return the current root for a Psi tree node."""
        while node.parent is not None:
            node = node.parent
        return node

    def _rank_leaf_pairs(self, leaves: list[_PsiTreeNode]) -> np.ndarray:
        """Rank all leaf pairs by original embedding-space cosine similarity."""
        node_embeddings = np.asarray([leaf.embedding for leaf in leaves], dtype=np.float64)
        node_embeddings = self._normalize_embeddings(node_embeddings)
        similarities = node_embeddings @ node_embeddings.T
        lower = np.tril_indices(len(leaves), -1)
        ordered = np.argsort(similarities[lower], axis=0)[::-1]
        return np.stack([lower[0][ordered], lower[1][ordered]], axis=-1)

    @staticmethod
    def _normalize_embeddings(node_embeddings: np.ndarray) -> np.ndarray:
        """Normalize embeddings for cosine operations while tolerating zero vectors."""
        node_embeddings = np.asarray(node_embeddings, dtype=np.float64)
        norms = np.linalg.norm(node_embeddings, axis=1, keepdims=True)
        return node_embeddings / np.maximum(norms, 1e-12)

    def _split_psi_buckets(self, nodes: list[_PsiTreeNode]) -> list[list[_PsiTreeNode]]:
        """Split large Psi inputs so exact pair ranking is bounded per bucket."""
        if len(nodes) <= self._psi_bucket_size:
            return [nodes]

        node_embeddings = self._normalize_embeddings(np.asarray([node.embedding for node in nodes], dtype=np.float64))
        groups = [np.arange(len(nodes), dtype=int)]
        buckets = []

        while groups:
            group = np.asarray(groups.pop(), dtype=int)
            if len(group) <= self._psi_bucket_size:
                buckets.append(group.tolist())
                continue

            fanout = min(max(2, int(np.ceil(len(group) / self._psi_bucket_size))), len(group), 32)
            group_embeddings = node_embeddings[group]
            center_idx = np.linspace(0, len(group_embeddings) - 1, num=fanout, dtype=int)
            centers = group_embeddings[center_idx].copy()

            for _ in range(5):
                labels = np.argmax(group_embeddings @ centers.T, axis=1)
                for center_id in range(fanout):
                    mask = labels == center_id
                    if not np.any(mask):
                        continue
                    center = group_embeddings[mask].mean(axis=0)
                    norm = np.linalg.norm(center)
                    centers[center_id] = center / norm if norm > 0 else center

            labels = np.argmax(group_embeddings @ centers.T, axis=1)
            split_groups = [group[labels == center_id].tolist() for center_id in range(fanout)]
            split_groups = [bucket for bucket in split_groups if bucket]
            if len(split_groups) <= 1:
                split_groups = [group[start : start + self._psi_bucket_size].tolist() for start in range(0, len(group), self._psi_bucket_size)]
            groups.extend(split_groups)

        buckets = [bucket for bucket in buckets if bucket]
        buckets.sort(key=lambda bucket: (len(bucket), bucket[0]))
        return [[nodes[idx] for idx in bucket] for bucket in buckets]

    def _assign_prototype_embeddings(self, node: _PsiTreeNode) -> np.ndarray:
        """Assign mean child embeddings to internal Psi nodes for bucket-level ranking."""
        if not node.children:
            return np.asarray(node.embedding, dtype=np.float64)
        embeddings = np.asarray([self._assign_prototype_embeddings(child) for child in node.children], dtype=np.float64)
        node.embedding = embeddings.mean(axis=0)
        return node.embedding

    @staticmethod
    def _iter_nodes(root: _PsiTreeNode):
        """Yield nodes in a Psi tree using a stack traversal."""
        stack = [root]
        while stack:
            node = stack.pop()
            yield node
            stack.extend(node.children)

    def _create_psi_parent(self, index: int, children: list[_PsiTreeNode]) -> _PsiTreeNode:
        """Create a parent node and attach the provided children to it."""
        parent = _PsiTreeNode(index=index, children=children)
        for child in children:
            child.parent = parent
        return parent

    def _rebalance_psi_tree(self, root: _PsiTreeNode, next_index: int) -> tuple[_PsiTreeNode, int]:
        """Group oversized Psi tree nodes so fanout stays within max_cluster."""
        max_children = max(2, int(self._max_cluster or 2))

        def rebalance(node: _PsiTreeNode):
            """Recursively group children when a Psi node exceeds fanout."""
            nonlocal next_index

            for child in list(node.children):
                rebalance(child)

            while len(node.children) > max_children:
                original_children = len(node.children)
                grouped_children = []
                for start in range(0, len(node.children), max_children):
                    batch = node.children[start : start + max_children]
                    if len(batch) == 1:
                        grouped_children.append(batch[0])
                        batch[0].parent = node
                    else:
                        grouped_children.append(self._create_psi_parent(next_index, batch))
                        grouped_children[-1].parent = node
                        next_index += 1
                node.children = grouped_children
                logging.info(
                    "RAPTOR Psi rebalance: node=%s children=%d grouped_to=%d max_cluster=%d",
                    node.index,
                    original_children,
                    len(grouped_children),
                    max_children,
                )

        rebalance(root)
        return self._root(root), next_index

    def _build_exact_psi_structure(
        self,
        nodes: list[_PsiTreeNode],
        next_index: int,
        task_id: str = "",
    ) -> tuple[_PsiTreeNode, int, int]:
        """Build an exact Psi subtree for a bounded node set."""
        if len(nodes) == 1:
            return nodes[0], next_index, 0

        ranked_pairs = self._rank_leaf_pairs(nodes)
        union_find = _PsiUnionFind(len(nodes))
        merges = 0
        for left_idx, right_idx in ranked_pairs:
            self._check_task_canceled(task_id, "Psi tree construction")
            if union_find.union(int(left_idx), int(right_idx)):
                merges += 1
            if merges == len(nodes) - 1:
                break

        local_nodes = {idx: node for idx, node in enumerate(nodes)}
        tree = union_find.tree
        children_by_parent = {}
        for child_idx, parent_idx in enumerate(tree):
            if child_idx not in local_nodes:
                local_nodes[child_idx] = _PsiTreeNode(index=next_index)
                next_index += 1
            if parent_idx == -1:
                continue
            children_by_parent.setdefault(parent_idx, []).append(child_idx)
            if parent_idx not in local_nodes:
                local_nodes[parent_idx] = _PsiTreeNode(index=next_index)
                next_index += 1

        for parent_idx, child_indices in children_by_parent.items():
            parent = local_nodes[parent_idx]
            parent.children = [local_nodes[child_idx] for child_idx in child_indices]
            for child in parent.children:
                child.parent = parent

        roots = [local_nodes[idx] for idx, parent_idx in enumerate(tree) if parent_idx == -1 and idx in local_nodes]
        root = max(roots, key=lambda node: node.index)
        return root, next_index, merges

    def _build_bucketed_psi_structure(
        self,
        nodes: list[_PsiTreeNode],
        next_index: int,
        task_id: str = "",
    ) -> tuple[_PsiTreeNode, int, int]:
        """Build large Psi trees by exact-ranking bounded buckets, then bucket roots."""
        buckets = self._split_psi_buckets(nodes)
        logging.info(
            "RAPTOR Psi bucketed build: nodes=%d buckets=%d bucket_size=%d exact_max_leaves=%d",
            len(nodes),
            len(buckets),
            self._psi_bucket_size,
            self._psi_exact_max_leaves,
        )

        bucket_roots = []
        merges = 0
        for bucket in buckets:
            bucket_root, next_index, bucket_merges = self._build_psi_structure_from_nodes(bucket, next_index, task_id)
            self._assign_prototype_embeddings(bucket_root)
            bucket_roots.append(bucket_root)
            merges += bucket_merges

        if len(bucket_roots) == 1:
            return bucket_roots[0], next_index, merges

        root, next_index, root_merges = self._build_psi_structure_from_nodes(bucket_roots, next_index, task_id)
        return root, next_index, merges + root_merges

    def _build_psi_structure_from_nodes(
        self,
        nodes: list[_PsiTreeNode],
        next_index: int,
        task_id: str = "",
    ) -> tuple[_PsiTreeNode, int, int]:
        """Build Psi structure exactly for small sets and bucket large sets."""
        if len(nodes) <= self._psi_exact_max_leaves:
            return self._build_exact_psi_structure(nodes, next_index, task_id)
        return self._build_bucketed_psi_structure(nodes, next_index, task_id)

    def _build_psi_structure(self, chunks, task_id: str = "") -> tuple[_PsiTreeNode, list[_PsiTreeNode]]:
        """Build the Psi merge tree from original chunk embeddings.

        ``chunks`` is expected in the normalized 3-tuple shape
        ``(text, vec, source_chunk_ids)`` — leaves are seeded with
        their own source ids, internal nodes get their ids set during
        layer materialization in ``_build_psi_layers``.
        """
        leaves = [
            _PsiTreeNode(
                index=i,
                text=item[0],
                embedding=np.asarray(item[1]),
                source_chunk_ids=list(item[2] if len(item) > 2 else []),
            )
            for i, item in enumerate(chunks)
        ]
        if len(leaves) == 1:
            return leaves[0], leaves

        root, next_index, merges = self._build_psi_structure_from_nodes(leaves, len(leaves), task_id)
        root, _ = self._rebalance_psi_tree(root, next_index)
        logging.info(
            "RAPTOR Psi tree built: leaves=%d merges=%d root_fanout=%d",
            len(leaves),
            merges,
            len(root.children),
        )
        return root, leaves

    @staticmethod
    def _psi_layers(root: _PsiTreeNode) -> dict[int, list[_PsiTreeNode]]:
        """Collect non-leaf Psi nodes by height for bottom-up summarization."""
        layers = {}

        def height(node: _PsiTreeNode) -> int:
            """Return node height while collecting internal nodes by layer."""
            if not node.children:
                return 0
            node_height = max(height(child) for child in node.children) + 1
            layers.setdefault(node_height, []).append(node)
            return node_height

        height(root)
        return layers

    async def _build_psi_layers(self, chunks, callback=None, task_id: str = ""):
        """Materialize Psi tree layers as summary chunks."""
        layers = [(0, len(chunks))]
        root, _ = self._build_psi_structure(chunks, task_id=task_id)

        for layer_idx, (_, nodes) in enumerate(sorted(self._psi_layers(root).items()), start=1):
            layer_start = len(chunks)

            async def summarize_node(node: _PsiTreeNode):
                """Summarize one Psi internal node if its children have text.

                Also propagates leaf provenance: the node's
                ``source_chunk_ids`` becomes the order-preserving deduped
                union of every child's ``source_chunk_ids``. Because
                children at this layer have already been processed (leaves
                first, then bottom-up), each child carries the full set
                of leaf ids underneath it — so the union here is the
                complete leaf set this summary covers.
                """
                texts = [child.text for child in node.children if child.text]
                if not texts:
                    logging.warning("RAPTOR Psi node %s skipped because it has no child text to summarize", node.index)
                    return None
                result = await self._summarize_texts(texts, callback, task_id)
                if result is None:
                    logging.warning("RAPTOR Psi node %s skipped because summarization failed", node.index)
                    return None
                _, node.text, node.embedding = result
                merged_ids: list[str] = []
                seen: set[str] = set()
                for child in node.children:
                    for src in child.source_chunk_ids:
                        if src and src not in seen:
                            seen.add(src)
                            merged_ids.append(src)
                node.source_chunk_ids = merged_ids
                return node

            tasks = [asyncio.create_task(summarize_node(node)) for node in nodes]
            try:
                summarized_nodes = await asyncio.gather(*tasks, return_exceptions=False)
            except Exception as e:
                logging.error(f"Error in RAPTOR Psi tree processing: {e}")
                for task in tasks:
                    task.cancel()
                await asyncio.gather(*tasks, return_exceptions=True)
                raise

            summarized_nodes = [node for node in summarized_nodes if node is not None]
            for node in summarized_nodes:
                chunks.append((node.text, node.embedding, list(node.source_chunk_ids)))

            if len(chunks) > layer_start:
                layers.append((layer_start, len(chunks)))
                logging.info(
                    "RAPTOR Psi layer materialized: layer=%d nodes=%d summaries=%d",
                    layer_idx,
                    len(nodes),
                    len(chunks) - layer_start,
                )
                if callback:
                    callback(msg="Build one Psi-RAG layer: {} -> {}".format(len(nodes), len(chunks) - layer_start))
            else:
                logging.warning("RAPTOR Psi layer %d produced no summaries; stopping materialization", layer_idx)
                break

        return chunks, layers

    async def __call__(
        self,
        chunks,
        random_state,
        callback=None,
        task_id: str = "",
        is_tree: bool = False,
    ):
        """Build summary chunks and layer boundaries for RAPTOR retrieval.

        ``chunks`` accepts either the legacy 2-tuple shape
        ``(text, vec)`` or the provenance-carrying 3-tuple shape
        ``(text, vec, source_chunk_ids)`` where ``source_chunk_ids`` is
        the list of original chunk ids that produced this entry. Output
        always uses the 3-tuple shape so every appended summary carries
        its leaves' ids. ``[]`` is left in the slot for a leaf whose id
        was missing — see the caller for the normalization rules.

        Return shapes:
          * ``is_tree=False`` (default) — original behavior: returns
            ``(chunks, layers)`` where ``chunks`` is the flat list
            (originals + summaries) and ``layers`` is the per-level
            index range ``[(start, end), ...]``.
          * ``is_tree=True`` — returns a hierarchical tree dict via
            ``_materialize_tree``. Supported for the classic builder
            only; raises ``NotImplementedError`` for PSI_TREE_BUILDER
            (PSI's hyperedge-driven summarization doesn't form a strict
            parent-of relation). Returns ``None`` when there's nothing
            to materialize.
        """
        if len(chunks) <= 1:
            return None if is_tree else ([], [])

        # Normalize input to the 3-tuple shape. Reject empties / bad
        # vectors at the same time the legacy path used to.
        def _normalize(item):
            if len(item) >= 3:
                text, vec, src = item[0], item[1], item[2]
            else:
                text, vec = item[0], item[1]
                src = []
            if not text or vec is None or len(vec) <= 0:
                return None
            # Defensive: a leaf should carry a list of strings. Drop
            # falsy entries so we don't propagate empty ids upward.
            if isinstance(src, (list, tuple)):
                src = [s for s in src if s]
            else:
                src = [src] if src else []
            return (text, vec, list(src), "")

        normalized = [t for t in (_normalize(c) for c in chunks) if t is not None]
        if len(normalized) <= 1:
            return None if is_tree else (normalized, [(0, len(normalized))])
        chunks = normalized

        if self._tree_builder == PSI_TREE_BUILDER:
            if is_tree:
                raise NotImplementedError(
                    "is_tree=True is not supported for PSI_TREE_BUILDER",
                )
            logging.info("RAPTOR: using %s tree builder for %d chunks", self._tree_builder, len(chunks))
            return await self._build_psi_layers(chunks, callback, task_id)

        # ``parent_child_map`` records each summary's immediate
        # children so ``_materialize_tree`` can walk back into a tree
        # when ``is_tree`` is set. Always populated (cheap) so the
        # tree path is just a return-shape choice at the end.
        parent_child_map: dict[int, list[int]] = {}
        n_originals = len(chunks)

        layers = [(0, len(chunks))]
        start, end = 0, len(chunks)

        @timeout(60 * 20)
        async def summarize(ck_idx: list[int]):
            """Summarize one classic RAPTOR cluster into the chunk list.

            On success appends ``(summary_text, summary_vec, src_ids)``
            where ``src_ids`` is the order-preserving deduped union of
            the ``source_chunk_ids`` of every chunk indexed in
            ``ck_idx`` — i.e. the full leaf set that contributed to
            the cluster, even through nested summaries.
            """
            nonlocal chunks

            texts = [chunks[i][0] for i in ck_idx]
            result = await self._summarize_texts(texts, callback, task_id)
            if result is not None:
                # ``dict.fromkeys`` is the cheapest way to de-dup a
                # list of strings while preserving first-seen order.
                merged_ids: list[str] = []
                seen: set[str] = set()
                for i in ck_idx:
                    for src in chunks[i][2]:
                        if src and src not in seen:
                            seen.add(src)
                            merged_ids.append(src)
                summary_ti, summary_text, summary_vec = result
                chunks.append((summary_text, summary_vec, merged_ids, summary_ti))
                # Index of the just-appended summary; map it to its
                # immediate children for the tree materializer below.
                parent_child_map[len(chunks) - 1] = list(ck_idx)

        while end - start > 1:
            self._check_task_canceled(task_id, "layer processing")

            # ``chunks`` is a mix of 3-tuples (layer-0 originals from
            # _normalize) and 4-tuples (summaries appended by
            # summarize). Vector is always at index 1 in both shapes,
            # so use positional access — the older ``_, embd, _, _``
            # form crashed on layer-0 entries.
            embeddings = [entry[1] for entry in chunks[start:end]]
            if end - start <= self._small_layer_collapse:
                # Too few nodes for meaningful sub-clustering. Skip the
                # clustering pass entirely and summarize the whole layer
                # into one parent, so the upper tree doesn't descend one
                # node per layer (N -> N-1 -> N-2 -> ... each a full
                # clustering + summarize pass).
                await summarize(list(range(start, end)))
                produced = len(chunks) - end
                if produced == 0:
                    logging.warning("RAPTOR layer produced no summaries; stopping materialization")
                    break
                logging.info(
                    "RAPTOR small-N collapse: layer of %d node(s) [%d:%d] collapsed into %d summary; stopping at tree top",
                    end - start, start, end, produced,
                )
                layers.append((end, len(chunks)))
                if callback:
                    callback(msg="Cluster one layer: {} -> {} (small-N collapse)".format(end - start, produced))
                break

            n_clusters, lbls = self.clustering(
                embeddings,
                random_state=random_state,
                task_id=task_id,
            )

            # Loop-termination guarantee. The outer ``while end - start > 1``
            # relies on each layer strictly shrinking the input count. If
            # the clusterer degenerates and returns one cluster per input,
            # every "cluster" is a single chunk, ``summarize()`` produces
            # one summary per input, and ``produced == end - start`` —
            # the same count carries into the next iteration and the loop
            # spins forever, logging "Cluster one layer: N -> N".
            #
            # Collapse everything at this level into a single cluster so
            # the layer produces exactly one summary. The tree gets a
            # taller-than-usual "single trunk" segment at this depth
            # instead of an infinite loop; downstream consumers only care
            # that ``layers`` is monotonically shrinking.
            if n_clusters >= len(embeddings):
                logging.warning(
                    "RAPTOR clustering did not reduce input count (%d inputs → %d clusters); collapsing this layer into a single summary to prevent a non-terminating loop",
                    len(embeddings),
                    n_clusters,
                )
                n_clusters = 1
                lbls = [0] * len(embeddings)

            tasks = []
            for c in range(n_clusters):
                ck_idx = [i + start for i in range(len(lbls)) if lbls[i] == c]
                assert len(ck_idx) > 0
                self._check_task_canceled(task_id, "before cluster processing")
                tasks.append(asyncio.create_task(summarize(ck_idx)))
            try:
                await asyncio.gather(*tasks, return_exceptions=False)
            except Exception as e:
                logging.error(f"Error in RAPTOR cluster processing: {e}")
                for t in tasks:
                    t.cancel()
                await asyncio.gather(*tasks, return_exceptions=True)
                raise

            produced = len(chunks) - end
            assert produced <= n_clusters, "{} vs. {}".format(produced, n_clusters)
            if produced < n_clusters:
                logging.warning(
                    "RAPTOR layer produced %d/%d cluster summaries; skipped %d cluster(s) due to errors",
                    produced,
                    n_clusters,
                    n_clusters - produced,
                )
            if produced == 0:
                logging.warning("RAPTOR layer produced no summaries; stopping materialization")
                break
            layers.append((end, len(chunks)))
            if callback:
                callback(msg="Cluster one layer: {} -> {}".format(end - start, produced))
            start = end
            end = len(chunks)

        if is_tree:
            return self._materialize_tree(chunks, layers, parent_child_map, n_originals), []
        return chunks, layers

    @staticmethod
    def _materialize_tree(chunks, layers, parent_child_map, n_originals):
        """Walk ``parent_child_map`` from the top layer down to layer-1
        and emit the user-facing tree dict. See ``__call__``'s
        ``is_tree=True`` contract for the shape.
        chunks: [(summary_text, summary_vec, merged_ids, summary_ti)]"""
        if not layers or len(chunks) == 0:
            return None
        top_start, top_end = layers[-1]
        if top_end <= top_start:
            return None

        def _title_at(idx: int) -> str:
            # Summary tuples are (text, vec, merged_ids, summary_ti)
            # — title is the 4th slot. Layer-0 originals are 3-tuples
            # and don't appear as tree nodes themselves (they collapse
            # into source_chunk_ids on their layer-1 parent).
            return chunks[idx][3] if len(chunks[idx]) >= 4 else ""

        def _desc_at(idx: int) -> str:
            return chunks[idx][0] if chunks[idx] else ""

        def _build_node(idx: int) -> dict:
            children_idx = parent_child_map.get(idx, [])
            # If every immediate child is a layer-0 original, this
            # node is a "leaf" in the tree contract — collapse to
            # source_chunk_ids.
            if children_idx and all(c < n_originals for c in children_idx):
                ids: list[str] = []
                seen: set[str] = set()
                for c in children_idx:
                    for s in chunks[c][2]:
                        if s and s not in seen:
                            seen.add(s)
                            ids.append(s)
                return {"title": _title_at(idx), "source_chunk_ids": ids, "description": _desc_at(idx)}
            return {"children": [_build_node(c) for c in children_idx], "title": _title_at(idx), "description": _desc_at(idx)}

        top_nodes = [_build_node(i) for i in range(top_start, top_end)]
        if len(top_nodes) == 1:
            return top_nodes[0]
        # Multiple top-layer summaries — clustering didn't collapse to
        # a single root. Wrap in a synthetic root so the caller always
        # sees one dict.
        return {"title": "(root)", "children": top_nodes}
