"""Summarize + tag job — module 5 phase 3, plan step 3 (kb/phases/kb-m5-faza3.md,
§5 step 3, §2 decision 3: two-track A/B pilot).

Pipeline: `document_chunk.text WHERE excluded_reason IS NULL ORDER BY chunk_index` per
`source='paperless'` envelope (duplicates via `entities[type=duplicate_of]` skipped whole)
-> LLM (`--backend ollama|anthropic`) forced-JSON `{"summary": ..., "tags": [...]}`
-> `INSERT document_summary` (`services/kb-postgres/init/004_summaries.sql`). `envelope` and
`document_chunk` are read-only here; this job only ever `INSERT`s into `document_summary`.

Two backends write the same table under different `model` values — `UNIQUE (envelope_id,
model)` exists precisely so both A/B tracks coexist (plan §2 decision 3). A second, separate
mode (`--embed-summaries`) embeds existing summaries with bge-m3 via Ollama, reusing
`kb_retrieval.embed.embed_chunk` 1:1 — writing and embedding are split so either can be re-run
without redoing the other (SOLARIA asleep blocks embedding, not writing; API downtime blocks
writing, not embedding).

Long documents (plan §1.2 outlier: ~360k chars) exceed practical LLM context: chunks are
grouped ~20-at-a-time into partial summaries (plain text, no JSON), then a final pass
synthesizes those partials with the normal JSON prompt. Counted separately
(`documents_mapreduce`), never silently skipped.

Runs on PIHA (rsync `src/` + `tags-vocab.yaml` to `/tmp`, per plan §1.4) against
kb-postgres@PIHA and OLLAMA_URL=http://solaria:11434, or on SOLARIA directly for the local
backend. Anthropic backend needs `ANTHROPIC_API_KEY` in the environment for the run only —
never logged, never written to disk.

Usage:
    # Dry run (default) — fetch, chunk-concat, count; no LLM calls, no DB writes:
    documents-ingest-summarize --dsn ... --backend ollama --model gemma3:12b

    # Real run, local track:
    documents-ingest-summarize --dsn ... --backend ollama --model gemma3:12b --apply

    # Real run, API track (ANTHROPIC_API_KEY must be set):
    documents-ingest-summarize --dsn ... --backend anthropic --model claude-haiku-4-5 --apply

    # Embed existing summaries (bge-m3, second pass):
    documents-ingest-summarize --dsn ... --embed-summaries --embed-model bge-m3 --apply
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import sys
import unicodedata
from pathlib import Path
from typing import Optional, Protocol

import aiohttp
import asyncpg
import structlog
import yaml
from anthropic import AsyncAnthropic

from documents_ingest.chunk_embed import EmbeddingDimensionError, EXPECTED_DIM
from kb_retrieval.embed import _vector_literal, embed_chunk

_log = structlog.get_logger(__name__)

DEFAULT_OLLAMA_URL = "http://localhost:11434"
DEFAULT_EMBED_MODEL = "bge-m3"
DEFAULT_MAX_TOKENS = 2048
DEFAULT_CHUNKS_PER_GROUP = 20

# Ollama's *runtime* context window defaults to a few thousand tokens regardless of a
# model's advertised max (gemma3:12b advertises 131072) -- confirmed empirically during the
# pilot: a 93k-char OWU document was silently truncated to prompt_eval_count=2051 tokens with
# no num_ctx set, and the resulting summary described the document's tail (a RODO clause)
# instead of its actual insurance terms. `options.num_ctx` must be set explicitly on every
# call, sized to the prompt -- never left to Ollama's default.
OLLAMA_CHARS_PER_TOKEN = 3.0  # conservative for diacritic-heavy Polish text
OLLAMA_MIN_NUM_CTX = 4096
OLLAMA_MAX_NUM_CTX = 131072  # gemma3:12b's advertised context_length
OLLAMA_CTX_RESPONSE_RESERVE = 4096  # headroom for template + generated output


def compute_num_ctx(prompt_chars: int) -> int:
    """Size Ollama's num_ctx to the actual prompt, rounded up to a 1024-token step (stable
    KV-cache allocation), clamped to [OLLAMA_MIN_NUM_CTX, OLLAMA_MAX_NUM_CTX]."""
    needed = int(prompt_chars / OLLAMA_CHARS_PER_TOKEN) + OLLAMA_CTX_RESPONSE_RESERVE
    needed = ((needed + 1023) // 1024) * 1024
    return max(OLLAMA_MIN_NUM_CTX, min(OLLAMA_MAX_NUM_CTX, needed))
# Plan §1.2: the pilot corpus has one ~360k-char outlier; everything else is far smaller.
# This threshold only exists to catch that shape of document, not to tune routine chunking.
DEFAULT_MAPREDUCE_THRESHOLD_CHARS = 200_000
DEFAULT_TAGS_VOCAB_PATH = Path(__file__).resolve().parent.parent.parent / "tags-vocab.yaml"

SYSTEM_PROMPT = (
    "Jesteś archiwistą domowej bazy wiedzy. Streszczasz dokumenty po polsku, zwięźle i "
    "faktograficznie: kto, co, kiedy, kwoty, numery umów/polis, terminy. Nie zgadujesz — "
    "czego nie ma w tekście, tego nie piszesz. Zwracasz wyłącznie JSON."
)

PARTIAL_SYSTEM_PROMPT = (
    "Jesteś archiwistą domowej bazy wiedzy. Streszczasz fragment dłuższego dokumentu po "
    "polsku w 2-4 zdaniach, wypisując wszystkie fakty: kwoty, daty, numery, strony. To "
    "streszczenie częściowe zostanie później połączone z innymi częściami tego samego "
    "dokumentu — nie pomijaj konkretów. Odpowiadasz czystym tekstem, bez JSON."
)

SUMMARY_JSON_SCHEMA = {
    "type": "object",
    "properties": {
        "summary": {"type": "string"},
        "tags": {"type": "array", "items": {"type": "string"}},
    },
    "required": ["summary", "tags"],
    "additionalProperties": False,
}

_INSERT_SQL = """
INSERT INTO document_summary (envelope_id, summary, tags, model)
VALUES ($1, $2, $3::jsonb, $4)
ON CONFLICT (envelope_id, model) DO NOTHING
"""

_TAG_CLEAN_RE = re.compile(r"[^\w-]+", re.UNICODE)
_TAG_DASH_COLLAPSE_RE = re.compile(r"-{2,}")


class LLMBackend(Protocol):
    """A backend answers one prompt at a time. `json_mode=True` requests forced-JSON output
    (final summary+tags call); `json_mode=False` is a plain-text call (map-reduce partials)."""

    async def complete(self, system: str, user: str, json_mode: bool = True) -> str: ...


class OllamaBackend:
    def __init__(self, session: aiohttp.ClientSession, base_url: str, model: str):
        self.session = session
        self.base_url = base_url
        self.model = model

    async def complete(self, system: str, user: str, json_mode: bool = True) -> str:
        num_ctx = compute_num_ctx(len(system) + len(user))
        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            "stream": False,
            "options": {"temperature": 0, "num_ctx": num_ctx},
        }
        if json_mode:
            payload["format"] = "json"
        async with self.session.post(f"{self.base_url}/api/chat", json=payload) as resp:
            resp.raise_for_status()
            data = await resp.json()
        return data["message"]["content"]


class AnthropicBackend:
    def __init__(self, client, model: str, max_tokens: int = DEFAULT_MAX_TOKENS):
        self.client = client
        self.model = model
        self.max_tokens = max_tokens

    async def complete(self, system: str, user: str, json_mode: bool = True) -> str:
        kwargs = {}
        if json_mode:
            kwargs["output_config"] = {
                "format": {"type": "json_schema", "schema": SUMMARY_JSON_SCHEMA}
            }
        response = await self.client.messages.create(
            model=self.model,
            max_tokens=self.max_tokens,
            temperature=0,
            system=system,
            messages=[{"role": "user", "content": user}],
            **kwargs,
        )
        return next((b.text for b in response.content if b.type == "text"), "")


def load_tags_vocab(path: Path) -> list[str]:
    """Controlled tag vocabulary (plan §2 decision 4) — a versioned YAML list."""
    with open(path, encoding="utf-8") as f:
        data = yaml.safe_load(f)
    return list(data.get("tags", []))


def build_user_prompt(vocab: list[str], content: str) -> str:
    return (
        "Słownik kontrolowany tagów (WYBIERZ Z TEJ LISTY w pierwszej kolejności, maks. 3 "
        "najtrafniejsze pozycje, dokładnie w tym zapisie):\n"
        f"{', '.join(vocab)}\n\n"
        "Dopiero jeśli żaden z powyższych tagów nie pasuje do dokumentu, możesz dodać maks. "
        "3 dodatkowe tagi spoza listy — one też MUSZĄ być PO POLSKU, małymi literami, "
        "kebab-case (np. \"wspolnota-mieszkaniowa\", nie \"community\" ani \"HOA\"). "
        "Nigdy nie zwracaj tagów po angielsku.\n\n"
        "Treść dokumentu:\n"
        f"{content}\n\n"
        'Zwróć JSON: {"summary": "<3-8 zdań PO POLSKU>", "tags": ["tag1", "tag2", ...]}'
    )


def build_partial_prompt(group_text: str, idx: int, total: int) -> str:
    return (
        f"Fragment dokumentu, część {idx}/{total}:\n\n{group_text}\n\n"
        "Streść ten fragment w 2-4 zdaniach po polsku, z konkretami (kwoty, daty, numery)."
    )


def is_duplicate(entities: list) -> bool:
    """`entities[type=duplicate_of]` (plan §3.2 / §5.1) — duplicate envelopes are skipped
    whole, never summarized even partially."""
    return any(
        isinstance(e, dict) and e.get("type") == "duplicate_of" for e in (entities or [])
    )


def normalize_tag(tag: str) -> str:
    """lowercase, kebab-case, NFC (plan §2 decision 4)."""
    t = unicodedata.normalize("NFC", tag.strip().lower())
    t = t.replace(" ", "-")
    t = _TAG_CLEAN_RE.sub("-", t)
    t = _TAG_DASH_COLLAPSE_RE.sub("-", t)
    return t.strip("-")


def normalize_and_validate_tags(raw_tags: list, vocab: list[str]) -> tuple[list[str], int]:
    """Normalizes + dedupes tags; vocab tags are kept in full, free-form tags capped at 3
    (plan §2 decision 4). Returns (final_tags, truncated_count) — truncated_count is
    informational only, not part of the stats balance invariant."""
    vocab_set = set(vocab)
    seen: list[str] = []
    for raw in raw_tags or []:
        if not isinstance(raw, str):
            continue
        norm = normalize_tag(raw)
        if norm and norm not in seen:
            seen.append(norm)

    vocab_tags = [t for t in seen if t in vocab_set]
    freeform = [t for t in seen if t not in vocab_set]
    truncated = max(0, len(freeform) - 3)
    return vocab_tags + freeform[:3], truncated


async def fetch_documents(conn: asyncpg.Connection, limit: Optional[int], offset: Optional[int]) -> list:
    """`source='paperless'` envelopes, ordered by id for stable --limit/--offset slicing."""
    query = "SELECT id, entities FROM envelope WHERE source = 'paperless' ORDER BY id"
    params: list = []
    if limit is not None:
        params.append(limit)
        query += f" LIMIT ${len(params)}"
    if offset:
        params.append(offset)
        query += f" OFFSET ${len(params)}"
    return await conn.fetch(query, *params)


async def fetch_active_chunk_texts(conn: asyncpg.Connection, envelope_id: str) -> list[str]:
    """`document_chunk.text WHERE excluded_reason IS NULL` (plan §5.1) — junk and duplicate
    chunks never enter a summary."""
    rows = await conn.fetch(
        "SELECT text FROM document_chunk WHERE envelope_id = $1 AND excluded_reason IS NULL "
        "ORDER BY chunk_index",
        envelope_id,
    )
    return [r["text"] for r in rows]


async def fetch_existing_summary_keys(conn: asyncpg.Connection, model: str) -> set[str]:
    """envelope_ids already summarized with this `model` — idempotency + dry-run preview."""
    rows = await conn.fetch("SELECT envelope_id FROM document_summary WHERE model = $1", model)
    return {r["envelope_id"] for r in rows}


def _decode_jsonb(value: object) -> object:
    if value is None:
        return None
    if isinstance(value, str):
        return json.loads(value)
    return value


async def get_summary_and_tags(
    backend: LLMBackend, system: str, user: str, retries: int = 1
) -> Optional[dict]:
    """Forced-JSON call with one retry on invalid JSON/schema (plan §5.1). Returns None
    (caller counts `llm_errors`) after `retries` failed attempts."""
    for attempt in range(retries + 1):
        try:
            raw = await backend.complete(system, user, json_mode=True)
            data = json.loads(raw)
        except Exception:
            data = None

        if (
            isinstance(data, dict)
            and isinstance(data.get("summary"), str)
            and isinstance(data.get("tags"), list)
        ):
            return data

        if attempt < retries:
            user = (
                user
                + "\n\nUWAGA: Poprzednia odpowiedź nie była poprawnym JSON zgodnym ze "
                "schematem. Zwróć WYŁĄCZNIE poprawny JSON, bez dodatkowego tekstu."
            )
    return None


async def summarize_mapreduce(
    backend: LLMBackend, chunks: list[str], group_size: int
) -> Optional[str]:
    """Groups chunks ~`group_size`-at-a-time into partial summaries (plain text), returns
    their concatenation as the "content" fed into the final JSON summarization pass. Returns
    None if any partial call fails — caller counts the whole document as `llm_errors`, never
    silently drops a partial."""
    groups = [chunks[i : i + group_size] for i in range(0, len(chunks), group_size)]
    partials: list[str] = []
    for idx, group in enumerate(groups, start=1):
        group_text = "\n\n".join(group)
        try:
            partial = await backend.complete(
                PARTIAL_SYSTEM_PROMPT,
                build_partial_prompt(group_text, idx, len(groups)),
                json_mode=False,
            )
        except Exception:
            _log.warning("skip.mapreduce_partial_error", group_index=idx, group_total=len(groups))
            return None
        partial = partial.strip()
        if not partial:
            _log.warning("skip.mapreduce_partial_empty", group_index=idx, group_total=len(groups))
            return None
        partials.append(partial)
    return "\n\n".join(partials)


async def insert_summary(
    conn: asyncpg.Connection, envelope_id: str, summary: str, tags: list[str], model: str
) -> str:
    """Returns asyncpg's command tag so the caller can tell a real insert from an
    ON CONFLICT no-op (second line of defense behind the pre-fetched `existing` set)."""
    return await conn.execute(_INSERT_SQL, envelope_id, summary, json.dumps(tags), model)


def _rows_affected(command_tag: str) -> int:
    return int(command_tag.rsplit(" ", 1)[-1])


async def run_summarize(
    dsn: str,
    backend_name: str,
    model: str,
    ollama_url: str = DEFAULT_OLLAMA_URL,
    anthropic_api_key: Optional[str] = None,
    tags_vocab_path: Path = DEFAULT_TAGS_VOCAB_PATH,
    limit: Optional[int] = None,
    offset: Optional[int] = None,
    apply: bool = False,
    max_tokens: int = DEFAULT_MAX_TOKENS,
    chunks_per_group: int = DEFAULT_CHUNKS_PER_GROUP,
    mapreduce_threshold_chars: int = DEFAULT_MAPREDUCE_THRESHOLD_CHARS,
) -> dict:
    """Summarize+tag one --limit/--offset slice of `source='paperless'` envelopes into
    `document_summary` under `model`.

    Stats must balance:
        documents_fetched = duplicates_skipped + no_active_chunks + already_summarized
                             + summarized + llm_errors

    `duplicates_skipped` (envelope has `entities[type=duplicate_of]`) and `no_active_chunks`
    (zero `document_chunk` rows with `excluded_reason IS NULL`) are both counted before any
    LLM call — dry-run and apply agree on this split even though dry-run never calls the LLM.
    A failed DB insert (rare — the pre-fetched `existing` set is the first line of defense
    against duplicate work) is folded into `llm_errors`, same as `chunk_embed.py`'s
    `chunks_errors` covers both embed and insert failures.
    """
    stats = {
        "documents_fetched": 0,
        "duplicates_skipped": 0,
        "no_active_chunks": 0,
        "already_summarized": 0,
        "summarized": 0,
        "llm_errors": 0,
        "documents_mapreduce": 0,
        "tags_truncated": 0,
    }

    conn = await asyncpg.connect(dsn)
    try:
        docs = await fetch_documents(conn, limit, offset)
        existing = await fetch_existing_summary_keys(conn, model)
        vocab = load_tags_vocab(tags_vocab_path)

        session: Optional[aiohttp.ClientSession] = None
        client = None
        backend: Optional[LLMBackend] = None
        if apply:
            if backend_name == "ollama":
                session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=300))
                backend = OllamaBackend(session, ollama_url, model)
            else:
                client = AsyncAnthropic(api_key=anthropic_api_key)
                backend = AnthropicBackend(client, model, max_tokens)

        try:
            for row in docs:
                stats["documents_fetched"] += 1
                envelope_id = row["id"]
                entities = _decode_jsonb(row["entities"]) or []

                if is_duplicate(entities):
                    stats["duplicates_skipped"] += 1
                    continue

                if envelope_id in existing:
                    stats["already_summarized"] += 1
                    continue

                chunks = await fetch_active_chunk_texts(conn, envelope_id)
                content = "\n\n".join(chunks).strip()
                if not content:
                    stats["no_active_chunks"] += 1
                    continue

                if not apply:
                    stats["summarized"] += 1
                    continue

                assert backend is not None
                if len(content) > mapreduce_threshold_chars:
                    stats["documents_mapreduce"] += 1
                    content = await summarize_mapreduce(backend, chunks, chunks_per_group)
                    if content is None:
                        stats["llm_errors"] += 1
                        continue

                user_prompt = build_user_prompt(vocab, content)
                data = await get_summary_and_tags(backend, SYSTEM_PROMPT, user_prompt)
                if data is None:
                    stats["llm_errors"] += 1
                    _log.warning("skip.llm_error", envelope_id=envelope_id)
                    continue

                tags, truncated = normalize_and_validate_tags(data["tags"], vocab)
                stats["tags_truncated"] += truncated

                try:
                    command_tag = await insert_summary(conn, envelope_id, data["summary"], tags, model)
                except Exception:
                    _log.warning("skip.insert_error", envelope_id=envelope_id, exc_info=True)
                    stats["llm_errors"] += 1
                    continue

                existing.add(envelope_id)
                if _rows_affected(command_tag) == 0:
                    # ON CONFLICT no-op: the pre-fetched `existing` set should make this rare;
                    # treat it the same as already-summarized rather than adding a bucket the
                    # plan's balance formula doesn't have.
                    stats["already_summarized"] += 1
                else:
                    stats["summarized"] += 1
        finally:
            if session is not None:
                await session.close()
            if client is not None:
                await client.close()
    finally:
        await conn.close()

    balance = (
        stats["duplicates_skipped"]
        + stats["no_active_chunks"]
        + stats["already_summarized"]
        + stats["summarized"]
        + stats["llm_errors"]
    )
    if balance != stats["documents_fetched"]:
        _log.error("stats_mismatch", **stats)

    _log.info("run_complete", apply=apply, backend=backend_name, model=model, **stats)
    return stats


async def run_embed_summaries(
    dsn: str,
    ollama_url: str = DEFAULT_OLLAMA_URL,
    embed_model: str = DEFAULT_EMBED_MODEL,
    model_filter: Optional[str] = None,
    limit: Optional[int] = None,
    offset: Optional[int] = None,
    apply: bool = False,
) -> dict:
    """Second, separate pass: embeds `document_summary.summary` (bge-m3, via Ollama) for
    rows with `embedding IS NULL`. Reuses `chunk_embed.embed_chunk` — same Ollama
    `/api/embeddings` call, same dimension guard. `--model` (the summarizing model) filters
    which track to embed; omit to embed both tracks' pending rows in one run.

    Stats must balance: summaries_fetched = embedded + errors
    """
    stats = {"summaries_fetched": 0, "embedded": 0, "errors": 0}

    query = "SELECT id, summary FROM document_summary WHERE embedding IS NULL"
    params: list = []
    if model_filter:
        params.append(model_filter)
        query += f" AND model = ${len(params)}"
    query += " ORDER BY id"
    if limit is not None:
        params.append(limit)
        query += f" LIMIT ${len(params)}"
    if offset:
        params.append(offset)
        query += f" OFFSET ${len(params)}"

    conn = await asyncpg.connect(dsn)
    try:
        rows = await conn.fetch(query, *params)

        session: Optional[aiohttp.ClientSession] = None
        if apply:
            session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120))

        try:
            for row in rows:
                stats["summaries_fetched"] += 1
                if not apply:
                    stats["embedded"] += 1
                    continue

                assert session is not None
                try:
                    embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, row["summary"])
                except Exception:
                    _log.warning("skip.embed_error", summary_id=row["id"], exc_info=True)
                    stats["errors"] += 1
                    continue

                if len(embedding) != EXPECTED_DIM:
                    raise EmbeddingDimensionError(
                        f"ollama model={embed_model!r} returned dim={len(embedding)}, "
                        f"expected {EXPECTED_DIM}"
                    )

                try:
                    await conn.execute(
                        "UPDATE document_summary SET embedding = $1::vector, embedding_model = $2 "
                        "WHERE id = $3",
                        _vector_literal(embedding),
                        embed_model,
                        row["id"],
                    )
                except Exception:
                    _log.warning("skip.update_error", summary_id=row["id"], exc_info=True)
                    stats["errors"] += 1
                    continue

                stats["embedded"] += 1
        finally:
            if session is not None:
                await session.close()
    finally:
        await conn.close()

    if stats["embedded"] + stats["errors"] != stats["summaries_fetched"]:
        _log.error("stats_mismatch", **stats)

    _log.info("run_complete", apply=apply, embed_model=embed_model, **stats)
    return stats


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Summarize+tag source='paperless' documents into document_summary, or "
        "embed existing summaries with --embed-summaries (module 5, phase 3, plan §5)."
    )
    parser.add_argument("--dsn", default=os.environ.get("KB_DSN"),
                        help="asyncpg DSN for kb-postgres (or set KB_DSN env var)")
    parser.add_argument("--limit", type=int, default=None, help="Max documents/summaries to process")
    parser.add_argument("--offset", type=int, default=0, help="Slice offset, ordered by id")
    parser.add_argument("--apply", action="store_true",
                        help="Actually call the LLM/Ollama and write rows. Default is dry-run.")

    parser.add_argument("--embed-summaries", action="store_true",
                        help="Embed existing document_summary rows (bge-m3) instead of writing new ones.")
    parser.add_argument("--embed-model", default=os.environ.get("SUMMARY_EMBED_MODEL", DEFAULT_EMBED_MODEL),
                        help=f"Ollama embedding model for --embed-summaries (default: {DEFAULT_EMBED_MODEL})")

    parser.add_argument("--backend", choices=["ollama", "anthropic"],
                        help="LLM backend for writing summaries (required unless --embed-summaries)")
    parser.add_argument("--model", help="Backend model name, e.g. gemma3:12b or claude-haiku-4-5 "
                        "(writing mode); or a filter on document_summary.model (--embed-summaries)")
    parser.add_argument("--ollama-url", default=os.environ.get("OLLAMA_URL", DEFAULT_OLLAMA_URL),
                        help=f"Ollama base URL (default: {DEFAULT_OLLAMA_URL}, or set OLLAMA_URL)")
    parser.add_argument("--anthropic-api-key", default=os.environ.get("ANTHROPIC_API_KEY"),
                        help="Anthropic API key (or set ANTHROPIC_API_KEY env var) — never logged")
    parser.add_argument("--tags-vocab", type=Path, default=DEFAULT_TAGS_VOCAB_PATH,
                        help=f"Path to tags-vocab.yaml (default: {DEFAULT_TAGS_VOCAB_PATH})")
    parser.add_argument("--max-tokens", type=int, default=DEFAULT_MAX_TOKENS,
                        help=f"Anthropic max_tokens (default: {DEFAULT_MAX_TOKENS})")
    parser.add_argument("--chunks-per-group", type=int, default=DEFAULT_CHUNKS_PER_GROUP,
                        help=f"Map-reduce group size in chunks (default: {DEFAULT_CHUNKS_PER_GROUP})")
    parser.add_argument("--mapreduce-threshold-chars", type=int, default=DEFAULT_MAPREDUCE_THRESHOLD_CHARS,
                        help=f"Content length above which map-reduce kicks in (default: {DEFAULT_MAPREDUCE_THRESHOLD_CHARS})")
    args = parser.parse_args()

    if not args.dsn:
        _log.error("missing_dsn", hint="pass --dsn or set KB_DSN")
        sys.exit(1)

    if args.embed_summaries:
        stats = asyncio.run(
            run_embed_summaries(
                dsn=args.dsn,
                ollama_url=args.ollama_url,
                embed_model=args.embed_model,
                model_filter=args.model,
                limit=args.limit,
                offset=args.offset,
                apply=args.apply,
            )
        )
        mode = "APPLY" if args.apply else "DRY-RUN"
        _log.info("summary", mode=mode, **stats)
        balanced = stats["embedded"] + stats["errors"] == stats["summaries_fetched"]
        sys.exit(1 if (stats["errors"] > 0 or not balanced) else 0)

    if not args.backend or not args.model:
        _log.error("missing_backend_or_model", hint="--backend and --model are required unless --embed-summaries")
        sys.exit(1)
    if args.backend == "anthropic" and args.apply and not args.anthropic_api_key:
        _log.error("missing_anthropic_api_key", hint="pass --anthropic-api-key or set ANTHROPIC_API_KEY")
        sys.exit(1)
    if not args.tags_vocab.exists():
        _log.error("missing_tags_vocab", path=str(args.tags_vocab))
        sys.exit(1)

    try:
        stats = asyncio.run(
            run_summarize(
                dsn=args.dsn,
                backend_name=args.backend,
                model=args.model,
                ollama_url=args.ollama_url,
                anthropic_api_key=args.anthropic_api_key,
                tags_vocab_path=args.tags_vocab,
                limit=args.limit,
                offset=args.offset,
                apply=args.apply,
                max_tokens=args.max_tokens,
                chunks_per_group=args.chunks_per_group,
                mapreduce_threshold_chars=args.mapreduce_threshold_chars,
            )
        )
    except EmbeddingDimensionError as exc:
        _log.error("dim_mismatch_abort", error=str(exc))
        sys.exit(1)

    mode = "APPLY" if args.apply else "DRY-RUN"
    _log.info("summary", mode=mode, backend=args.backend, model=args.model, **stats)

    balanced = stats["documents_fetched"] == (
        stats["duplicates_skipped"] + stats["no_active_chunks"] + stats["already_summarized"]
        + stats["summarized"] + stats["llm_errors"]
    )
    failed = stats["llm_errors"] > 0 or not balanced
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()