#!/usr/bin/env python3
"""Chunk a document set, embed it through an OpenAI-compatible server, and index it locally.

Purpose: the ingest half of the private question-answering project. Reads .md and .txt files,
    splits them on their own heading structure, packs the pieces into overlapping windows with
    the heading path prepended, embeds each window through /v1/embeddings, and writes the text
    and the vectors into one SQLite file using the sqlite-vec extension. Re-running it
    re-ingests only the files whose contents changed.
Platform: all (pure Python over HTTP; the embedding server may be on any track)
Minimum memory: 8 GB on the machine running the embedding model; this script needs very little
Assumes: Python 3.9 or later, `sqlite-vec` installed in the active environment, SQLite 3.41 or
    later, and an OpenAI-compatible embeddings endpoint reachable at --embed-url. Start one
    with:  llama-server --model <embedding-model.gguf> --embedding --port 8090
    The database file is created if it does not exist and is safe to delete: everything in it
    can be rebuilt from the documents.

Usage: python3 ingest.py --docs sample-docs --db qa-index.db --embed-url http://127.0.0.1:8090/v1
       python3 ingest.py --docs ~/notes --db qa-index.db --rebuild --labbook labbook.md
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import sqlite3
import struct
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Iterable, List, Optional, Tuple

try:
    import sqlite_vec
except ImportError:  # pragma: no cover - environment check, not logic
    sys.exit("sqlite-vec is not installed. Run: uv pip install sqlite-vec")

TEXT_SUFFIXES = {".md", ".txt"}

# Chunk sizes are given in words because this script has no tokeniser. For English prose a
# word is roughly 1.3 tokens, so 300 words is about 400 tokens. Adjust for your own documents
# rather than trusting the ratio: the page explains how to check it.
DEFAULT_WORDS = 300
DEFAULT_OVERLAP = 60


# --------------------------------------------------------------------------------------
# Reading and chunking
# --------------------------------------------------------------------------------------

def find_documents(root: Path) -> List[Path]:
    if root.is_file():
        return [root]
    return sorted(p for p in root.rglob("*") if p.suffix.lower() in TEXT_SUFFIXES and p.is_file())


def split_sections(text: str) -> List[Tuple[str, str]]:
    """Split Markdown on ATX headings, returning (heading path, body) pairs.

    A file with no headings comes back as a single section with an empty path, which is the
    correct behaviour for a plain .txt file.
    """
    lines = text.splitlines()
    sections: List[Tuple[str, str]] = []
    stack: List[str] = []
    buffer: List[str] = []

    def flush() -> None:
        body = "\n".join(buffer).strip()
        if body:
            sections.append((" > ".join(stack), body))
        buffer.clear()

    for line in lines:
        match = re.match(r"^(#{1,6})\s+(.*\S)\s*$", line)
        if match:
            flush()
            level = len(match.group(1))
            del stack[level - 1:]
            stack.append(match.group(2))
        else:
            buffer.append(line)
    flush()
    return sections


def pack_paragraphs(body: str, target_words: int, overlap_words: int) -> List[str]:
    """Pack paragraphs into windows of about target_words, overlapping by overlap_words."""
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", body) if p.strip()]
    windows: List[str] = []
    current: List[str] = []
    current_words = 0

    for paragraph in paragraphs:
        words = len(paragraph.split())
        if current and current_words + words > target_words:
            windows.append("\n\n".join(current))
            # Carry the tail of the window forward so a fact on a boundary appears whole
            # in one of the two neighbours.
            tail: List[str] = []
            tail_words = 0
            for previous in reversed(current):
                if tail_words >= overlap_words:
                    break
                tail.insert(0, previous)
                tail_words += len(previous.split())
            current = tail
            current_words = tail_words
        current.append(paragraph)
        current_words += words

    if current:
        windows.append("\n\n".join(current))
    return windows


def heading_path(title: str, heading: str) -> str:
    """Join the file's name to its heading path, without repeating a title-like top heading."""
    if not heading:
        return title
    first = heading.split(" > ")[0].lower()
    if all(word in first for word in title.lower().split()):
        return heading
    return f"{title} > {heading}"


def chunk_document(path: Path, root: Path, target_words: int, overlap_words: int) -> List[dict]:
    """One document to a list of chunk records, with the heading path prepended to the text."""
    text = path.read_text(encoding="utf-8", errors="replace")
    source = str(path.relative_to(root)) if root.is_dir() else path.name
    title = path.stem.replace("-", " ")
    records: List[dict] = []
    for heading, body in split_sections(text):
        full_heading = heading_path(title, heading)
        for window in pack_paragraphs(body, target_words, overlap_words):
            records.append({
                "source": source,
                "heading": full_heading,
                # The heading path is part of the embedded text on purpose: a passage that
                # never repeats its own topic words matches no question about them.
                "text": f"{full_heading}\n\n{window}",
            })
    for index, record in enumerate(records):
        record["chunk_index"] = index
    return records


def file_hash(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


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

def embed_batch(texts: List[str], url: str, model: str, api_key: Optional[str],
                timeout: int) -> List[List[float]]:
    endpoint = url.rstrip("/") + "/embeddings"
    payload = json.dumps({"model": model, "input": texts}).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(endpoint, data=payload, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            body = json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError(f"{endpoint} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {endpoint}: {exc.reason}") from exc

    rows = sorted(body["data"], key=lambda row: row.get("index", 0))
    return [row["embedding"] for row in rows]


def serialise(vector: Iterable[float]) -> bytes:
    """The compact float32 blob sqlite-vec stores. Same layout as sqlite_vec.serialize_float32."""
    values = list(vector)
    return struct.pack(f"{len(values)}f", *values)


# --------------------------------------------------------------------------------------
# The index
# --------------------------------------------------------------------------------------

def open_db(path: str) -> sqlite3.Connection:
    db = sqlite3.connect(path)
    db.enable_load_extension(True)
    sqlite_vec.load(db)
    db.enable_load_extension(False)
    db.execute("pragma journal_mode=wal")
    return db


def ensure_base_schema(db: sqlite3.Connection) -> None:
    """The two ordinary tables. Created before anything is embedded, because the incremental
    path has to read them even on a run where nothing changed."""
    db.execute("create table if not exists meta (key text primary key, value text not null)")
    db.execute("""create table if not exists chunks (
                    id integer primary key,
                    source text not null,
                    heading text not null,
                    chunk_index integer not null,
                    text text not null,
                    source_hash text not null)""")
    db.execute("create index if not exists chunks_source on chunks(source)")
    db.commit()


def ensure_vec_schema(db: sqlite3.Connection, dim: int, model: str) -> None:
    """The vector table, whose dimension is fixed at creation, plus the model safety catch."""
    db.execute(f"""create virtual table if not exists vec_chunks using vec0(
                     chunk_id integer primary key,
                     embedding float[{dim}] distance_metric=cosine)""")

    stored = dict(db.execute("select key, value from meta").fetchall())
    if stored:
        # An index built with a different model or dimension cannot be queried with this one.
        # Distances between vectors from different models are meaningless, and nothing
        # downstream would notice, so this is the safety catch.
        if stored.get("embedding_model") != model or int(stored.get("dimension", 0)) != dim:
            raise SystemExit(
                f"this index was built with {stored.get('embedding_model')} at dimension "
                f"{stored.get('dimension')}; you are ingesting with {model} at {dim}. "
                "Re-run with --rebuild, or use a different --db."
            )
    else:
        db.executemany("insert into meta(key, value) values (?, ?)", [
            ("embedding_model", model),
            ("dimension", str(dim)),
            ("created", time.strftime("%Y-%m-%dT%H:%M:%S")),
        ])
    db.commit()


def delete_source(db: sqlite3.Connection, source: str) -> int:
    ids = [row[0] for row in db.execute("select id from chunks where source = ?", (source,))]
    if not ids:
        return 0
    marks = ",".join("?" for _ in ids)
    db.execute(f"delete from vec_chunks where chunk_id in ({marks})", ids)
    db.execute("delete from chunks where source = ?", (source,))
    return len(ids)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--docs", required=True, help="directory or file to ingest")
    parser.add_argument("--db", default="qa-index.db", help="SQLite index to create or update")
    parser.add_argument("--embed-url", default="http://127.0.0.1:8090/v1",
                        help="OpenAI-compatible base URL of the embedding server")
    parser.add_argument("--embed-model", default="qwen3-embedding",
                        help="model name or alias the embedding server answers to")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--words", type=int, default=DEFAULT_WORDS, help="target words per chunk")
    parser.add_argument("--overlap", type=int, default=DEFAULT_OVERLAP, help="overlap in words")
    parser.add_argument("--batch", type=int, default=16, help="chunks per embedding request")
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--rebuild", action="store_true", help="delete the index and start again")
    parser.add_argument("--dry-run", action="store_true", help="chunk and report; embed nothing")
    parser.add_argument("--labbook", default=None, help="append one JSON line to this file")
    args = parser.parse_args()

    root = Path(args.docs).expanduser()
    if not root.exists():
        sys.exit(f"--docs {root} does not exist")

    documents = find_documents(root)
    if not documents:
        sys.exit(f"no .md or .txt files under {root}")

    if args.rebuild and Path(args.db).exists():
        Path(args.db).unlink()
        print(f"==> removed {args.db}")

    started = time.time()
    all_chunks = {}
    for path in documents:
        records = chunk_document(path, root, args.words, args.overlap)
        if not records:
            print(f"    {path}: no text, skipped")
            continue
        digest = file_hash(path)
        for record in records:
            record["source_hash"] = digest
        all_chunks[records[0]["source"]] = (digest, records)
        print(f"    {path}: {len(records)} chunk(s)")

    total_chunks = sum(len(records) for _, records in all_chunks.values())
    print(f"==> {len(documents)} document(s), {total_chunks} chunk(s)")

    if args.dry_run:
        longest = max((len(r["text"].split()) for _, rs in all_chunks.values() for r in rs), default=0)
        print(f"    longest chunk: {longest} words. Nothing was embedded (--dry-run).")
        return

    db = open_db(args.db)
    ensure_base_schema(db)
    ingested, skipped = 0, 0
    dimension = None
    stored_dim = db.execute("select value from meta where key = 'dimension'").fetchone()
    if stored_dim:
        dimension = int(stored_dim[0])

    for source, (digest, records) in sorted(all_chunks.items()):
        existing = db.execute(
            "select distinct source_hash from chunks where source = ?", (source,)).fetchall()
        if existing and existing[0][0] == digest:
            skipped += len(records)
            continue

        vectors: List[List[float]] = []
        for start in range(0, len(records), args.batch):
            batch = [r["text"] for r in records[start:start + args.batch]]
            vectors.extend(embed_batch(batch, args.embed_url, args.embed_model,
                                       args.api_key, args.timeout))
        if dimension is None:
            dimension = len(vectors[0])
        elif len(vectors[0]) != dimension:
            sys.exit(
                f"the embedding server returned a {len(vectors[0])}-dimension vector where "
                f"this index expects {dimension}. Re-run with --rebuild if you changed model."
            )
        ensure_vec_schema(db, dimension, args.embed_model)

        removed = delete_source(db, source)
        if removed:
            print(f"    {source}: replacing {removed} stale chunk(s)")
        for record, vector in zip(records, vectors):
            cursor = db.execute(
                "insert into chunks(source, heading, chunk_index, text, source_hash) "
                "values (?, ?, ?, ?, ?)",
                (record["source"], record["heading"], record["chunk_index"],
                 record["text"], record["source_hash"]))
            db.execute("insert into vec_chunks(chunk_id, embedding) values (?, ?)",
                       (cursor.lastrowid, serialise(vector)))
        db.commit()
        ingested += len(records)
        print(f"    {source}: {len(records)} chunk(s) embedded and indexed")

    elapsed = time.time() - started
    indexed = db.execute("select count(*) from chunks").fetchone()[0]
    print(f"==> {ingested} chunk(s) ingested, {skipped} unchanged, {indexed} in the index "
          f"({elapsed:.1f} s)")

    if args.labbook:
        record = {
            "lab": "part-10/project-private-document-qa/ingest",
            "docs": str(root),
            "documents": len(documents),
            "chunks_total": indexed,
            "chunks_ingested": ingested,
            "chunks_unchanged": skipped,
            "embed_model": args.embed_model,
            "dimension": dimension,
            "words_per_chunk": args.words,
            "overlap_words": args.overlap,
            "seconds": round(elapsed, 1),
            "date": time.strftime("%Y-%m-%d"),
        }
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record) + "\n")
        print(f"    recorded in {args.labbook}")


if __name__ == "__main__":
    main()