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

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 jobs/documents-ingest/

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) 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`. Note: that UNIQUE
constraint is on (envelope_id, chunk_index) only, not including `model` — re-embedding the
same pilot with a *different* model would hit this path and its embedding would be
discarded (a real chunk of work wasted, correctly reported, but not persisted). Out of
scope for this pilot (single model, bge-m3); flagged as a follow-up if/when a second model
is ever indexed (the real fix is `UNIQUE (envelope_id, chunk_index, model)` at the schema
layer).
"""
from __future__ import annotations

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

import aiohttp
import asyncpg
import structlog

_log = structlog.get_logger(__name__)

DEFAULT_OLLAMA_URL = "http://localhost:11434"
DEFAULT_MODEL = "bge-m3"
EXPECTED_DIM = 1024

# Plan decision 3: ~600 tok/chunk, ~150 tok overlap. No local bge-m3 tokenizer available
# offline, so tokens are approximated at ~4 chars/token (common heuristic) -> ~2400 chars
# target, ~600 chars overlap. Paragraph boundaries are preferred; a hard character-based
# sliding window is only used when a single paragraph alone exceeds the target (plan §1.2:
# Paperless OCR text has no page-break markers to split on instead).
CHARS_PER_TOKEN = 4
TARGET_TOKENS = 600
OVERLAP_TOKENS = 150
TARGET_CHARS = TARGET_TOKENS * CHARS_PER_TOKEN
OVERLAP_CHARS = OVERLAP_TOKENS * CHARS_PER_TOKEN

_PARA_SPLIT = re.compile(r"\n\s*\n")

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


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


def split_paragraphs(text: str) -> list[str]:
    """Split on blank-line boundaries; drop empty fragments."""
    return [p.strip() for p in _PARA_SPLIT.split(text) if p.strip()]


def hard_split(text: str, size: int = TARGET_CHARS, overlap: int = OVERLAP_CHARS) -> list[str]:
    """Fixed-size sliding window over raw characters — fallback for a paragraph too big to fit."""
    if overlap >= size:
        # step = size - overlap would be <= 0 below, so `start` would never advance past `n`
        # and the while loop would spin forever, growing `chunks` without bound.
        raise ValueError(f"chunk_overlap ({overlap}) must be smaller than chunk_size ({size})")
    if len(text) <= size:
        return [text]
    step = size - overlap
    chunks = []
    start = 0
    n = len(text)
    while start < n:
        end = min(start + size, n)
        chunks.append(text[start:end])
        if end == n:
            break
        start += step
    return chunks


def chunk_text(text: str, size: int = TARGET_CHARS, overlap: int = OVERLAP_CHARS) -> list[str]:
    """Paragraph-preferring chunker with a hard character-fallback (plan §2 decision 3).

    Empty/whitespace-only text (e.g. the 26 empty_content Paperless documents, plan §1.2)
    returns []. A document shorter than one chunk returns exactly one chunk (the full text).
    """
    text = (text or "").strip()
    if not text:
        return []

    paragraphs = split_paragraphs(text)

    chunks: list[str] = []
    buffer = ""
    for para in paragraphs:
        if len(para) > size:
            if buffer:
                chunks.append(buffer)
                buffer = ""
            chunks.extend(hard_split(para, size, overlap))
            continue

        candidate = f"{buffer}\n\n{para}" if buffer else para
        if len(candidate) <= size:
            buffer = candidate
            continue

        chunks.append(buffer)
        tail = buffer[-overlap:] if buffer else ""
        buffer = f"{tail}\n\n{para}" if tail else para

    if buffer:
        chunks.append(buffer)
    return chunks


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


def _vector_literal(embedding: list[float]) -> str:
    """Render a Python float list as a pgvector input literal, e.g. '[0.1,0.2,...]'."""
    return "[" + ",".join(repr(v) for v in embedding) + "]"


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 embed_chunk(
    session: aiohttp.ClientSession, base_url: str, model: str, text: str
) -> tuple[list[float], float]:
    """POST /api/embeddings on Ollama for one chunk. Returns (embedding, elapsed_seconds)."""
    t0 = time.monotonic()
    async with session.post(f"{base_url}/api/embeddings", json={"model": model, "prompt": text}) as resp:
        resp.raise_for_status()
        data = await resp.json()
    elapsed = time.monotonic() - t0

    embedding = data.get("embedding")
    if not embedding:
        raise ValueError(f"ollama response missing 'embedding': {data!r}")
    return embedding, elapsed


async def insert_chunk(
    conn: asyncpg.Connection, envelope_id: str, chunk_index: int, text: str,
    embedding: list[float], model: str,
) -> 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."""
    return await conn.execute(
        _INSERT_SQL, envelope_id, chunk_index, text, _vector_literal(embedding), model
    )


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_conflict_skipped + chunks_errors

    `chunks_conflict_skipped` counts inserts where `ON CONFLICT (envelope_id, chunk_index)
    DO NOTHING` silently discarded the row (the pre-fetched `existing` set is the first line
    of defense against this and should make it rare; see the module docstring's note on the
    UNIQUE constraint not including `model`) — 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_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

                    if not apply:
                        stats["chunks_inserted"] += 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_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_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()