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

Schema, connection factory and the sqlite-vec extension load.

`chunks` is the canonical retrieval unit: `chunks_fts` gives it a BM25 rank and
`vec_chunks` a KNN rank, both keyed on `chunks.id`, so RRF merges two rankings of
the *same* set.
"""

from __future__ import annotations

import sqlite3
from pathlib import Path

import sqlite_vec  # ty: ignore[unresolved-import]

# Compiled into the vec0 table definition. `wiki_embed` guards config.dims against it —
# a mismatch means the index was built for a different model.
EMBEDDING_DIMS = 1024

SCHEMA = f"""
PRAGMA journal_mode = WAL;
PRAGMA foreign_keys = ON;

CREATE TABLE IF NOT EXISTS meta (
    key   TEXT PRIMARY KEY,
    value TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS sources (
    source_id    TEXT PRIMARY KEY,
    kind         TEXT NOT NULL CHECK(kind IN ('git', 'workspace')),
    indexed_rev  TEXT,
    last_sync_at TEXT
);

CREATE TABLE IF NOT EXISTS files (
    source_id  TEXT NOT NULL REFERENCES sources(source_id),
    path       TEXT NOT NULL,
    title      TEXT,
    tags       TEXT,
    headings   TEXT,
    sha256     TEXT NOT NULL,
    size       INTEGER NOT NULL,
    mtime      REAL,
    indexed_at TEXT NOT NULL,
    PRIMARY KEY (source_id, path)
);

CREATE TABLE IF NOT EXISTS chunks (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    source_id   TEXT NOT NULL,
    path        TEXT NOT NULL,
    chunk_idx   INTEGER NOT NULL,
    breadcrumb  TEXT NOT NULL,
    text        TEXT NOT NULL,
    embedded_at TEXT,
    UNIQUE (source_id, path, chunk_idx),
    FOREIGN KEY (source_id, path) REFERENCES files(source_id, path) ON DELETE CASCADE
);

CREATE INDEX IF NOT EXISTS idx_chunks_pending ON chunks(id) WHERE embedded_at IS NULL;

CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
    breadcrumb, text,
    content='chunks', content_rowid='id',
    tokenize='unicode61 remove_diacritics 2'
);

CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
    INSERT INTO chunks_fts(rowid, breadcrumb, text) VALUES (new.id, new.breadcrumb, new.text);
END;

CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
    INSERT INTO chunks_fts(chunks_fts, rowid, breadcrumb, text)
        VALUES('delete', old.id, old.breadcrumb, old.text);
END;

-- The WHEN guard keeps `UPDATE chunks SET embedded_at` (sync step 4b) from rewriting
-- an FTS row whose indexed text did not change.
CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks
WHEN old.text IS NOT new.text OR old.breadcrumb IS NOT new.breadcrumb
BEGIN
    INSERT INTO chunks_fts(chunks_fts, rowid, breadcrumb, text)
        VALUES('delete', old.id, old.breadcrumb, old.text);
    INSERT INTO chunks_fts(rowid, breadcrumb, text) VALUES (new.id, new.breadcrumb, new.text);
END;

CREATE VIRTUAL TABLE IF NOT EXISTS vec_chunks USING vec0(
    embedding float[{EMBEDDING_DIMS}] distance_metric=cosine
);
"""


def get_db(path: Path) -> sqlite3.Connection:
    """Return an autocommit connection with sqlite-vec loaded and foreign keys on."""
    conn = sqlite3.connect(path, isolation_level=None)
    conn.enable_load_extension(True)
    sqlite_vec.load(conn)
    conn.enable_load_extension(False)
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA foreign_keys = ON")
    conn.row_factory = sqlite3.Row
    _migrate(conn)
    return conn


def _migrate(conn: sqlite3.Connection) -> None:
    """Idempotently bring an existing DB up to the current schema.

    init_db only runs the full SCHEMA on a missing file, so live DBs never see
    later additions. Each step must be a no-op once applied.
    """
    # No migrations yet — schema_version 1 is the initial shape.


def init_db(path: Path) -> None:
    """Create tables, indexes and triggers if they don't exist."""
    path.parent.mkdir(parents=True, exist_ok=True)
    conn = get_db(path)
    try:
        conn.executescript(SCHEMA)
    finally:
        conn.close()