#!/usr/bin/env python3
"""Offline, synthetic sqlite-vec storage + nearest-neighbor verification.

Creates only a temporary database, never opens user files or talks to a network.
Requires sqlite-vec installed in the chosen Python environment.
The 384-dimension vectors are *synthetic*, not AI-generated text embeddings.
"""
from __future__ import annotations

import argparse
import json
import math
import sqlite3
import struct
import sys
import tempfile
from pathlib import Path
from typing import Sequence

DIMS = 384


def serialize_f32(vector: Sequence[float]) -> bytes:
    if len(vector) != DIMS or not all(math.isfinite(float(x)) for x in vector):
        raise ValueError("Expected exactly 384 finite float32 dimensions")
    return struct.pack(f"<{DIMS}f", *vector)


def load_sqlite_vec(db: sqlite3.Connection) -> str:
    import sqlite_vec

    db.enable_load_extension(True)
    try:
        sqlite_vec.load(db)
    finally:
        db.enable_load_extension(False)
    return str(db.execute("SELECT vec_version()").fetchone()[0])


def new_database(db: sqlite3.Connection) -> None:
    with db:
        db.execute(
            """CREATE TABLE memory_docs (
                id INTEGER PRIMARY KEY,
                category TEXT NOT NULL,
                content TEXT NOT NULL
            )"""
        )
        db.execute(
            """CREATE VIRTUAL TABLE vec_memory USING vec0(
                doc_id INTEGER PRIMARY KEY,
                embedding float[384]
            )"""
        )


def vector(a: float, b: float) -> list[float]:
    return [a, b] + [0.0] * (DIMS - 2)


def insert_memory(
    db: sqlite3.Connection, category: str, content: str, embedding: Sequence[float]
) -> int:
    if not category.strip() or not content.strip():
        raise ValueError("Category and content must not be empty")
    packed = serialize_f32(embedding)
    with db:
        doc_id = db.execute(
            "INSERT INTO memory_docs(category, content) VALUES(?, ?)",
            (category, content),
        ).lastrowid
        db.execute(
            "INSERT INTO vec_memory(doc_id, embedding) VALUES(?, ?)",
            (doc_id, packed),
        )
    return int(doc_id)


def search(
    db: sqlite3.Connection, query: Sequence[float], limit: int = 3
) -> list[dict[str, object]]:
    if not 1 <= limit <= 20:
        raise ValueError("limit must be 1..20")
    packed = serialize_f32(query)
    rows = db.execute(
        """SELECT d.id, d.category, d.content, v.distance
           FROM vec_memory v
           JOIN memory_docs d ON v.doc_id=d.id
           WHERE v.embedding MATCH ? AND k=?
           ORDER BY v.distance ASC""",
        (packed, limit),
    ).fetchall()
    return [
        {
            "id": int(i),
            "category": category,
            "content": content,
            "l2_distance": round(float(distance), 6),
        }
        for i, category, content, distance in rows
    ]


def offline_self_test() -> dict[str, object]:
    checks: list[str] = []
    with tempfile.TemporaryDirectory(prefix="sqlite-vec-smoke-") as tmp:
        path = Path(tmp) / "synthetic.db"
        with sqlite3.connect(path) as db:
            version = load_sqlite_vec(db)
            new_database(db)
            checks.append("create_schema")
            a = insert_memory(db, "decision", "first synthetic record", vector(1, 0))
            b = insert_memory(db, "warning", "second synthetic record", vector(0, 1))
            c = insert_memory(db, "test", "third synthetic record", vector(0.6, 0.4))
            assert len({a, b, c}) == 3
            checks.append("insert_vectors")
            results = search(db, vector(0.95, 0.05), 3)
            assert [x["id"] for x in results] == [a, c, b], results
            assert results[0]["l2_distance"] < results[1]["l2_distance"]
            checks.append("l2_order_and_doc_join")
            assert len(search(db, vector(0.95, 0.05), 1)) == 1
            checks.append("top_k_limit")
        with sqlite3.connect(path) as db:
            load_sqlite_vec(db)
            assert [x["id"] for x in search(db, vector(0.95, 0.05), 3)] == [
                a, c, b
            ]
            checks.append("persist_and_reopen")
            before = db.execute("SELECT count(*) FROM memory_docs").fetchone()[0]
            try:
                with db:
                    db.execute(
                        "INSERT INTO memory_docs(category, content) VALUES(?, ?)",
                        ("test", "will roll back"),
                    )
                    raise RuntimeError("synthetic rollback")
            except RuntimeError:
                pass
            assert db.execute("SELECT count(*) FROM memory_docs").fetchone()[0] == before
            checks.append("transaction_rollback")
            for invalid in (vector(1, 0)[:-1], [float("nan")] * DIMS):
                try:
                    serialize_f32(invalid)
                except ValueError:
                    continue
                raise AssertionError("Invalid vector was accepted")
            checks.append("reject_invalid_vectors")
            for invalid_limit in (0, 21):
                try:
                    search(db, vector(1, 0), invalid_limit)
                except ValueError:
                    continue
                raise AssertionError("Invalid limit was accepted")
            checks.append("reject_unbounded_search")
            assert not db.execute(
                "SELECT count(*) FROM memory_docs WHERE content LIKE '%API_KEY%'"
            ).fetchone()[0]
            checks.append("synthetic_only")
    return {
        "result": "PASS",
        "sqlite_vec_version": version,
        "synthetic_vectors": 3,
        "dimensions": DIMS,
        "metric": "sqlite-vec vec0 default L2 (Euclidean)",
        "checks": checks,
        "network_calls": 0,
        "user_data_reads": 0,
        "user_data_writes": 0,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true", required=True)
    args = parser.parse_args()
    try:
        print(json.dumps(offline_self_test(), indent=2))
        return 0
    except (sqlite3.Error, ValueError, AssertionError, ImportError) as e:
        print(json.dumps({"result": "FAIL", "error_type": type(e).__name__}), file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())