#
#  Copyright 2025 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.
#
"""Shared helpers for the knowlege_compile pipelines (structure + wiki).

Both ``structure.py`` (compile_structure_from_text / merge_compiled_structures)
and ``wiki.py`` (the MAP→REDUCE→PLAN→REFINE artifact pipeline) need the same set
of plumbing: encode-through-LLMBundle, stable id minting, search-tokenizer
pairs, order-preserving chunk-id unions, defensive LLMBundle validation, the
``chat_mdl.max_length * INPUT_UTILIZATION - prompt_overhead`` token-budget
calculation, and thin ES I/O wrappers.

Anything in this module is meant to be:
  - LLMBundle-aware but provider-agnostic;
  - Safe to import from either pipeline without circular references;
  - Synchronous unless an awaitable behaviour is required.

Heavier shared logic that is conceptually identical but happens to differ in
shape between the two pipelines (e.g. pairwise-cosine dedup, LLM "are these
the same?" batching) intentionally stays in each pipeline file for now —
extract those only when their shapes converge.
"""

from __future__ import annotations

import asyncio
import logging
import string
from typing import Any, Awaitable, Callable, Iterable, Optional

import xxhash

from common.misc_utils import thread_pool_exec
from common.token_utils import num_tokens_from_string
from rag.nlp import rag_tokenizer
from rag.prompts.generator import INPUT_UTILIZATION, gen_json, split_chunks


def knowledge_compile_gen_conf(chat_mdl, gen_conf: Optional[dict] = None) -> dict:
    """Add model-specific reasoning controls for knowledge compilation only."""
    conf = dict(gen_conf or {})
    model_config = getattr(chat_mdl, "model_config", None)
    if not isinstance(model_config, dict):
        model_config = {}
    model_name = str(model_config.get("llm_name") or getattr(chat_mdl, "llm_name", "")).lower()

    if "deepseek-v4" in model_name:
        conf["max_completion_tokens"] = 32768
        extra_body = conf.get("extra_body")
        extra_body = dict(extra_body) if isinstance(extra_body, dict) else {}
        extra_body["thinking"] = {"type": "disabled"}
        conf["extra_body"] = extra_body
    elif "qwen3" in model_name:
        # chat_model.py maps this flag to the provider-specific request body.
        # -preview variants (e.g. qwen3.8-max-preview) only accept
        # enable_thinking=True on their API endpoint.
        conf["enable_thinking"] = True if "-preview" in model_name else False
    else:
        # LiteLLM maps this common control for providers that support it and
        # drops it for providers that do not. Keep model-specific overrides
        # above because some APIs use a different, explicit request field.
        conf["reasoning_effort"] = "none"

    return conf


# ---------------------------------------------------------------------------
# ID minting
# ---------------------------------------------------------------------------


def stable_row_id(*parts) -> str:
    """xxh64 hexdigest of ``":".join(parts)`` — stable per part tuple, used
    as the ES row id when we want idempotent upserts.

    ``None`` parts become empty strings, everything else is ``str()``-ified.
    """
    key = ":".join("" if p is None else str(p) for p in parts)
    return xxhash.xxh64(key.encode("utf-8", "surrogatepass")).hexdigest()


# ---------------------------------------------------------------------------
# Embedding
# ---------------------------------------------------------------------------


async def encode(embd_mdl, texts: list[str]) -> list:
    """``LLMBundle.encode`` wrapped in ``thread_pool_exec``.

    Returns the embeddings list (drops the ``used_tokens`` count); empty
    input returns ``[]``. Caller is responsible for ensuring ``embd_mdl``
    is a real bundle — use :func:`ensure_llm_bundle` to validate at entry.
    """
    if not texts:
        return []
    embeddings, _ = await thread_pool_exec(embd_mdl.encode, texts)
    return list(embeddings)


# ---------------------------------------------------------------------------
# Tokenization for keyword search
# ---------------------------------------------------------------------------


def tokenize_for_search(text: str) -> tuple[str, str]:
    """Returns ``(content_ltks, content_sm_ltks)`` for a piece of text.

    Empty / non-string input returns ``("", "")``. Used wherever we write a
    searchable ES row that needs both tokenizations.
    """
    if not isinstance(text, str) or not text:
        return "", ""
    ltks = rag_tokenizer.tokenize(text)
    if not ltks:
        return "", ""
    sm = rag_tokenizer.fine_grained_tokenize(ltks)
    return ltks, sm


# ---------------------------------------------------------------------------
# Order-preserving union of string lists
# ---------------------------------------------------------------------------


def union_ordered(*lists: Optional[Iterable]) -> list[str]:
    """Concatenate iterables and dedupe, preserving first-seen order.
    Falsy values and non-strings are silently dropped.
    """
    seen_set: set[str] = set()
    seen: list[str] = []
    for lst in lists:
        if not lst:
            continue
        for v in lst:
            if not v or not isinstance(v, str):
                continue
            if v in seen_set:
                continue
            seen_set.add(v)
            seen.append(v)
    return seen


# ---------------------------------------------------------------------------
# Token-budget calculation for split_chunks
# ---------------------------------------------------------------------------


def make_input_budget(
    chat_mdl,
    *prompts: str,
    floor: int = 1024,
    utilization: float = INPUT_UTILIZATION,
) -> int:
    """``chat_mdl.max_length * utilization - num_tokens(sum of prompts)``,
    floored at ``floor``.

    Mirrors the budget idiom used by ``compile_structure_from_text`` and
    ``wiki_map_from_chunks``: caller passes the constant prompt scaffolding
    (system prompt + user template) — ``split_chunks`` then sizes batches
    to leave that much room.
    """
    overhead = num_tokens_from_string("".join(p or "" for p in prompts))
    budget = int(chat_mdl.max_length * utilization) - overhead
    return max(budget, floor)


# ---------------------------------------------------------------------------
# Defensive LLMBundle validation
# ---------------------------------------------------------------------------


def ensure_llm_bundle(mdl, method: str, *, label: str = "model"):
    """Return ``mdl`` if it exposes ``method``; otherwise try to unwrap a
    tuple, otherwise return ``None`` and log an error.

    Common cause for tuple inputs at call sites: ``LLMBundle.encode()`` and
    similar methods return ``(embeddings, used_tokens)``. If a caller stores
    the *result* of ``encode()`` into a variable named like
    ``embedding_model`` and passes that in, we end up with a tuple here.
    We unwrap with a warning so the pipeline keeps working while the caller
    is fixed.
    """
    if hasattr(mdl, method):
        return mdl
    if isinstance(mdl, tuple) and mdl and hasattr(mdl[0], method):
        logging.warning(
            "%s arrived as a %s; unwrapping to first element (check the call site — was %s()'s return value passed instead of the LLMBundle?)",
            label,
            type(mdl).__name__,
            method,
        )
        return mdl[0]
    logging.error(
        "%s has no .%s method (type=%s); aborting",
        label,
        method,
        type(mdl).__name__,
    )
    return None


# ---------------------------------------------------------------------------
# ES I/O wrappers
# ---------------------------------------------------------------------------


async def doc_storage_search(
    select_fields: list[str],
    condition: dict,
    *,
    tenant_id: str,
    kb_ids: list[str],
    match_expressions: list | None = None,
    offset: int = 0,
    limit: int = 1000,
    label: str = "doc_storage_search",
) -> dict:
    """Thin wrapper around ``docStoreConn.search`` + ``get_fields``.

    Returns ``{row_id: row_dict}``. Returns ``{}`` on failure (with a
    logged exception). ``label`` is included in the failure log so each
    call site is identifiable.
    """
    from common import settings
    from common.doc_store.doc_store_base import OrderByExpr
    from rag.nlp import search as _rag_search

    index = _rag_search.index_name(tenant_id)
    try:
        res = await thread_pool_exec(
            settings.docStoreConn.search,
            select_fields,
            [],
            condition,
            match_expressions or [],
            OrderByExpr(),
            offset,
            limit,
            index,
            kb_ids,
        )
        return settings.docStoreConn.get_fields(res, select_fields) or {}
    except Exception:
        logging.exception("%s failed (condition=%r)", label, condition)
        return {}


async def doc_storage_insert(
    rows: list[dict],
    tenant_id: str,
    kb_id: str,
    *,
    label: str = "doc_storage_insert",
) -> None:
    """Bulk insert wrapped in ``thread_pool_exec``. Logs on failure."""
    if not rows:
        return
    from common import settings
    from rag.nlp import search as _rag_search

    index = _rag_search.index_name(tenant_id)
    try:
        await thread_pool_exec(settings.docStoreConn.insert, rows, index, kb_id)
    except Exception:
        logging.exception("%s failed (%d row(s))", label, len(rows))


async def doc_storage_delete(
    condition: dict,
    tenant_id: str,
    kb_id: str,
    *,
    label: str = "doc_storage_delete",
) -> None:
    """Bulk delete wrapped in ``thread_pool_exec``. Best-effort; logs on
    failure (some callers rely on id-based upsert as a fallback)."""
    from common import settings
    from rag.nlp import search as _rag_search

    index = _rag_search.index_name(tenant_id)
    try:
        await thread_pool_exec(settings.docStoreConn.delete, condition, index, kb_id)
    except Exception:
        logging.debug("%s failed (condition=%r); caller may rely on id-upsert", label, condition)


async def doc_storage_upsert_one(
    filter_condition: dict,
    row: dict,
    tenant_id: str,
    kb_id: str,
    *,
    label: str = "doc_storage_upsert_one",
) -> None:
    """Delete-by-filter then insert. Used when an in-place update would
    require knowing the existing row's id and we'd rather drop+re-create.

    Best-effort delete (failures are debug-logged) followed by the insert.
    Set ``row["id"]`` to a stable value derived from the filter
    (:func:`stable_row_id`) so id-based dedup at the connector catches any
    race that bypasses the delete.
    """
    await doc_storage_delete(filter_condition, tenant_id, kb_id, label=f"{label}.delete")
    await doc_storage_insert([row], tenant_id, kb_id, label=f"{label}.insert")


# ---------------------------------------------------------------------------
# Doc-vector field discovery
# ---------------------------------------------------------------------------


def find_vec_field(doc: dict) -> tuple[Optional[str], Optional[list]]:
    """Locate the ``q_<dim>_vec`` field on an ES doc dict. Returns
    ``(field_name, vec)`` or ``(None, None)`` if the doc carries no
    embedding."""
    for k, v in doc.items():
        if isinstance(k, str) and k.startswith("q_") and k.endswith("_vec"):
            return k, v
    return None, None


# ---------------------------------------------------------------------------
# Chunked-LLM pipeline engine
# ---------------------------------------------------------------------------
#
# Both artifact MAP and compile_structure_from_text follow the same outer shape:
#
#   1. Filter chunks (drop empty text, optionally skip a "resume" set);
#   2. Pack remaining chunks into batches via ``split_chunks`` sized to leave
#      room for the prompt scaffolding;
#   3. Run an LLM-driven ``process_batch`` over each batch in parallel under
#      an ``asyncio.Semaphore(max_workers)``;
#   4. Aggregate the per-batch results into a single value.
#
# The inner LLM call shape diverges between the pipelines — artifact uses a
# single ``gen_json`` per batch with ``[CHUNK_ID Cn]``-labelled bodies,
# structure uses two ``gen_json`` calls (nodes then edges) with ``---``
# separators and no per-chunk attribution. That divergence lives in each
# pipeline's ``process_batch`` closure; this engine only owns the scaffold.


def _default_chunk_text(chunk: dict) -> str:
    if not isinstance(chunk, dict):
        return ""
    text = chunk.get("text") or chunk.get("content_with_weight") or chunk.get("content") or ""
    return text if isinstance(text, str) else ""


def _default_label(position_in_batch: int) -> str:
    return f"C{position_in_batch + 1}"


def build_chunk_batches(
    chunks: list[dict],
    chat_mdl,
    *,
    prompt_overhead_tokens: int,
    resume_chunk_ids: Optional[set[str]] = None,
    scrub_text: Optional[Callable[[str], str]] = None,
    label_fn: Callable[[int], str] = _default_label,
    chunk_text_picker: Optional[Callable[[dict], str]] = None,
    budget_floor: int = 1024,
    batch_size_cap: Optional[int] = None,
    window_fraction: Optional[float] = None,
) -> tuple[list[list[dict]], dict]:
    """Filter chunks, pack into batches, return per-batch entries.

    Each batch entry is ``{"label": str, "chunk_id": str, "text": str}``
    where ``label`` is per-batch positional (default ``C1``, ``C2``, …) and
    ``text`` is the post-scrub chunk body. Empty or resume-skipped chunks
    are dropped.

    Two packing modes:
      - **Default (split_chunks)**: ``input_budget`` derived from
        ``chat_mdl.max_length * INPUT_UTILIZATION - prompt_overhead_tokens``.
        Used by ``structure.py`` and the legacy artifact MAP path.
      - **Cap+fraction (greedy)**: when ``batch_size_cap`` is provided,
        chunks are packed greedily with two cutoffs — chunk-count exceeds
        ``batch_size_cap`` OR accumulated tokens exceed
        ``chat_mdl.max_length * window_fraction``. This is the artifact
        compilation rule (BS=8, window=0.5).

    Returns ``(batches, info)`` where ``info`` is a small stats dict.
    """
    if not chunks:
        return [], {"total": 0, "kept": 0, "skipped_resume": 0, "skipped_empty": 0, "input_budget": 0, "n_batches": 0}

    picker = chunk_text_picker or _default_chunk_text
    resume_set = resume_chunk_ids or set()

    chunk_ids: list[str] = []
    chunk_texts: list[str] = []
    skipped_resume = 0
    skipped_empty = 0

    for chunk in chunks:
        cid = chunk.get("id") or chunk.get("chunk_id")
        if not cid:
            skipped_empty += 1
            continue
        if cid in resume_set:
            skipped_resume += 1
            continue
        text = picker(chunk)
        if not text or not text.strip():
            skipped_empty += 1
            continue
        if scrub_text is not None:
            text = scrub_text(text)
            if not text or not text.strip():
                skipped_empty += 1
                continue
        chunk_ids.append(cid)
        chunk_texts.append(text)

    if not chunk_texts:
        return [], {
            "total": len(chunks),
            "kept": 0,
            "skipped_resume": skipped_resume,
            "skipped_empty": skipped_empty,
            "input_budget": 0,
            "n_batches": 0,
        }

    batches: list[list[dict]] = []
    input_budget: int

    if batch_size_cap is not None:
        # Artifact mode — greedy bin-packing with chunk-count + token caps.
        fraction = window_fraction if window_fraction is not None else 0.5
        token_cap = max(int(chat_mdl.max_length * fraction), budget_floor)
        input_budget = token_cap

        current: list[dict] = []
        current_tks = 0
        for idx, text in enumerate(chunk_texts):
            tks = num_tokens_from_string(text)
            would_overflow_count = len(current) >= batch_size_cap
            would_overflow_tokens = current and (current_tks + tks > token_cap)
            if would_overflow_count or would_overflow_tokens:
                batches.append(current)
                current = []
                current_tks = 0
            current.append(
                {
                    "label": label_fn(len(current)),
                    "chunk_id": chunk_ids[idx],
                    "text": text,
                }
            )
            current_tks += tks
        if current:
            batches.append(current)
    else:
        input_budget = max(
            int(chat_mdl.max_length * INPUT_UTILIZATION) - prompt_overhead_tokens,
            budget_floor,
        )

        raw_batches = split_chunks(chunk_texts, input_budget) or []
        for batch in raw_batches:
            packed: list[dict] = []
            for position, item in enumerate(batch):
                for idx, text in item.items():
                    packed.append(
                        {
                            "label": label_fn(position),
                            "chunk_id": chunk_ids[idx],
                            "text": text,
                        }
                    )
            if packed:
                batches.append(packed)

    info = {
        "total": len(chunks),
        "kept": len(chunk_texts),
        "skipped_resume": skipped_resume,
        "skipped_empty": skipped_empty,
        "input_budget": input_budget,
        "n_batches": len(batches),
    }
    return batches, info


async def run_chunked_pipeline(
    batches: list[list[dict]],
    *,
    process_batch: Callable[..., Awaitable[Any]],
    aggregate: Optional[Callable[[list[Any]], Any]] = None,
    max_workers: int = 6,
    callback: Optional[Callable] = None,
    log_prefix: str = "chunked_pipeline",
) -> Any:
    """Run ``process_batch`` over each batch in parallel.

    ``process_batch`` is called as
    ``await process_batch(entries: list[dict], batch_idx: int, total: int)``
    and may return anything; ``aggregate`` (if given) is called with the
    list of per-batch results and its return value is the engine's return.
    Without ``aggregate`` the raw per-batch results list is returned.

    Cancel-on-error semantics: if any task raises, all sibling tasks are
    cancelled and the exception propagates.
    """
    if not batches:
        return aggregate([]) if aggregate else []

    total = len(batches)
    worker_count = max_workers if max_workers and max_workers > 0 else 6
    work_queue: asyncio.Queue[tuple[int, list[dict]] | None] = asyncio.Queue(maxsize=worker_count)
    results: list[Any] = [None] * total
    completed: list[bool] = [False] * total

    async def _producer() -> None:
        for idx, entries in enumerate(batches):
            await work_queue.put((idx, entries))
        for _ in range(worker_count):
            await work_queue.put(None)

    async def _worker() -> None:
        while True:
            item = await work_queue.get()
            if item is None:
                work_queue.task_done()
                return
            idx, entries = item
            try:
                results[idx] = await process_batch(entries, idx, total)
                completed[idx] = True
            finally:
                work_queue.task_done()

    producer = asyncio.create_task(_producer())
    workers = [asyncio.create_task(_worker()) for _ in range(worker_count)]
    try:
        await asyncio.gather(producer, *workers)
    except BaseException:
        producer.cancel()
        for t in workers:
            t.cancel()
        await asyncio.gather(producer, *workers, return_exceptions=True)
        raise

    if callback:
        try:
            callback(1.0, f"{log_prefix}: {total} batch(es) complete")
        except Exception:
            logging.debug("%s: completion callback failed", log_prefix, exc_info=True)

    ordered_results = [results[idx] for idx in range(total) if completed[idx]]
    return aggregate(ordered_results) if aggregate else ordered_results


# ---------------------------------------------------------------------------
# Bulk dedup engine — exact + embedding + LLM disambiguation
# ---------------------------------------------------------------------------
#
# Replaces wiki's _wiki_exact_dedup_entities / _wiki_exact_dedup_concepts /
# _wiki_embedding_dedup_entities / _wiki_resolve_ambiguous_entities /
# _wiki_apply_merges with one parameterised engine. structure.py's
# merge_compiled_structures uses a different algorithm (incremental
# kept-set + per-pair LLM judgement) and stays as-is.


_PUNCT_TABLE = str.maketrans("", "", string.punctuation)


DEFAULT_DISAMBIGUATE_SYSTEM = "You are a named-entity resolution assistant. Return only JSON."


def normalize_key(name) -> str:
    """Lowercase + strip whitespace + strip ASCII punctuation. Used as the
    bucket key for exact dedup."""
    if not isinstance(name, str):
        return ""
    return name.lower().strip().translate(_PUNCT_TABLE)


def _exact_dedup_by_key(
    items: list[dict],
    *,
    name_key: str,
    type_key: Optional[str] = None,
    aggregate_extra: Optional[Callable[[list[dict]], dict]] = None,
) -> list[dict]:
    """Group items by ``(normalize(item[name_key]), item.get(type_key))``.

    Canonical record per group:
      - ``<name_key>``: the most-common spelling across the group
      - ``<type_key>`` (if given): the group's shared value
      - ``aliases``: sorted union of every name + every input alias, minus
        the canonical name
      - ``mention_count``: sum of input ``mention_count`` values (defaults
        to ``1`` per missing)
      - ``chunk_ids``: order-preserving union
      - ``_norm``: the normalized key (stripped by ``bulk_dedup_items``)
      - any extras from ``aggregate_extra(group)``
    """
    groups: dict[tuple, list[dict]] = {}
    for it in items:
        if not isinstance(it, dict):
            continue
        norm = normalize_key(it.get(name_key, ""))
        if not norm:
            continue
        key = (norm, it.get(type_key) if type_key else None)
        groups.setdefault(key, []).append(it)

    canonical: list[dict] = []
    for (norm, type_val), group in groups.items():
        name_counts: dict[str, int] = {}
        for it in group:
            n = it.get(name_key, "")
            if isinstance(n, str) and n:
                name_counts[n] = name_counts.get(n, 0) + 1
        best = max(name_counts, key=lambda k: name_counts[k]) if name_counts else ""

        aliases: set[str] = set()
        chunk_id_lists: list[list] = []
        mention_count = 0
        for it in group:
            n = it.get(name_key, "")
            if isinstance(n, str) and n:
                aliases.add(n)
            for a in it.get("aliases") or []:
                if isinstance(a, str) and a:
                    aliases.add(a)
            chunk_id_lists.append(it.get("chunk_ids") or [])
            mention_count += int(it.get("mention_count") or 1)
        aliases.discard(best)

        record: dict = {
            name_key: best,
            "aliases": sorted(aliases),
            "mention_count": mention_count,
            "chunk_ids": union_ordered(*chunk_id_lists),
            "_norm": norm,
        }
        if type_key:
            record[type_key] = type_val
        if aggregate_extra is not None:
            try:
                extras = aggregate_extra(group) or {}
                if isinstance(extras, dict):
                    record.update(extras)
            except Exception:
                logging.exception("bulk_dedup: aggregate_extra failed for group %r", norm)
        canonical.append(record)

    return canonical


async def _embedding_dedup(
    canonical: list[dict],
    embd_mdl,
    *,
    name_key: str,
    type_key: Optional[str] = None,
    merge_threshold: float = 0.90,
    ambiguous_low: float = 0.75,
) -> tuple[dict[int, int], list[tuple[int, int]], Optional[list]]:
    """Blockwise pairwise cosine; same-type-only when ``type_key`` given.

    Returns ``(merged_into, ambiguous_pairs, vectors)``. ``merged_into``
    is a union-find map ``index → parent_index``. ``ambiguous_pairs`` is the
    [ambiguous_low, merge_threshold) bucket (after removing pairs already
    linked by auto-merges). ``vectors`` is ``None`` on embedding failure
    (caller should skip dedup).
    """
    n = len(canonical)
    if n <= 1:
        return {}, [], []

    names = [it.get(name_key, "") for it in canonical]
    try:
        vectors = await encode(embd_mdl, names)
    except Exception:
        logging.exception("bulk_dedup: embedding batch failed")
        return {}, [], None
    if vectors is None or len(vectors) != n:
        return {}, [], None

    try:
        import numpy as np

        matrix = np.asarray([list(v) for v in vectors], dtype=float)
        if matrix.ndim != 2 or matrix.shape[0] != n:
            raise ValueError(f"invalid embedding matrix shape: {matrix.shape}")

        # Normalize once, then calculate only bounded blocks of the upper
        # triangle. The old implementation materialized an N x N matrix,
        # which becomes a significant memory spike for large KBs.
        norms = np.linalg.norm(matrix, axis=1, keepdims=True)
        matrix = np.divide(matrix, norms, out=np.zeros_like(matrix), where=norms > 0)
    except Exception:
        logging.exception("bulk_dedup: pairwise cosine failed; skipping")
        return {}, [], vectors

    merged_into: dict[int, int] = {}

    def _root(i: int) -> int:
        while i in merged_into:
            i = merged_into[i]
        return i

    auto_pairs: list[tuple[int, int]] = []
    ambiguous_pairs: list[tuple[int, int]] = []
    block_size = 1024
    groups: dict[Any, list[int]] = {}
    for index, item in enumerate(canonical):
        groups.setdefault(item.get(type_key) if type_key else None, []).append(index)

    for group_indices in groups.values():
        for left_start in range(0, len(group_indices), block_size):
            left_indices = group_indices[left_start : left_start + block_size]
            left_vectors = matrix[left_indices]
            for right_start in range(left_start, len(group_indices), block_size):
                right_indices = group_indices[right_start : right_start + block_size]
                sims = left_vectors @ matrix[right_indices].T
                if right_start == left_start:
                    candidate_mask = np.triu(sims >= ambiguous_low, k=1)
                else:
                    candidate_mask = sims >= ambiguous_low
                rows, cols = np.nonzero(candidate_mask)
                for row, col in zip(rows.tolist(), cols.tolist(), strict=True):
                    score = float(sims[row, col])
                    if score >= merge_threshold:
                        auto_pairs.append((left_indices[row], right_indices[col]))
                    else:
                        ambiguous_pairs.append((left_indices[row], right_indices[col]))

    for i, j in auto_pairs:
        ri, rj = _root(i), _root(j)
        if ri == rj:
            continue
        if canonical[ri].get("mention_count", 0) >= canonical[rj].get("mention_count", 0):
            merged_into[rj] = ri
        else:
            merged_into[ri] = rj

    still_ambiguous = [(i, j) for i, j in ambiguous_pairs if _root(i) != _root(j)]
    return merged_into, still_ambiguous, vectors


async def _resolve_ambiguous_pairs(
    canonical: list[dict],
    ambiguous_pairs: list[tuple[int, int]],
    merged_into: dict[int, int],
    chat_mdl,
    *,
    name_key: str,
    type_key: Optional[str] = None,
    batch_size: int = 50,
    llm_timeout: int = 60,
    system_prompt: str = DEFAULT_DISAMBIGUATE_SYSTEM,
) -> dict[int, int]:
    """LLM-judged disambiguation in bounded concurrent waves.

    Each wave filters pairs using the latest union-find roots, then runs at
    most one batch per available pooled LLM slot. Results are applied in one
    serial pass before the next wave starts, avoiding redundant calls for
    pairs merged by an earlier wave.
    """
    if not ambiguous_pairs:
        return merged_into

    def _root(i: int) -> int:
        while i in merged_into:
            i = merged_into[i]
        return i

    async def _resolve_batch(batch: list[tuple[int, int]]) -> tuple[list[tuple[int, int]], Optional[list]]:
        lines: list[str] = []
        for k, (i, j) in enumerate(batch):
            a_type = f" ({canonical[i].get(type_key, '')})" if type_key else ""
            b_type = f" ({canonical[j].get(type_key, '')})" if type_key else ""
            lines.append(f'{k + 1}. "{canonical[i].get(name_key, "")}"{a_type} vs "{canonical[j].get(name_key, "")}"{b_type}')

        user_prompt = (
            "For each pair below, determine if they refer to the same real-world entity.\n"
            f"Return a JSON array of exactly {len(batch)} booleans "
            "(true = same entity, false = different).\n"
            "Return ONLY the JSON array.\n\n" + "\n".join(lines)
        )

        try:
            res = await asyncio.wait_for(
                gen_json(
                    system_prompt,
                    user_prompt,
                    chat_mdl,
                    gen_conf=knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0}),
                ),
                timeout=llm_timeout,
            )
        except asyncio.TimeoutError:
            logging.warning("bulk_dedup: disambiguation timed out (%d pairs)", len(batch))
            return batch, None
        except Exception:
            logging.exception("bulk_dedup: disambiguation call failed (%d pairs)", len(batch))
            return batch, None

        decisions = None
        if isinstance(res, list):
            decisions = res
        elif isinstance(res, dict):
            for v in res.values():
                if isinstance(v, list):
                    decisions = v
                    break
        if not isinstance(decisions, list):
            logging.warning("bulk_dedup: disambiguation returned unexpected shape: %r", type(res))
            return batch, None
        return batch, decisions

    pool = getattr(chat_mdl, "_pool", None)
    wave_batch_count = max(1, int(getattr(pool, "max_concurrency", 10)))
    remaining = list(ambiguous_pairs)
    while remaining:
        eligible = [(i, j) for i, j in remaining if _root(i) != _root(j)]
        if not eligible:
            break
        wave_pairs = eligible[: wave_batch_count * batch_size]
        batches = [wave_pairs[start : start + batch_size] for start in range(0, len(wave_pairs), batch_size)]
        results = await asyncio.gather(*(_resolve_batch(batch) for batch in batches))
        for batch, decisions in results:
            if decisions is None:
                continue
            for k, (i, j) in enumerate(batch):
                verdict = decisions[k] if k < len(decisions) else False
                if not verdict:
                    continue
                ri, rj = _root(i), _root(j)
                if ri == rj:
                    continue
                if canonical[ri].get("mention_count", 0) >= canonical[rj].get("mention_count", 0):
                    merged_into[rj] = ri
                else:
                    merged_into[ri] = rj
        remaining = eligible[len(wave_pairs) :]

    return merged_into


def _apply_dedup_merges(
    canonical: list[dict],
    merged_into: dict[int, int],
    *,
    name_key: str,
) -> list[dict]:
    """Union-find collapse: sum ``mention_count``, union ``aliases`` and
    ``chunk_ids`` per canonical."""

    def _root(i: int) -> int:
        while i in merged_into:
            i = merged_into[i]
        return i

    roots: set[int] = {_root(i) for i in range(len(canonical))}
    out: list[dict] = []
    for ri in roots:
        base = dict(canonical[ri])
        aliases: set[str] = set(base.get("aliases") or [])
        chunk_id_lists: list[list] = [base.get("chunk_ids") or []]
        mention_count = int(base.get("mention_count") or 0)
        for i, it in enumerate(canonical):
            if i == ri or _root(i) != ri:
                continue
            mention_count += int(it.get("mention_count") or 0)
            aliases.update(it.get("aliases") or [])
            n = it.get(name_key)
            if isinstance(n, str) and n:
                aliases.add(n)
            chunk_id_lists.append(it.get("chunk_ids") or [])
        aliases.discard(base.get(name_key) or "")
        base["aliases"] = sorted(aliases)
        base["mention_count"] = mention_count
        base["chunk_ids"] = union_ordered(*chunk_id_lists)
        out.append(base)
    return out


async def bulk_dedup_items(
    items: list[dict],
    *,
    name_key: str,
    type_key: Optional[str] = None,
    chat_mdl=None,
    embd_mdl=None,
    merge_threshold: float = 0.90,
    ambiguous_low: float = 0.75,
    ambiguous_batch_size: int = 50,
    disambiguate_system_prompt: str = DEFAULT_DISAMBIGUATE_SYSTEM,
    llm_timeout: int = 60,
    aggregate_extra: Optional[Callable[[list[dict]], dict]] = None,
    strip_norm_key: bool = True,
) -> list[dict]:
    """Three-phase dedup → canonical items.

    Phase 1 (always): exact dedup by ``(normalize(item[name_key]),
    item.get(type_key))`` — groups by normalized key, sums mention_count,
    unions aliases and chunk_ids, optionally adds extras via
    ``aggregate_extra(group)``.

    Phase 2 (when ``embd_mdl`` is provided AND ``len(canonical) > 1``):
    blockwise pairwise cosine over the canonical ``name_key`` values.
    Pairs at similarity ≥ ``merge_threshold`` auto-merge; pairs in
    ``[ambiguous_low, merge_threshold)`` move to phase 3. When ``type_key``
    is given, pairs are only considered when both endpoints share the same
    type. Embedding failures cause this phase (and 3) to be skipped.

    Phase 3 (when ``chat_mdl`` is provided AND ambiguous pairs remain):
    batched LLM disambiguation via ``gen_json`` — each batch asks for a
    JSON array of booleans. True verdicts join the union-find.

    Apply: union-find collapse — sum mention_count, union aliases /
    chunk_ids per canonical.

    Setting both ``chat_mdl`` and ``embd_mdl`` to ``None`` makes this an
    exact-dedup-only call (which is what artifact uses for concepts).
    """
    canonical = _exact_dedup_by_key(
        items,
        name_key=name_key,
        type_key=type_key,
        aggregate_extra=aggregate_extra,
    )

    if len(canonical) > 1 and embd_mdl is not None:
        merged_into, ambig, vectors = await _embedding_dedup(
            canonical,
            embd_mdl,
            name_key=name_key,
            type_key=type_key,
            merge_threshold=merge_threshold,
            ambiguous_low=ambiguous_low,
        )
        if vectors is None:
            logging.warning("bulk_dedup: embedding phase skipped — keeping exact-dedup result")
        else:
            if chat_mdl is not None and ambig:
                merged_into = await _resolve_ambiguous_pairs(
                    canonical,
                    ambig,
                    merged_into,
                    chat_mdl,
                    name_key=name_key,
                    type_key=type_key,
                    batch_size=ambiguous_batch_size,
                    llm_timeout=llm_timeout,
                    system_prompt=disambiguate_system_prompt,
                )
            canonical = _apply_dedup_merges(canonical, merged_into, name_key=name_key)

    if strip_norm_key:
        for it in canonical:
            it.pop("_norm", None)
    return canonical


__all__ = [
    "stable_row_id",
    "encode",
    "tokenize_for_search",
    "union_ordered",
    "make_input_budget",
    "ensure_llm_bundle",
    "doc_storage_search",
    "doc_storage_insert",
    "doc_storage_delete",
    "doc_storage_upsert_one",
    "find_vec_field",
    # New engines
    "normalize_key",
    "build_chunk_batches",
    "run_chunked_pipeline",
    "bulk_dedup_items",
    "DEFAULT_DISAMBIGUATE_SYSTEM",
]
