"""Chunk + embed job — module 5 phase 2, plan step 6 (kb/phases/kb-m5-faza2.md,
§6 step 6, decision 3).

`chunk_text`/`hard_split`/`split_paragraphs`/`TARGET_CHARS`/`OVERLAP_CHARS` moved to
`kb_mail.chunking` in module 5 faza mailowa, Krok 0 (kb/phases/kb-m5-faza-mailowa.md, §3)
so `jobs/mail-body-ingest` shares the exact same chunker instead of a copy-pasted drift; re-exported
here unchanged so nothing importing them from this module breaks.

Pipeline: `envelope(source='paperless').entities[type=content].text` -> chunk (paragraph-
preferring, ~600 tok/chunk, ~150 tok overlap, hard char-fallback for oversized paragraphs)
-> `POST /api/embeddings` (Ollama on SOLARIA, model `bge-m3`) -> `INSERT document_chunk`
(`services/kb-postgres/init/002_chunks.sql`, untouched by this change).

Scope is deliberately narrow to the pilot (plan §6 step 7): only `source='paperless'`
envelopes. Mail chunking/embedding is a later phase (plan §1.1) and will reuse this same
`chunk_text()` + `document_chunk` table, not a new pipeline.

Runs on SOLARIA (needs Ollama on localhost) against kb-postgres@PIHA over Tailscale:

Install (from repo root):
    pip install -e packages/kb-mail/
    pip install -e packages/kb-retrieval/
    pip install -e jobs/documents-ingest/

`embed_chunk`/`_vector_literal`/`DEFAULT_MODEL`/`DEFAULT_OLLAMA_URL` moved to
`kb_retrieval.embed` in module 5 phase 4 (kb/phases/kb-m5-faza4.md, §3, decision 1) so
`kb-query` (Docker service) can share the same client without pulling in this job's `anthropic`
dependency; re-exported here unchanged so nothing importing them from this module breaks.

Usage:
    # Dry run (default) — chunk and count, no Ollama calls, no DB writes:
    documents-ingest-embed --dsn postgresql://kb:<pw>@piha:5433/kb

    # Real run:
    documents-ingest-embed --dsn ... --apply

    # Smoke-test slice:
    documents-ingest-embed --dsn ... --apply --limit 10

DSN can come from KB_DSN, Ollama URL from OLLAMA_URL (default http://localhost:11434).

Idempotency: a pre-fetched set of existing (envelope_id, chunk_index) pairs for this
`model` skips chunks already embedded — no re-embedding, no wasted Ollama calls on rerun.
`insert_chunk`'s own `ON CONFLICT (envelope_id, chunk_index, model) DO NOTHING` is the
second line of defense; its command tag is checked so a silently-skipped row is counted as
`chunks_conflict_skipped`, never miscounted as `chunks_inserted`.

OCR-junk filter (module 5, phase 3, plan §3.1): `is_ocr_junk()` screens each chunk before
embedding. A junk chunk is INSERTed with `excluded_reason='ocr_junk'` and `embedding=NULL`
— no Ollama call, no HNSW entry — and counted as `chunks_junk_flagged`. The three signals
and their thresholds are calibrated in-code (see `is_ocr_junk` docstring) against the
2683-chunk pilot corpus; retrieval and all downstream consumers filter on
`WHERE excluded_reason IS NULL`.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import sys
from typing import Optional

import aiohttp
import asyncpg
import structlog
from kb_mail.chunking import OVERLAP_CHARS, TARGET_CHARS, chunk_text, hard_split, split_paragraphs
from kb_retrieval.embed import DEFAULT_MODEL, DEFAULT_OLLAMA_URL, _vector_literal, embed_chunk

_log = structlog.get_logger(__name__)

EXPECTED_DIM = 1024

_INSERT_SQL = """
INSERT INTO document_chunk (envelope_id, chunk_index, text, embedding, model, excluded_reason)
VALUES ($1, $2, $3, $4::vector, $5, $6)
ON CONFLICT (envelope_id, chunk_index, model) DO NOTHING
"""

# Plan §3.1, signal 1: refined from "any C0 control char" to a raw-count threshold during
# calibration on the 2683-chunk pilot corpus (2026-07-16) — legit OCR text carries a handful
# of stray control bytes (paperless:119 mojibake, a confirmed retrieval hit; an English paper
# with 1-2 stray bytes per chunk), while true binary/barcode noise carries 10-63 per chunk.
# A threshold of 5 cleanly separates the two on this corpus.
CONTROL_CHAR_JUNK_THRESHOLD = 5
# Plan §3.1, signal 2: share of characters outside the "wordy" class (alnum + PL diacritics +
# common punctuation + whitespace). On the pilot corpus no legitimate chunk crosses 20%, so
# this never fires alone here — kept for corpora where junk isn't diluted by legible headers.
NON_WORDY_RATIO_THRESHOLD = 0.30
# Plan §3.1, signal 3: share of whitespace-split tokens that look like words. Legit text
# clusters well above 40%; junk (scrambled fonts, dot-leader ToCs) sits below 25% on the
# pilot corpus.
WORDLIKE_RATIO_THRESHOLD = 0.25

_CONTROL_CHAR_RE = re.compile(r"[\x00-\x08\x0b-\x1f]")
_NON_WORDY_CHAR_RE = re.compile(
    r"[^a-zA-Z0-9ąćęłńóśźżĄĆĘŁŃÓŚŹŻ\s.,;:!?\"'`()\[\]{}<>/\\|@#$%^&*_+=~-]"
)
_WORDLIKE_TOKEN_RE = re.compile(r"^[a-zA-Ząćęłńóśźż-]{2,30}$")


def is_ocr_junk(text: str) -> bool:
    """Three-signal OCR-junk heuristic (plan §3.1, decision 1): binary noise (ground barcodes,
    scrambled fonts), not merely ugly text — partial mojibake alone is not junk (paperless:119
    stays a confirmed retrieval hit despite it). Chunk-level, never document-level."""
    if not text:
        return False
    if len(_CONTROL_CHAR_RE.findall(text)) >= CONTROL_CHAR_JUNK_THRESHOLD:
        return True
    if len(_NON_WORDY_CHAR_RE.findall(text)) / len(text) > NON_WORDY_RATIO_THRESHOLD:
        return True
    tokens = text.split()
    if not tokens:
        return False
    wordlike = sum(1 for t in tokens if _WORDLIKE_TOKEN_RE.match(t))
    return (wordlike / len(tokens)) < WORDLIKE_RATIO_THRESHOLD


class EmbeddingDimensionError(RuntimeError):
    """Ollama returned a vector of the wrong dimension for the target `document_chunk` schema."""


def extract_content(entities: list) -> str:
    """Pull `entities[type=content].text` out of a document envelope (plan §4.2). Missing
    or empty content both yield ""."""
    for entity in entities or []:
        if isinstance(entity, dict) and entity.get("type") == "content":
            return entity.get("text") or ""
    return ""


def _decode_jsonb(value: object) -> object:
    """asyncpg may return jsonb as a str or an already-decoded object."""
    if value is None:
        return None
    if isinstance(value, str):
        return json.loads(value)
    return value


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_existing_chunk_keys(conn: asyncpg.Connection, model: str) -> set[tuple[str, int]]:
    """(envelope_id, chunk_index) pairs already embedded with this model — idempotency + dry-run preview."""
    rows = await conn.fetch(
        "SELECT envelope_id, chunk_index FROM document_chunk WHERE model = $1", model
    )
    return {(r["envelope_id"], r["chunk_index"]) for r in rows}


async def insert_chunk(
    conn: asyncpg.Connection, envelope_id: str, chunk_index: int, text: str,
    embedding: Optional[list[float]], model: str, excluded_reason: Optional[str] = None,
) -> str:
    """Returns asyncpg's command tag (e.g. 'INSERT 0 1' or 'INSERT 0 0' if ON CONFLICT
    DO NOTHING skipped the row) so the caller can tell a real insert from a no-op.
    `embedding=None` (junk chunks, `excluded_reason='ocr_junk'`) inserts a NULL vector —
    pgvector's HNSW index skips NULLs automatically."""
    vector_literal = _vector_literal(embedding) if embedding is not None else None
    return await conn.execute(
        _INSERT_SQL, envelope_id, chunk_index, text, vector_literal, model, excluded_reason
    )


def _rows_affected(command_tag: str) -> int:
    """Parse the row count out of an asyncpg INSERT command tag ('INSERT <oid> <rows>')."""
    return int(command_tag.rsplit(" ", 1)[-1])


async def run(
    dsn: str,
    ollama_url: str = DEFAULT_OLLAMA_URL,
    model: str = DEFAULT_MODEL,
    limit: Optional[int] = None,
    offset: Optional[int] = None,
    apply: bool = False,
    chunk_size: int = TARGET_CHARS,
    chunk_overlap: int = OVERLAP_CHARS,
) -> dict:
    """Chunk + embed one --limit/--offset slice of `source='paperless'` envelopes.

    dry-run (apply=False) chunks and counts everything — no Ollama calls, no DB writes;
    `chunks_inserted` reports what *would* be written (mirrors the rest of this job family).
    Returns stats that must always balance:

        documents_fetched  = empty_content + documents_chunked
        chunks_total        = chunks_already_embedded + chunks_inserted + chunks_junk_flagged
                               + chunks_conflict_skipped + chunks_errors

    Each chunk is screened by `is_ocr_junk()` before embedding (plan §3.1): a junk chunk
    skips Ollama entirely and is INSERTed with `excluded_reason='ocr_junk'`,
    `embedding=NULL`, counted as `chunks_junk_flagged` — never as `chunks_inserted`.

    `chunks_conflict_skipped` counts inserts where `ON CONFLICT (envelope_id, chunk_index,
    model) DO NOTHING` silently discarded the row (the pre-fetched `existing` set is the
    first line of defense against this and should make it rare) — tracked separately so a
    silent no-op is never miscounted as a successful write.

    Every embedding response's dimension is checked against `EXPECTED_DIM` (raises
    EmbeddingDimensionError and aborts the whole run on mismatch) — never silently indexes
    a vector that doesn't match the VECTOR(1024) column. A DB write failure for one chunk
    (bad bytes, a dropped connection) is isolated and counted as `chunks_errors`, same as an
    embed failure — it never aborts the run for the rest of the slice.
    """
    stats = {
        "documents_fetched": 0,
        "empty_content": 0,
        "documents_chunked": 0,
        "chunks_total": 0,
        "chunks_already_embedded": 0,
        "chunks_inserted": 0,
        "chunks_junk_flagged": 0,
        "chunks_conflict_skipped": 0,
        "chunks_errors": 0,
        "embed_calls": 0,
        "embed_seconds_total": 0.0,
    }

    conn = await asyncpg.connect(dsn)
    try:
        docs = await fetch_documents(conn, limit, offset)
        existing = await fetch_existing_chunk_keys(conn, model)

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

        try:
            for row in docs:
                stats["documents_fetched"] += 1
                envelope_id = row["id"]
                entities = _decode_jsonb(row["entities"]) or []
                content = extract_content(entities)
                chunks = chunk_text(content, chunk_size, chunk_overlap)

                if not chunks:
                    stats["empty_content"] += 1
                    continue

                stats["documents_chunked"] += 1
                for idx, chunk in enumerate(chunks):
                    stats["chunks_total"] += 1
                    key = (envelope_id, idx)
                    if key in existing:
                        stats["chunks_already_embedded"] += 1
                        continue

                    junk = is_ocr_junk(chunk)

                    if not apply:
                        if junk:
                            stats["chunks_junk_flagged"] += 1
                        else:
                            stats["chunks_inserted"] += 1
                        continue

                    if junk:
                        try:
                            command_tag = await insert_chunk(
                                conn, envelope_id, idx, chunk, None, model, excluded_reason="ocr_junk"
                            )
                        except Exception:
                            _log.warning(
                                "skip.insert_error", envelope_id=envelope_id, chunk_index=idx, exc_info=True
                            )
                            stats["chunks_errors"] += 1
                            continue

                        existing.add(key)
                        if _rows_affected(command_tag) == 0:
                            _log.warning(
                                "chunk.conflict_skipped", envelope_id=envelope_id, chunk_index=idx, model=model
                            )
                            stats["chunks_conflict_skipped"] += 1
                        else:
                            stats["chunks_junk_flagged"] += 1
                        continue

                    assert session is not None
                    try:
                        embedding, elapsed = await embed_chunk(session, ollama_url, model, chunk)
                    except Exception:
                        _log.warning(
                            "skip.embed_error", envelope_id=envelope_id, chunk_index=idx, exc_info=True
                        )
                        stats["chunks_errors"] += 1
                        continue

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

                    stats["embed_calls"] += 1
                    stats["embed_seconds_total"] += elapsed

                    try:
                        command_tag = await insert_chunk(conn, envelope_id, idx, chunk, embedding, model)
                    except Exception:
                        _log.warning(
                            "skip.insert_error", envelope_id=envelope_id, chunk_index=idx, exc_info=True
                        )
                        stats["chunks_errors"] += 1
                        continue

                    existing.add(key)
                    if _rows_affected(command_tag) == 0:
                        _log.warning(
                            "chunk.conflict_skipped", envelope_id=envelope_id, chunk_index=idx, model=model
                        )
                        stats["chunks_conflict_skipped"] += 1
                    else:
                        stats["chunks_inserted"] += 1
        finally:
            if session is not None:
                await session.close()
    finally:
        await conn.close()

    balance_docs = stats["empty_content"] + stats["documents_chunked"]
    balance_chunks = (
        stats["chunks_already_embedded"] + stats["chunks_inserted"] + stats["chunks_junk_flagged"]
        + stats["chunks_conflict_skipped"] + stats["chunks_errors"]
    )
    if balance_docs != stats["documents_fetched"] or balance_chunks != stats["chunks_total"]:
        _log.error("stats_mismatch", **stats)

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


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Chunk source='paperless' envelope content and embed via Ollama/bge-m3 "
        "into document_chunk (module 5, phase 2 — plan §6 step 6)."
    )
    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("--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("--model", default=os.environ.get("OLLAMA_EMBED_MODEL", DEFAULT_MODEL),
                        help=f"Ollama embedding model (default: {DEFAULT_MODEL})")
    parser.add_argument("--limit", type=int, default=None,
                        help="Max documents to process (default: all)")
    parser.add_argument("--offset", type=int, default=0,
                        help="Slice offset, ordered by envelope id (default: 0)")
    parser.add_argument("--chunk-size", type=int, default=TARGET_CHARS,
                        help=f"Target chunk size in characters (default: {TARGET_CHARS} ~= {TARGET_TOKENS} tok)")
    parser.add_argument("--chunk-overlap", type=int, default=OVERLAP_CHARS,
                        help=f"Chunk overlap in characters (default: {OVERLAP_CHARS} ~= {OVERLAP_TOKENS} tok)")
    parser.add_argument("--apply", action="store_true",
                        help="Actually call Ollama and insert chunks. Default is dry-run (chunk + count only).")
    args = parser.parse_args()

    if not args.dsn:
        _log.error("missing_dsn", hint="pass --dsn or set KB_DSN")
        sys.exit(1)
    if args.chunk_overlap >= args.chunk_size:
        _log.error(
            "invalid_chunk_params", chunk_size=args.chunk_size, chunk_overlap=args.chunk_overlap,
            hint="--chunk-overlap must be smaller than --chunk-size",
        )
        sys.exit(1)

    try:
        stats = asyncio.run(
            run(
                dsn=args.dsn,
                ollama_url=args.ollama_url,
                model=args.model,
                limit=args.limit,
                offset=args.offset,
                apply=args.apply,
                chunk_size=args.chunk_size,
                chunk_overlap=args.chunk_overlap,
            )
        )
    except EmbeddingDimensionError as exc:
        _log.error("dim_mismatch_abort", error=str(exc))
        sys.exit(1)

    mode = "APPLY" if args.apply else "DRY-RUN"
    avg_embed = (
        stats["embed_seconds_total"] / stats["embed_calls"] if stats["embed_calls"] else 0.0
    )
    _log.info("summary", mode=mode, avg_embed_seconds_per_chunk=round(avg_embed, 4), **stats)

    balanced = (
        stats["documents_fetched"] == stats["empty_content"] + stats["documents_chunked"]
        and stats["chunks_total"] == (
            stats["chunks_already_embedded"] + stats["chunks_inserted"] + stats["chunks_junk_flagged"]
            + stats["chunks_conflict_skipped"] + stats["chunks_errors"]
        )
    )
    failed = stats["chunks_errors"] > 0 or stats["chunks_conflict_skipped"] > 0 or not balanced
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()