#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = ["sqlite-vec==0.1.6"]
# ///
"""Data-access layer for the wiki skill.

Pure SQL plus lifecycle helpers. No printing, no argparse, no sys.exit.

`vec_chunks` is a vec0 virtual table and therefore NOT reachable by the foreign key
cascade that cleans up `chunks` and `chunks_fts`. Every path that removes chunks must
delete their vectors explicitly — that is why the deletes here go through
`_delete_vectors_for_file`.
"""

from __future__ import annotations

import json
import sqlite3
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from pathlib import Path

from sqlite_vec import serialize_float32  # ty: ignore[unresolved-import]
from wiki_db import get_db, init_db


@contextmanager
def connection(db_path: Path) -> Iterator[sqlite3.Connection]:
    """Open a connection, initialising the DB if missing."""
    if not db_path.exists():
        init_db(db_path)
    conn = get_db(db_path)
    try:
        yield conn
    finally:
        conn.close()


@contextmanager
def tx(conn: sqlite3.Connection) -> Iterator[sqlite3.Connection]:
    """Wrap an already-open connection in an explicit transaction.

    The connection is in autocommit mode (`isolation_level=None`), so transactions are
    ours to open — sqlite3's implicit handling would otherwise fail a nested BEGIN.
    """
    conn.execute("BEGIN")
    try:
        yield conn
        conn.execute("COMMIT")
    except Exception:
        conn.execute("ROLLBACK")
        raise


@contextmanager
def transaction(db_path: Path) -> Iterator[sqlite3.Connection]:
    """Open a connection wrapped in an explicit transaction."""
    with connection(db_path) as conn, tx(conn):
        yield conn


# ---------------------------------------------------------------------------
# meta — identity of the embedding space
# ---------------------------------------------------------------------------


def read_meta(conn: sqlite3.Connection) -> dict[str, str]:
    return {row["key"]: row["value"] for row in conn.execute("SELECT key, value FROM meta")}


def write_meta(conn: sqlite3.Connection, values: dict[str, str]) -> None:
    conn.executemany(
        "INSERT INTO meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value",
        sorted(values.items()),
    )


# ---------------------------------------------------------------------------
# sources
# ---------------------------------------------------------------------------


def upsert_source(conn: sqlite3.Connection, source_id: str, kind: str) -> None:
    conn.execute(
        "INSERT INTO sources (source_id, kind) VALUES (?, ?) ON CONFLICT(source_id) DO UPDATE SET kind = excluded.kind",
        (source_id, kind),
    )


def get_source(conn: sqlite3.Connection, source_id: str) -> sqlite3.Row | None:
    return conn.execute("SELECT * FROM sources WHERE source_id = ?", (source_id,)).fetchone()


def mark_synced(conn: sqlite3.Connection, source_id: str, indexed_rev: str | None, now: str) -> None:
    conn.execute(
        "UPDATE sources SET indexed_rev = ?, last_sync_at = ? WHERE source_id = ?",
        (indexed_rev, now, source_id),
    )


# ---------------------------------------------------------------------------
# files
# ---------------------------------------------------------------------------


def get_file(conn: sqlite3.Connection, source_id: str, path: str) -> sqlite3.Row | None:
    return conn.execute("SELECT * FROM files WHERE source_id = ? AND path = ?", (source_id, path)).fetchone()


def list_source_files(conn: sqlite3.Connection, source_id: str) -> dict[str, sqlite3.Row]:
    """Indexed files of one source, keyed by path — the basis for deletion detection."""
    rows = conn.execute("SELECT * FROM files WHERE source_id = ?", (source_id,))
    return {row["path"]: row for row in rows}


def upsert_file(
    conn: sqlite3.Connection,
    source_id: str,
    path: str,
    title: str | None,
    tags: list[str],
    headings: list[str],
    sha256: str,
    size: int,
    mtime: float | None,
    now: str,
) -> None:
    conn.execute(
        "INSERT INTO files (source_id, path, title, tags, headings, sha256, size, mtime, indexed_at) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) "
        "ON CONFLICT(source_id, path) DO UPDATE SET "
        "title = excluded.title, tags = excluded.tags, headings = excluded.headings, "
        "sha256 = excluded.sha256, size = excluded.size, mtime = excluded.mtime, "
        "indexed_at = excluded.indexed_at",
        (
            source_id,
            path,
            title,
            json.dumps(tags, ensure_ascii=False),
            json.dumps(headings, ensure_ascii=False),
            sha256,
            size,
            mtime,
            now,
        ),
    )


def touch_file_stat(conn: sqlite3.Connection, source_id: str, path: str, size: int, mtime: float | None) -> None:
    """Refresh the cheap change-detection stats after a sha256 match (content unchanged)."""
    conn.execute(
        "UPDATE files SET size = ?, mtime = ? WHERE source_id = ? AND path = ?",
        (size, mtime, source_id, path),
    )


def delete_file(conn: sqlite3.Connection, source_id: str, path: str) -> None:
    """Drop a file and everything derived from it, vec0 rows included."""
    _delete_vectors_for_file(conn, source_id, path)
    conn.execute("DELETE FROM files WHERE source_id = ? AND path = ?", (source_id, path))


def list_toc_files(conn: sqlite3.Connection, source_id: str | None = None, tag: str | None = None) -> list[sqlite3.Row]:
    sql = "SELECT source_id, path, title, tags FROM files"
    clauses: list[str] = []
    params: list[str] = []
    if source_id:
        clauses.append("source_id = ?")
        params.append(source_id)
    if tag:
        clauses.append("EXISTS (SELECT 1 FROM json_each(files.tags) WHERE value = ?)")
        params.append(tag)
    if clauses:
        sql += " WHERE " + " AND ".join(clauses)
    sql += " ORDER BY source_id, path"
    return list(conn.execute(sql, params))


# ---------------------------------------------------------------------------
# chunks
# ---------------------------------------------------------------------------


def _delete_vectors_for_file(conn: sqlite3.Connection, source_id: str, path: str) -> None:
    ids = [
        row["id"] for row in conn.execute("SELECT id FROM chunks WHERE source_id = ? AND path = ?", (source_id, path))
    ]
    conn.executemany("DELETE FROM vec_chunks WHERE rowid = ?", [(i,) for i in ids])


def replace_chunks(conn: sqlite3.Connection, source_id: str, path: str, chunks: Sequence[tuple[str, str]]) -> None:
    """Swap a file's chunks for a freshly built set, as (breadcrumb, text) in order.

    New chunks land with `embedded_at IS NULL`; the embed pass picks them up, so an
    unreachable Ollama degrades to FTS-only instead of failing the sync.
    """
    _delete_vectors_for_file(conn, source_id, path)
    conn.execute("DELETE FROM chunks WHERE source_id = ? AND path = ?", (source_id, path))
    conn.executemany(
        "INSERT INTO chunks (source_id, path, chunk_idx, breadcrumb, text) VALUES (?, ?, ?, ?, ?)",
        [(source_id, path, idx, breadcrumb, text) for idx, (breadcrumb, text) in enumerate(chunks)],
    )


def pending_chunks(conn: sqlite3.Connection, limit: int) -> list[sqlite3.Row]:
    return list(
        conn.execute(
            "SELECT id, breadcrumb, text FROM chunks WHERE embedded_at IS NULL ORDER BY id LIMIT ?",
            (limit,),
        )
    )


def has_pending(conn: sqlite3.Connection) -> bool:
    return conn.execute("SELECT 1 FROM chunks WHERE embedded_at IS NULL LIMIT 1").fetchone() is not None


def store_embedding(conn: sqlite3.Connection, chunk_id: int, vector: Sequence[float], now: str) -> None:
    conn.execute("DELETE FROM vec_chunks WHERE rowid = ?", (chunk_id,))
    conn.execute(
        "INSERT INTO vec_chunks (rowid, embedding) VALUES (?, ?)",
        (chunk_id, serialize_float32(list(vector))),
    )
    conn.execute("UPDATE chunks SET embedded_at = ? WHERE id = ?", (now, chunk_id))


def reset_index(conn: sqlite3.Connection) -> None:
    """Drop every indexed artifact, keeping the source rows. Used by `--full`."""
    ids = [row["id"] for row in conn.execute("SELECT id FROM chunks")]
    conn.executemany("DELETE FROM vec_chunks WHERE rowid = ?", [(i,) for i in ids])
    conn.execute("DELETE FROM files")
    conn.execute("UPDATE sources SET indexed_rev = NULL")


def index_stats(conn: sqlite3.Connection) -> dict[str, int]:
    return {
        "files": conn.execute("SELECT count(*) AS n FROM files").fetchone()["n"],
        "chunks": conn.execute("SELECT count(*) AS n FROM chunks").fetchone()["n"],
        "vectors": conn.execute("SELECT count(*) AS n FROM vec_chunks").fetchone()["n"],
        "pending": conn.execute("SELECT count(*) AS n FROM chunks WHERE embedded_at IS NULL").fetchone()["n"],
    }


def orphan_vector_ids(conn: sqlite3.Connection) -> list[int]:
    """Vector rowids with no surviving chunk — must always be empty (regression guard)."""
    rows = conn.execute("SELECT rowid AS rid FROM vec_chunks WHERE rowid NOT IN (SELECT id FROM chunks)")
    return [row["rid"] for row in rows]


# ---------------------------------------------------------------------------
# retrieval
# ---------------------------------------------------------------------------


def bm25_ranked_ids(conn: sqlite3.Connection, match_expr: str, limit: int) -> list[int]:
    rows = conn.execute(
        "SELECT rowid AS rid FROM chunks_fts WHERE chunks_fts MATCH ? ORDER BY rank LIMIT ?",
        (match_expr, limit),
    )
    return [row["rid"] for row in rows]


def knn_ranked_ids(conn: sqlite3.Connection, vector: Sequence[float], k: int) -> list[int]:
    rows = conn.execute(
        "SELECT rowid AS rid FROM vec_chunks WHERE embedding MATCH ? AND k = ? ORDER BY distance",
        (serialize_float32(list(vector)), k),
    )
    return [row["rid"] for row in rows]


def fetch_chunks(conn: sqlite3.Connection, ids: Sequence[int]) -> dict[int, sqlite3.Row]:
    if not ids:
        return {}
    placeholders = ",".join("?" * len(ids))
    rows = conn.execute(
        f"SELECT id, source_id, path, chunk_idx, breadcrumb, text FROM chunks WHERE id IN ({placeholders})",
        list(ids),
    )
    return {row["id"]: row for row in rows}