#
#  Copyright 2026 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.
#

"""KB-wide wiki / artifact compilation.

Extracted from ``rag.svr.task_executor_refactor.task_handler`` where the
same pipeline previously lived as a set of ``_wiki_*`` / ``_persist_wiki_*``
methods and one ``_run_wiki`` orchestrator. The public entry point is
:func:`run_wiki`.

The pipeline runs MAP per (doc, template) — each MAP call resumes from
its own ``wiki_map_extract`` ES rows — then REDUCE / PLAN / REFINE
KB-wide via ``rag.advanced_rag.knowlege_compile.wiki``. Refined pages
land in ES as searchable ``wiki_page`` rows, one
``wiki_page_topic`` row per topic, and ``wiki_entity`` /
``wiki_relation`` rows for the dataset Artifact tab's canvas graph.

Design notes:

* ``load_chunks_for_doc`` is injected rather than imported to keep the
  module decoupled from ``TaskHandler``'s streaming chunk loader.
* The eligibility loop resolves each doc's
  ``parser_config.compilation_template_group_id`` to a template list
  via the shared parser-config helper and
  ``CompilationTemplateGroupService.resolve_template_ids``.
* The persistence helpers (``persist_wiki_pages_to_es`` etc.) are
  exposed at module level for testing but are only called from
  :func:`run_wiki` in production.
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
from typing import AsyncIterator, Callable, Dict, List, Optional

import xxhash

from common import settings
from common.constants import LLMType
from common.misc_utils import thread_pool_exec
from rag.nlp import search
from rag.advanced_rag.knowlege_compile.structure import LLMCallPool
from rag.advanced_rag.knowlege_compile.wiki import (
    wiki_map_from_chunks,
    wiki_plan_from_reduction,
    wiki_reduce_from_extracts,
    wiki_refine_from_plan,
)
from rag.svr.task_executor_refactor.task_context import TaskContext


# ----- tunables ------------------------------------------------------
# Artifact-MAP tuning: how many chunks to feed per ``wiki_map_from_chunks``
# invocation. The function does its own per-call resume-set load + ES
# persist, so smaller batches mean more (small) ES round-trips but a flat
# memory footprint. 64 keeps the resume-set re-reads cheap while leaving
# room for the function's internal split_chunks packing to do real work.
WIKI_MAP_BATCH_CHUNKS = 64

# The pool limits actual MAP LLM calls rather than the surrounding batch
# tasks. This lets the next waiting batch start as soon as a model call
# finishes, without waiting for the previous batch's ES persistence work.
WIKI_MAP_LLM_POOL_SIZE = 20

# Global MAP admission limit: active calls plus calls waiting in the pool.
WIKI_MAP_MAX_PENDING = 25

# Keep only a small number of outer batches buffered. With 20 workers this
# bounds the in-memory MAP work to roughly 25 batches (20 active + 5 waiting).
WIKI_MAP_QUEUE_SIZE = 5

# REFINE pages are independent. Keep enough page workers to feed the shared
# pool; the pool itself still caps actual LLM requests at 20.
WIKI_REFINE_WORKERS = WIKI_MAP_LLM_POOL_SIZE

# Per-node cap on ``source_chunk_ids`` carried by the canvas graph blob.
# Pages can accumulate hundreds of source chunks; the graph response is
# meant for fast canvas rendering, not full provenance audit, so we trim
# each node's list. The full per-page list is still available on the
# ``wiki_page`` row the UI deep-links into.
WIKI_GRAPH_MAX_CHUNK_IDS_PER_NODE = 64

# Title + comments stamped on every ``wiki_commit`` row produced by
# :func:`run_wiki`'s regeneration path (as opposed to the dialog-edit
# path, which carries the user's own title/comments).
WIKI_REGEN_COMMIT_TITLE = "Regenerated by artifact compilation"
WIKI_REGEN_COMMIT_COMMENTS_TEMPLATE = "Auto-update via run_wiki (action={action})"
WIKI_MAP_COMPILE_KWD = "wiki_map_extract"
WIKI_REDUCE_COMPILE_KWD = "wiki_reduce_result"
WIKI_PLAN_COMPILE_KWD = "wiki_compilation_plan"
WIKI_DRAFT_COMPILE_KWD = "wiki_page_draft"
WIKI_PAGE_COMPILE_KWD = "wiki_page"
WIKI_PAGE_TOPIC_COMPILE_KWD = "wiki_page_topic"
WIKI_DERIVED_COMPILE_KWDS = (
    WIKI_REDUCE_COMPILE_KWD,
    WIKI_PLAN_COMPILE_KWD,
    WIKI_DRAFT_COMPILE_KWD,
    WIKI_PAGE_COMPILE_KWD,
    WIKI_PAGE_TOPIC_COMPILE_KWD,
    "wiki_entity",
    "wiki_relation",
    "wiki_page_graph",
)


# ----- helpers -------------------------------------------------------


def _parser_config_compilation_template_ids(parser_config, tenant_id: str) -> list[str]:
    """Resolve a doc's ``parser_config`` to compile-template ids by
    looking up configured groups. Returns ``[]`` if no group resolves."""
    from rag.svr.task_executor_refactor.chunk_post_processor import (
        _parser_config_compilation_template_group_ids,
    )
    from api.db.services.compilation_template_group_service import (
        CompilationTemplateGroupService,
    )

    template_ids: list[str] = []
    seen: set[str] = set()
    for group_id in _parser_config_compilation_template_group_ids(parser_config):
        for template_id in CompilationTemplateGroupService.resolve_template_ids(group_id, tenant_id):
            if template_id in seen:
                continue
            seen.add(template_id)
            template_ids.append(template_id)
    return template_ids


def _normalize_compilation_template_group_ids(raw) -> list[str]:
    if isinstance(raw, str):
        raw = [raw]
    if not isinstance(raw, list):
        return []
    ids: list[str] = []
    seen: set[str] = set()
    for group_id in raw:
        if not isinstance(group_id, str):
            continue
        group_id = group_id.strip()
        if group_id and group_id not in seen:
            seen.add(group_id)
            ids.append(group_id)
    return ids


def _extract_pipeline_compiler_group_ids(dsl) -> list[str]:
    if isinstance(dsl, str):
        try:
            dsl = json.loads(dsl)
        except Exception:
            return []
    if not isinstance(dsl, dict):
        return []
    components = dsl.get("components")
    if not isinstance(components, dict):
        return []

    group_ids: list[str] = []
    seen: set[str] = set()
    for component in components.values():
        if not isinstance(component, dict):
            continue
        obj = component.get("obj") if isinstance(component.get("obj"), dict) else {}
        component_name = obj.get("component_name") or component.get("component_name") or component.get("name")
        if not isinstance(component_name, str) or component_name.lower() != "compiler":
            continue
        candidates = [
            obj.get("params") if isinstance(obj.get("params"), dict) else {},
            obj,
            component.get("params") if isinstance(component.get("params"), dict) else {},
            component,
        ]
        for candidate in candidates:
            for key in ("compilation_template_group_ids", "compilation_template_group_id"):
                for group_id in _normalize_compilation_template_group_ids(candidate.get(key)):
                    if group_id not in seen:
                        seen.add(group_id)
                        group_ids.append(group_id)
    return group_ids


def _pipeline_compilation_template_ids(pipeline_id: str, tenant_id: str) -> list[str]:
    pipeline_id = (pipeline_id or "").strip()
    if not pipeline_id:
        return []
    from api.db.services.canvas_service import UserCanvasService
    from api.db.services.compilation_template_group_service import (
        CompilationTemplateGroupService,
    )

    ok, canvas = UserCanvasService.get_by_id(pipeline_id)
    if not ok or not canvas:
        return []
    template_ids: list[str] = []
    seen: set[str] = set()
    for group_id in _extract_pipeline_compiler_group_ids(getattr(canvas, "dsl", None)):
        for template_id in CompilationTemplateGroupService.resolve_template_ids(group_id, tenant_id):
            if template_id in seen:
                continue
            seen.add(template_id)
            template_ids.append(template_id)
    return template_ids


async def _wiki_existing_map_doc_ids(tenant_id: str, kb_id: str) -> set[str]:
    from common.doc_store.doc_store_base import OrderByExpr

    index = search.index_name(tenant_id)
    if not settings.docStoreConn.index_exist(index, kb_id):
        return set()

    doc_ids: set[str] = set()
    select_fields = ["id", "doc_id"]
    offset = 0
    page_size = 1000
    while True:
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                {"compile_kwd": [WIKI_MAP_COMPILE_KWD]},
                [],
                OrderByExpr(),
                offset,
                page_size,
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields) or {}
        except Exception:
            logging.exception("wiki: failed to scan MAP doc ids for kb=%s", kb_id)
            return doc_ids
        if not field_map:
            break
        for row in field_map.values():
            doc_id = row.get("doc_id")
            if isinstance(doc_id, str) and doc_id:
                doc_ids.add(doc_id)
        if len(field_map) < page_size:
            break
        offset += page_size
    return doc_ids


async def _wiki_delete_deleted_doc_state(
    tenant_id: str,
    kb_id: str,
    deleted_doc_ids: set[str],
) -> None:
    if not deleted_doc_ids:
        return

    index = search.index_name(tenant_id)
    if not settings.docStoreConn.index_exist(index, kb_id):
        return

    # 1. MAP resume rows are keyed by the real doc_id — delete outright.
    try:
        await thread_pool_exec(
            settings.docStoreConn.delete,
            {
                "compile_kwd": [WIKI_MAP_COMPILE_KWD],
                "doc_id": sorted(deleted_doc_ids),
            },
            index,
            kb_id,
        )
    except Exception:
        logging.exception(
            "wiki: failed to delete MAP rows for removed docs in kb=%s docs=%s",
            kb_id,
            sorted(deleted_doc_ids),
        )
        return

    # 2. Derived KB-scoped rows: reference-counted self-healing backstop for
    # the eager delete-time cleanup (DocumentService.remove_wiki_products).
    # Read every row referencing any deleted doc, drop the ones left with no
    # surviving owner, and shrink the rest to their surviving doc set. This
    # replaces the former blunt "delete every derived row" wipe so products
    # shared with still-present docs survive a peer's deletion.
    from common.doc_store.doc_store_base import OrderByExpr

    deleted = set(deleted_doc_ids)
    derived_kwds = list(WIKI_DERIVED_COMPILE_KWDS)
    select_fields = ["id", "source_doc_ids"]
    to_delete: list[str] = []
    to_shrink: list[tuple[str, list[str]]] = []
    offset = 0
    page_size = 1000
    while True:
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                {"compile_kwd": derived_kwds, "source_doc_ids": sorted(deleted)},
                [],
                OrderByExpr(),
                offset,
                page_size,
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields) or {}
        except Exception:
            logging.exception("wiki: failed to scan derived rows for removed docs in kb=%s", kb_id)
            return
        if not field_map:
            break
        for row_id, row in field_map.items():
            raw = row.get("source_doc_ids")
            if isinstance(raw, str):
                owners = [raw] if raw else []
            elif isinstance(raw, list):
                owners = [d for d in raw if isinstance(d, str) and d]
            else:
                owners = []
            remaining = [d for d in owners if d not in deleted]
            if remaining:
                to_shrink.append((row_id, remaining))
            else:
                to_delete.append(row_id)
        if len(field_map) < page_size:
            break
        offset += page_size

    # Drop rows with no surviving owner (delete by id in batches).
    for i in range(0, len(to_delete), page_size):
        try:
            await thread_pool_exec(
                settings.docStoreConn.delete,
                {"id": to_delete[i : i + page_size]},
                index,
                kb_id,
            )
        except Exception:
            logging.exception("wiki: failed to drop orphaned derived rows in kb=%s", kb_id)

    # Shrink rows still owned by surviving docs to just those docs.
    for row_id, remaining in to_shrink:
        try:
            await thread_pool_exec(
                settings.docStoreConn.update,
                {"id": row_id},
                {"source_doc_ids": remaining},
                index,
                kb_id,
            )
        except Exception:
            logging.exception("wiki: failed to shrink source_doc_ids for row=%s in kb=%s", row_id, kb_id)

    logging.info(
        "wiki: ref-counted cleanup for %d deleted doc(s) in kb=%s (dropped=%d, shrunk=%d)",
        len(deleted),
        kb_id,
        len(to_delete),
        len(to_shrink),
    )


def _wiki_topic_from_page(page: Dict, fallback: str = "") -> str:
    for key in ("topic", "title", "page_type"):
        value = page.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()
    return fallback.strip()


def _wiki_topic_slug(topic: str) -> str:
    slug = re.sub(r"[^a-z0-9]+", "-", topic.lower()).strip("-")
    if not slug:
        slug = xxhash.xxh64(topic.encode("utf-8", "surrogatepass")).hexdigest()[:12]
    return f"topic/{slug[:80]}"


async def _ensure_wiki_topic_rows(
    ctx: TaskContext,
    index: str,
    kb_id_str: str,
    topics_by_name: dict[str, str],
    topic_doc_ids: dict[str, list[str]] | None = None,
) -> None:
    if not topics_by_name:
        return

    from common.doc_store.doc_store_base import OrderByExpr
    from rag.nlp import rag_tokenizer

    topic_doc_ids = topic_doc_ids or {}
    topics = list(topics_by_name.keys())
    existing: set[str] = set()
    try:
        res = await thread_pool_exec(
            settings.docStoreConn.search,
            ["id", "topic_kwd"],
            [],
            {
                "compile_kwd": [WIKI_PAGE_TOPIC_COMPILE_KWD],
                "topic_kwd": topics,
            },
            [],
            OrderByExpr(),
            0,
            max(len(topics), 1),
            index,
            [ctx.kb_id],
        )
        field_map = settings.docStoreConn.get_fields(res, ["id", "topic_kwd"])
        for row in (field_map or {}).values():
            raw = row.get("topic_kwd")
            if isinstance(raw, list):
                existing.update(t for t in raw if isinstance(t, str) and t)
            elif isinstance(raw, str) and raw:
                existing.add(raw)
    except Exception:
        logging.exception(
            "wiki_persist: topic existence read failed for kb=%s; inserting stable topic ids",
            kb_id_str,
        )

    rows: list[dict] = []
    for topic, slug in topics_by_name.items():
        if topic in existing:
            # Topic row already present — refresh only its doc provenance so
            # delete-time ref-counting stays accurate as new docs contribute
            # to (or stop contributing to) the topic across recompiles.
            doc_ids = topic_doc_ids.get(topic) or []
            try:
                await thread_pool_exec(
                    settings.docStoreConn.update,
                    {"compile_kwd": WIKI_PAGE_TOPIC_COMPILE_KWD, "topic_kwd": topic},
                    {"source_doc_ids": list(doc_ids)},
                    index,
                    ctx.kb_id,
                )
            except Exception:
                logging.exception(
                    "wiki_persist: topic provenance update failed for kb=%s topic=%s",
                    kb_id_str,
                    topic,
                )
            continue
        topic_id = xxhash.xxh64(
            f"{kb_id_str}:{WIKI_PAGE_TOPIC_COMPILE_KWD}:{topic}".encode(
                "utf-8",
                "surrogatepass",
            ),
        ).hexdigest()
        content_ltks = rag_tokenizer.tokenize(topic)
        rows.append(
            {
                "id": topic_id,
                "kb_id": kb_id_str,
                "doc_id": kb_id_str,
                "compile_kwd": WIKI_PAGE_TOPIC_COMPILE_KWD,
                "topic_kwd": topic,
                "title_kwd": topic,
                "slug_kwd": slug,
                "source_doc_ids": list(topic_doc_ids.get(topic) or []),
                "content_with_weight": topic,
                "content_ltks": content_ltks,
                "content_sm_ltks": rag_tokenizer.fine_grained_tokenize(content_ltks),
                "available_int": 1,
            }
        )

    if rows:
        await thread_pool_exec(settings.docStoreConn.insert, rows, index, ctx.kb_id)


# ----- persistence ---------------------------------------------------


async def persist_wiki_pages_to_es(
    ctx: TaskContext,
    pages: List[Dict],
    embd_mdl,
) -> None:
    """Insert one ES row per generated artifact page using the
    knowledge-compilation schema:

      id                  xxh64(kb_id + ":" + slug)
      compile_kwd         "wiki_page"
      slug_kwd            page.slug
      title_kwd           page.title
      page_type_kwd       page.page_type
      topic_kwd           page.topic
      entity_names_kwd    page.entity_names
      outlinks_kwd        page.outlinks
      related_kb_pages_kwd page.related_kb_pages
      source_chunk_ids    page.source_chunk_ids
      source_doc_ids      page.source_doc_ids
      kb_id               ctx.kb_id
      content_with_weight rendered markdown
      content_ltks /
      content_sm_ltks     tokenize(content_md + summary)
      q_<dim>_vec         embed(summary)

    ``action`` is intentionally not stored — it's a planner artifact
    and has no meaning post-write.
    """
    if not pages:
        return

    from rag.nlp import rag_tokenizer
    from common.doc_store.doc_store_base import OrderByExpr
    from api.db.services.file_commit_service import (
        FileCommitService as WikiCommitService,
    )

    index = search.index_name(ctx.tenant_id)
    kb_id_str = str(ctx.kb_id)

    # Capture the prior rendered content for every slug we're about to
    # overwrite, so the per-page commit row downstream has a real diff
    # baseline. Single batch read by slug_kwd IN [...] — one round-trip
    # regardless of page count. Failures here degrade gracefully.
    target_slugs: list[str] = [(p.get("slug") or "").strip() for p in pages if isinstance(p.get("slug"), str) and p.get("slug")]
    prior_by_slug: dict[str, str] = {}
    if target_slugs:
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                ["id", "slug_kwd", "content_with_weight"],
                [],
                {"compile_kwd": [WIKI_PAGE_COMPILE_KWD], "slug_kwd": list(target_slugs)},
                [],
                OrderByExpr(),
                0,
                max(len(target_slugs), 1),
                index,
                [ctx.kb_id],
            )
            field_map = settings.docStoreConn.get_fields(
                res,
                ["id", "slug_kwd", "content_with_weight"],
            )
            for row in (field_map or {}).values():
                s = row.get("slug_kwd")
                c = row.get("content_with_weight")
                if isinstance(s, str) and isinstance(c, str):
                    prior_by_slug[s] = c
        except Exception:
            logging.exception(
                "wiki_persist: prior-content read failed for kb=%s; commit audit will treat all pages as creations",
                kb_id_str,
            )

    # Batch the summary embeddings in one model call. Empty strings are
    # swapped for a single space so the encoder doesn't reject them —
    # they still yield a vector but contribute nothing meaningful.
    summaries = [(p.get("summary") or "").strip() for p in pages]
    embed_inputs = [s if s else " " for s in summaries]
    try:
        embeddings, _ = await thread_pool_exec(embd_mdl.encode, embed_inputs)
    except Exception:
        logging.exception(
            "wiki_persist: summary embedding batch failed for kb=%s",
            kb_id_str,
        )
        return
    try:
        n_emb = len(embeddings) if embeddings is not None else 0
    except TypeError:
        n_emb = 0
    if n_emb != len(pages):
        logging.warning(
            "wiki_persist: embedding count %d != pages %d for kb=%s; aborting",
            n_emb,
            len(pages),
            kb_id_str,
        )
        return

    rows: List[Dict] = []
    topics_by_name: dict[str, str] = {}
    # Per-topic union of contributing doc ids, so topic rows can be
    # reference-counted at document-delete time like every other product.
    topic_doc_ids: dict[str, list[str]] = {}
    for page, vec in zip(pages, embeddings):
        slug = page.get("slug") or ""
        if not slug:
            continue
        title = page.get("title") or slug
        topic = _wiki_topic_from_page(page, title)
        summary = page.get("summary") or ""
        content_md = page.get("content_md_rendered") or page.get("content_md") or page.get("content_md_raw") or ""

        vec_list = vec.tolist() if hasattr(vec, "tolist") else list(vec)
        if not vec_list:
            logging.warning(
                "wiki_persist: empty embedding for slug=%s; skipping",
                slug,
            )
            continue

        text_for_search = (content_md + "\n\n" + summary).strip()
        content_ltks = rag_tokenizer.tokenize(text_for_search) if text_for_search else ""
        content_sm_ltks = rag_tokenizer.fine_grained_tokenize(content_ltks) if content_ltks else ""

        row_id = xxhash.xxh64(
            f"{kb_id_str}:{slug}".encode("utf-8", "surrogatepass"),
        ).hexdigest()

        rows.append(
            {
                "id": row_id,
                "kb_id": kb_id_str,
                "doc_id": kb_id_str,  # sentinel; KB-scoped row, real provenance in source_doc_ids
                "compile_kwd": WIKI_PAGE_COMPILE_KWD,
                "slug_kwd": slug,
                "title_kwd": title,
                "page_type_kwd": page.get("page_type") or "concept",
                "topic_kwd": topic,
                "entity_names_kwd": list(page.get("entity_names") or []),
                "outlinks_kwd": list(page.get("outlinks") or []),
                "outlinks_int": len(list(page.get("outlinks") or [])),
                "related_kb_pages_kwd": list(page.get("related_kb_pages") or []),
                "source_chunk_ids": list(page.get("source_chunk_ids") or []),
                "source_doc_ids": list(page.get("source_doc_ids") or []),
                "content_with_weight": content_md,
                # Summary kept verbatim alongside the rendered body so the
                # viewer can render it as a distinct (smaller) block above
                # the main content.
                "summary_with_weight": summary,
                "content_ltks": content_ltks,
                "content_sm_ltks": content_sm_ltks,
                f"q_{len(vec_list)}_vec": vec_list,
                "available_int": 1,
            }
        )
        if topic:
            topics_by_name.setdefault(topic, _wiki_topic_slug(topic))
            bucket = topic_doc_ids.setdefault(topic, [])
            seen_bucket = set(bucket)
            for d in page.get("source_doc_ids") or []:
                if isinstance(d, str) and d and d not in seen_bucket:
                    seen_bucket.add(d)
                    bucket.append(d)

    if not rows:
        return

    try:
        await thread_pool_exec(settings.docStoreConn.insert, rows, index, ctx.kb_id)
    except Exception:
        logging.exception(
            "wiki_persist: bulk insert failed for kb=%s (rows=%d)",
            kb_id_str,
            len(rows),
        )
        return

    if topics_by_name:
        try:
            await _ensure_wiki_topic_rows(ctx, index, kb_id_str, topics_by_name, topic_doc_ids)
        except Exception:
            logging.exception(
                "wiki_persist: topic row insert failed for kb=%s",
                kb_id_str,
            )

    # Audit trail: one ArtifactCommit row per page whose rendered
    # content actually changed (record_edit silently skips empty diffs).
    # Best-effort — commit failures log but don't fail the artifact
    # compile.
    for page in pages:
        slug = (page.get("slug") or "").strip()
        if not slug:
            continue
        if not prior_by_slug.get(slug, ""):
            continue
        content_md = page.get("content_md_rendered") or page.get("content_md") or page.get("content_md_raw") or ""
        action = (page.get("action") or "CREATE").upper()
        try:
            WikiCommitService.record_page_edit(
                tenant_id=ctx.tenant_id,
                kb_id=ctx.kb_id,
                page_type=page.get("page_type") or "concept",
                slug=slug,
                content_before=prior_by_slug.get(slug, ""),
                content_after=content_md,
                title=WIKI_REGEN_COMMIT_TITLE,
                comments=WIKI_REGEN_COMMIT_COMMENTS_TEMPLATE.format(action=action),
                user_id=None,  # system commit — no human author
            )
        except Exception:
            logging.exception(
                "wiki_persist: commit record failed for kb=%s slug=%s",
                kb_id_str,
                slug,
            )


def build_wiki_page_graph(
    pages: List[Dict],
    kb_id: str,
) -> tuple[List[Dict], List[Dict]]:
    """Project the REFINE-emitted page list onto per-entity and
    per-relation ES rows.

    Returns:
        (entity_rows, relation_rows)

    Both lists are ES-ready docs (one per node / per surviving edge)
    using the standard artifact envelope. They are BM25-only (no
    ``q_<dim>_vec``) — entities carry ``content_ltks`` derived from
    ``slug + " " + summary`` so name/summary lookups hit the lexical
    index.

    ``by_slug`` is kept internally only to filter dangling outlinks
    (a target slug not present as a node in this KB).
    """
    from rag.nlp import rag_tokenizer

    by_slug: Dict[str, Dict] = {}
    entity_rows: List[Dict] = []
    for p in pages or []:
        slug = (p.get("slug") or "").strip()
        if not slug:
            continue
        outlinks_raw = p.get("outlinks") or []
        # ``weight`` is the page's outlink count. Drives node size on
        # the canvas. Computed on the raw outlink list before dangling-
        # target filtering so visual weight reflects what the writer
        # actually emitted.
        weight = len(outlinks_raw) if isinstance(outlinks_raw, list) else 0

        # Per-node provenance: union of source chunks REFINE attributed
        # to this page. Dedup preserves first-seen order; cap the list
        # to keep the graph blob small.
        raw_chunk_ids = p.get("source_chunk_ids") or []
        seen_chunk_ids: dict[str, None] = {}
        for cid in raw_chunk_ids:
            if isinstance(cid, str) and cid and cid not in seen_chunk_ids:
                seen_chunk_ids[cid] = None
                if len(seen_chunk_ids) >= WIKI_GRAPH_MAX_CHUNK_IDS_PER_NODE:
                    break
        capped_chunk_ids = list(seen_chunk_ids.keys())

        page_type = p.get("page_type") or "concept"
        description = p.get("summary") or ""
        name = p.get("title") or slug
        aliases = list(p.get("entity_names") or [])

        # Per-node doc provenance: the documents that fed this page. Stamped
        # onto the entity row so document deletion can reference-count it
        # (drop the entity only when its last source doc is removed).
        page_doc_ids = [d for d in (p.get("source_doc_ids") or []) if isinstance(d, str) and d]

        by_slug[slug] = {
            "slug": slug,
            "name": name,
            "aliases": aliases,
            "description": description,
            "type": page_type,
            "weight": weight,
            "source_chunk_ids": capped_chunk_ids,
            "source_doc_ids": page_doc_ids,
        }

        # Per-entity ES row. content_ltks is built from slug + summary
        # so BM25 hits both the deep-link key and human prose.
        content_text = (slug + " " + description).strip()
        entity_payload = {
            "slug": slug,
            "name": name,
            "aliases": aliases,
            "description": description,
            "type": page_type,
            "weight": weight,
        }
        entity_rows.append(
            {
                "id": xxhash.xxh64(
                    f"wiki_entity:{kb_id}:{slug}".encode("utf-8", "surrogatepass"),
                ).hexdigest(),
                "kb_id": kb_id,
                "doc_id": kb_id,  # KB-scoped sentinel
                "available_int": 1,
                "compile_kwd": "wiki_entity",
                "type_kwd": "wiki_" + page_type,
                "slug_kwd": slug,
                "weight_int": int(weight),
                "source_chunk_ids": capped_chunk_ids,
                "source_doc_ids": page_doc_ids,
                "content_ltks": rag_tokenizer.tokenize(content_text) if content_text else "",
                "content_with_weight": json.dumps(entity_payload, ensure_ascii=False),
            }
        )

    relation_rows: List[Dict] = []
    for p in pages or []:
        src = (p.get("slug") or "").strip()
        if not src or src not in by_slug:
            continue
        for raw_target in p.get("outlinks") or []:
            if isinstance(raw_target, str):
                tgt = raw_target.strip()
            elif isinstance(raw_target, dict):
                tgt = str(raw_target.get("slug") or "").strip()
            else:
                tgt = ""
            if not tgt or tgt == src or tgt not in by_slug:
                continue
            # Edge provenance is the union of both endpoints' source docs, so
            # the relation is dropped only once neither endpoint traces to a
            # surviving document.
            edge_doc_ids: list[str] = []
            edge_seen: set[str] = set()
            for endpoint in (src, tgt):
                for d in by_slug[endpoint].get("source_doc_ids") or []:
                    if d not in edge_seen:
                        edge_seen.add(d)
                        edge_doc_ids.append(d)
            relation_payload = {"from": src, "to": tgt}
            relation_rows.append(
                {
                    "id": xxhash.xxh64(
                        f"wiki_relation:{kb_id}:{src}:{tgt}".encode("utf-8", "surrogatepass"),
                    ).hexdigest(),
                    "kb_id": kb_id,
                    "doc_id": kb_id,
                    "available_int": 1,
                    "compile_kwd": "wiki_relation",
                    "type_kwd": "wiki_relation",
                    "from_kwd": src,
                    "to_kwd": tgt,
                    "source_doc_ids": edge_doc_ids,
                    "content_with_weight": json.dumps(relation_payload, ensure_ascii=False),
                }
            )

    return entity_rows, relation_rows


async def persist_wiki_page_graph_to_es(
    ctx: TaskContext,
    pages: List[Dict],
) -> None:
    """Materialize and store the per-entity / per-relation ES rows
    derived from artifact pages.

    Writes two row types — both delete-then-insert for idempotent
    re-runs:

    1. ``compile_kwd="wiki_entity"`` — one row per page node,
       BM25-only via ``content_ltks``.
    2. ``compile_kwd="wiki_relation"`` — one row per surviving
       edge (dangling outlinks dropped by the builder).

    Also sweeps any leftover legacy ``wiki_page_graph`` blob so
    the index doesn't accumulate stale state.
    """
    kb_id_str = str(ctx.kb_id)
    entity_rows, relation_rows = build_wiki_page_graph(pages or [], kb_id_str)

    index = search.index_name(ctx.tenant_id)

    async def _replace_bucket(kwd: str, rows: List[Dict]) -> None:
        try:
            await thread_pool_exec(
                settings.docStoreConn.delete,
                {"compile_kwd": kwd},
                index,
                ctx.kb_id,
            )
        except Exception:
            logging.debug(
                "%s: prior delete failed; relying on id-upsert",
                kwd,
            )
        if not rows:
            return
        try:
            await thread_pool_exec(
                settings.docStoreConn.insert,
                rows,
                index,
                ctx.kb_id,
            )
        except Exception:
            logging.exception(
                "%s: insert failed for kb=%s (%d rows)",
                kwd,
                kb_id_str,
                len(rows),
            )

    async def _sweep_legacy_blob() -> None:
        try:
            await thread_pool_exec(
                settings.docStoreConn.delete,
                {"compile_kwd": "wiki_page_graph"},
                index,
                ctx.kb_id,
            )
        except Exception:
            logging.debug(
                "wiki_page_graph: legacy blob sweep failed for kb=%s",
                kb_id_str,
            )

    await asyncio.gather(
        _replace_bucket("wiki_entity", entity_rows),
        _replace_bucket("wiki_relation", relation_rows),
        _sweep_legacy_blob(),
    )


# ----- main entry ----------------------------------------------------


async def run_wiki(
    ctx: TaskContext,
    embedding_model,
    load_chunks_for_doc: Callable[..., AsyncIterator[list[dict]]],
) -> None:
    """KB-wide artifact compilation task.

    Runs after the user clicks the "Artifact" button in the dataset
    generate menu. Iterates every doc in the KB whose parser_config
    has a compilation template group resolving to an artifacts-kind
    child, runs MAP per-doc (which uses ES-stored resume rows to skip
    chunks already processed in a previous run), then runs REDUCE /
    PLAN / REFINE KB-wide and persists pages.

    Batching: each MAP call uses ``batch_size_cap=8`` and
    ``window_fraction=0.5`` — i.e. roll over to a new batch when the
    current batch reaches 8 chunks OR its accumulated token count
    exceeds 50% of the chat model's ``max_length``.
    """
    # Local imports so this module doesn't drag in the API service
    # layer at import time — that's a source of circular-import risk
    # given how much lives under ``api.db.services``.
    from api.db.services.document_service import DocumentService
    from api.db.services.knowledgebase_service import KnowledgebaseService
    from api.db.services.compilation_template_service import CompilationTemplateService
    from api.db.services.llm_service import LLMBundle
    from api.db.joint_services.tenant_model_service import (
        get_tenant_default_model_by_type,
        resolve_model_config,
    )
    from api.apps.restful_apis.chunk_api import _compilation_template_kind

    progress = ctx.progress_cb
    progress(0.0, "Loading documents for wiki compilation...")

    # 1. Resolve KB metadata for PLAN.
    ok, kb = KnowledgebaseService.get_by_id(ctx.kb_id)
    if not ok:
        progress(-1, f"KB {ctx.kb_id} not found.")
        return
    kb_name = kb.name
    kb_description = kb.description

    # 2. Pick docs eligible for artifact compilation (those whose
    # configured template group resolves to at least one artifacts-kind
    # child). The frontend Artifact button targets the KB, but the
    # per-doc opt-in is what gates inclusion.
    all_docs, _ = await thread_pool_exec(
        DocumentService.get_by_kb_id,
        kb_id=ctx.kb_id,
        page_number=0,
        items_per_page=0,
        orderby="create_time",
        desc=False,
        keywords="",
        run_status=[],
        types=[],
        suffix=[],
    )
    current_doc_ids = {str(d.get("id")) for d in all_docs or [] if d.get("id")}
    existing_map_doc_ids = await _wiki_existing_map_doc_ids(ctx.tenant_id, ctx.kb_id)
    deleted_doc_ids = existing_map_doc_ids - current_doc_ids
    if deleted_doc_ids:
        progress(
            0.02,
            f"Removing stale wiki state for {len(deleted_doc_ids)} deleted document(s)...",
        )
        await _wiki_delete_deleted_doc_state(
            ctx.tenant_id,
            ctx.kb_id,
            deleted_doc_ids,
        )

    eligible = []
    pipeline_template_ids_cache: dict[str, list[str]] = {}
    for d in all_docs or []:
        pc = d.get("parser_config") or {}
        template_ids: list[str] = []
        seen_template_ids: set[str] = set()
        for template_id in _parser_config_compilation_template_ids(pc, ctx.tenant_id):
            if template_id in seen_template_ids:
                continue
            seen_template_ids.add(template_id)
            template_ids.append(template_id)
        pipeline_id = (d.get("pipeline_id") or "").strip()
        if pipeline_id:
            if pipeline_id not in pipeline_template_ids_cache:
                pipeline_template_ids_cache[pipeline_id] = _pipeline_compilation_template_ids(pipeline_id, ctx.tenant_id)
            for template_id in pipeline_template_ids_cache[pipeline_id]:
                if template_id in seen_template_ids:
                    continue
                seen_template_ids.add(template_id)
                template_ids.append(template_id)

        for template_id in template_ids:
            template = CompilationTemplateService.get_saved(template_id, ctx.tenant_id)
            config = template.get("config") if template else {}
            kind = _compilation_template_kind(config.get("kind") if isinstance(config, dict) else "")
            if kind == "wiki":
                eligible.append((d, template_id))
                break
    if not eligible:
        progress(1.0, "No documents are configured for wiki compilation.")
        return

    # 3. Resolve chat models. MAP is per-(doc, template) so each pair
    # uses its template's own ``llm_id``. REDUCE / PLAN / REFINE are
    # KB-wide and need exactly one model — we pick the first eligible
    # template's ``llm_id`` as the canonical KB chat model.
    llm_bundle_cache: dict[str, LLMBundle] = {}

    def _bundle_for(llm_id: str | None) -> LLMBundle:
        key = (llm_id or "").strip() or "__tenant_default__"
        cached = llm_bundle_cache.get(key)
        if cached is not None:
            return cached
        try:
            if key == "__tenant_default__":
                cfg = get_tenant_default_model_by_type(ctx.tenant_id, LLMType.CHAT)
            else:
                cfg = resolve_model_config(
                    ctx.tenant_id,
                    LLMType.CHAT,
                    key,
                )
        except Exception:
            logging.exception(
                "wiki: chat model resolution failed for llm_id=%s (kb=%s); falling back to tenant default",
                key,
                ctx.kb_id,
            )
            cfg = get_tenant_default_model_by_type(ctx.tenant_id, LLMType.CHAT)
            key = "__tenant_default__"
            cached = llm_bundle_cache.get(key)
            if cached is not None:
                return cached
        bundle = LLMBundle(ctx.tenant_id, cfg, lang=ctx.language)
        llm_bundle_cache[key] = bundle
        return bundle

    def _stage_cb(prefix: str):
        def _cb(*args, **kwargs):
            try:
                if args and isinstance(args[0], (int, float)):
                    msg = args[1] if len(args) > 1 else kwargs.get("msg", "")
                    progress(msg=f"{prefix} {msg}")
                else:
                    msg = kwargs.get("msg") or (args[0] if args else "")
                    progress(msg=f"{prefix} {msg}")
            except Exception:
                logging.exception("wiki: progress callback failed")

        return _cb

    # 4. MAP per eligible doc. Each MAP call's own resume mechanism
    # (wiki_map_extract rows keyed by chunk_id) skips chunks that
    # were already processed in a prior run — this is the incremental
    # behavior the user asked for.
    #
    # Resolve templates before starting workers so the first eligible
    # template remains the deterministic source for KB-wide REDUCE/PLAN/
    # REFINE, independent of which document worker happens to finish first.
    resolved_eligible: list[tuple[dict, str, dict]] = []
    for doc, template_id in eligible:
        template = CompilationTemplateService.get_saved(template_id, ctx.tenant_id)
        if not template:
            logging.warning(
                "artifact: template %s not found for doc %s; skipping",
                template_id,
                doc["id"],
            )
            continue
        resolved_eligible.append((doc, template_id, template.get("config") or {}))

    if not resolved_eligible:
        progress(1.0, "No valid templates resolved for wiki compilation.")
        return

    # ``kb_chat_llm_id`` is captured from the first eligible template and
    # used as the canonical chat model for KB-wide REDUCE/PLAN/REFINE.
    # ``kb_writer_example`` follows the same first-template-wins rule.
    first_parser_cfg = resolved_eligible[0][2]
    first_parser_cfg = first_parser_cfg if isinstance(first_parser_cfg, dict) else {}
    kb_chat_llm_id: Optional[str] = (first_parser_cfg.get("llm_id") or "").strip() or None
    first_example = first_parser_cfg.get("example")
    kb_writer_example: Optional[str] = first_example if isinstance(first_example, str) and first_example.strip() else None
    map_llm_pool = LLMCallPool(WIKI_MAP_LLM_POOL_SIZE, max_pending=WIKI_MAP_MAX_PENDING)
    n_docs = len(resolved_eligible)
    map_queue: asyncio.Queue = asyncio.Queue(maxsize=WIKI_MAP_QUEUE_SIZE)
    doc_stats = {
        i: {
            "doc": doc,
            "batch_count": 0,
            "saw_any": False,
            "status": "ok",
            "agg": {"entities": 0, "concepts": 0, "claims": 0, "relations": 0},
            "delta": {"new": 0, "changed": 0, "unchanged": 0, "deleted": 0},
            "had_delta": False,
        }
        for i, (doc, _, _) in enumerate(resolved_eligible)
    }

    async def _produce_document(i: int, job: tuple[dict, str, dict]) -> None:
        doc, template_id, parser_cfg = job
        doc_id = doc["id"]
        stats = doc_stats[i]
        progress(0.05 + 0.6 * (i / n_docs), f"MAP {i + 1}/{n_docs}: {doc.get('name', doc_id)}")
        try:
            async for batch in load_chunks_for_doc(
                ctx.tenant_id,
                ctx.kb_id,
                doc_id,
                batch_size=WIKI_MAP_BATCH_CHUNKS,
            ):
                stats["saw_any"] = True
                stats["batch_count"] += 1
                await map_queue.put((i, doc, template_id, parser_cfg, batch, stats["batch_count"]))
        except Exception:
            logging.exception("wiki: loading MAP chunks failed for doc %s", doc_id)
            stats["status"] = "error"

    async def _map_worker() -> None:
        while True:
            item = await map_queue.get()
            try:
                if item is None:
                    return
                i, doc, template_id, parser_cfg, batch, batch_no = item
                doc_id = doc["id"]
                stats = doc_stats[i]
                map_llm_id = (parser_cfg.get("llm_id") or "").strip() if isinstance(parser_cfg, dict) else ""
                phase1 = await wiki_map_from_chunks(
                    chunks=batch,
                    chat_mdl=map_llm_pool.wrap(
                        _bundle_for(map_llm_id),
                        priority=30,
                        label=f"wiki-map:{doc_id}",
                        context=f"{ctx.kb_id}:{doc_id}:map",
                    ),
                    embd_mdl=embedding_model,
                    doc_id=doc_id,
                    tenant_id=ctx.tenant_id,
                    kb_id=ctx.kb_id,
                    language=ctx.language,
                    callback=_stage_cb(f"[wiki MAP {i + 1}/{n_docs} b{batch_no}]"),
                    parser_config=parser_cfg,
                    batch_size_cap=8,
                    window_fraction=0.5,
                    # Keep a bounded internal worker queue. The shared pool
                    # globally limits active + admitted waiting calls to
                    # WIKI_MAP_MAX_PENDING, while this prevents every outer
                    # batch from creating all of its sub-batch tasks at once.
                    max_workers=6,
                )
                for key in stats["agg"]:
                    stats["agg"][key] += len(phase1.get(key) or [])
                meta = phase1.get("_meta") or {}
                if isinstance(meta, dict):
                    for key in stats["delta"]:
                        stats["delta"][key] += int(meta.get(key, 0) or 0)
                    stats["had_delta"] |= bool(meta.get("had_delta"))
            except Exception:
                logging.exception(
                    "wiki: MAP failed for doc %s batch %d",
                    doc_id,
                    batch_no,
                )
                stats["status"] = "error"
            finally:
                map_queue.task_done()

    producers = [asyncio.create_task(_produce_document(i, job)) for i, job in enumerate(resolved_eligible)]
    workers = [asyncio.create_task(_map_worker()) for _ in range(WIKI_MAP_LLM_POOL_SIZE)]
    try:
        await asyncio.gather(*producers)
        await map_queue.join()
    finally:
        for task in producers + workers:
            if not task.done():
                task.cancel()
        await asyncio.gather(*producers, *workers, return_exceptions=True)

    for i, stats in doc_stats.items():
        doc = stats["doc"]
        doc_id = doc["id"]
        if not stats["saw_any"] and stats["status"] == "ok":
            stats["status"] = "empty"
            logging.info("wiki: no chunks for doc %s; skipping", doc_id)
        agg = stats["agg"]
        delta = stats["delta"]
        logging.info(
            "wiki: MAP doc=%s entities=%d concepts=%d claims=%d relations=%d (batches=%d, new=%d changed=%d unchanged=%d deleted=%d, delta=%s)",
            doc_id,
            agg["entities"],
            agg["concepts"],
            agg["claims"],
            agg["relations"],
            stats["batch_count"],
            delta["new"],
            delta["changed"],
            delta["unchanged"],
            delta["deleted"],
            stats["had_delta"],
        )

    # 5. REDUCE / PLAN / REFINE KB-wide. Each phase has its own
    # input_hash gate (REDUCE keys off the MAP-state hash, PLAN off
    # REDUCE's hash, REFINE off PLAN's hash) so re-runs without an
    # upstream delta short-circuit at the cache layer.
    kb_chat_mdl = _bundle_for(kb_chat_llm_id)

    try:
        progress(0.65, "Reducing extracts KB-wide...")
        await wiki_reduce_from_extracts(
            chat_mdl=map_llm_pool.wrap(
                kb_chat_mdl,
                priority=20,
                label="wiki-reduce",
                context=f"{ctx.kb_id}:reduce",
            ),
            embd_mdl=embedding_model,
            tenant_id=ctx.tenant_id,
            kb_id=ctx.kb_id,
            callback=_stage_cb("[wiki REDUCE]"),
        )

        progress(0.75, "Planning wiki pages...")
        await wiki_plan_from_reduction(
            chat_mdl=kb_chat_mdl,
            embd_mdl=embedding_model,
            tenant_id=ctx.tenant_id,
            kb_id=ctx.kb_id,
            kb_name=kb_name,
            kb_description=kb_description,
            callback=_stage_cb("[wiki PLAN]"),
        )

        progress(0.85, "Refining pages...")
        pages = await wiki_refine_from_plan(
            chat_mdl=map_llm_pool.wrap(
                kb_chat_mdl,
                priority=30,
                label="wiki-refine",
                context=f"{ctx.kb_id}:refine",
            ),
            embd_mdl=embedding_model,
            tenant_id=ctx.tenant_id,
            kb_id=ctx.kb_id,
            max_workers=WIKI_REFINE_WORKERS,
            callback=_stage_cb("[wiki REFINE]"),
            example=kb_writer_example,
        )
    except Exception:
        logging.exception("wiki: REDUCE/PLAN/REFINE failed for kb %s", ctx.kb_id)
        progress(-1, "Wiki pipeline failed during REDUCE/PLAN/REFINE.")
        return

    # 6. Persist searchable wiki_page rows.
    try:
        await persist_wiki_pages_to_es(ctx=ctx, pages=pages or [], embd_mdl=embedding_model)
    except Exception:
        logging.exception("wiki: ES persist failed for kb %s", ctx.kb_id)

    # 7. Materialize the canvas graph from the refined pages.
    try:
        await persist_wiki_page_graph_to_es(ctx=ctx, pages=pages or [])
    except Exception:
        logging.exception("wiki: page-graph persist failed for kb %s", ctx.kb_id)

    progress(1.0, f"Wiki compiled {len(pages or [])} page(s).")
