backend/database.py

import sqlite3
import struct
import json
from pathlib import Path
from contextlib import contextmanager
from typing import Optional

DB_PATH = Path(__file__).parent.parent / "data" / "bookmarks.db"

SCHEMA = """
CREATE TABLE IF NOT EXISTS bookmarks (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    url TEXT NOT NULL UNIQUE,
    title TEXT NOT NULL,
    description TEXT DEFAULT '',
    context TEXT DEFAULT '',
    folder_path TEXT DEFAULT '',
    domain TEXT DEFAULT '',
    favicon_url TEXT DEFAULT '',
    auto_tags TEXT DEFAULT '[]',
    user_tags TEXT DEFAULT '[]',
    content_snippet TEXT DEFAULT '',
    embedding_text TEXT DEFAULT '',
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    last_visited_at TIMESTAMP,
    visit_count INTEGER DEFAULT 0,
    is_dead BOOLEAN DEFAULT FALSE
);

CREATE INDEX IF NOT EXISTS idx_bookmarks_domain ON bookmarks(domain);
CREATE INDEX IF NOT EXISTS idx_bookmarks_created ON bookmarks(created_at);
"""

FTS_SCHEMA = """
CREATE VIRTUAL TABLE IF NOT EXISTS bookmarks_fts USING fts5(
    title, description, context, auto_tags, user_tags, content_snippet,
    content='bookmarks', content_rowid='id'
);
"""

VEC_SCHEMA = """
CREATE VIRTUAL TABLE IF NOT EXISTS vec_bookmarks USING vec0(
    id INTEGER PRIMARY KEY,
    embedding FLOAT[384]
);
"""

FTS_TRIGGERS = """
CREATE TRIGGER IF NOT EXISTS bookmarks_ai AFTER INSERT ON bookmarks BEGIN
    INSERT INTO bookmarks_fts(rowid, title, description, context, auto_tags, user_tags, content_snippet)
    VALUES (new.id, new.title, new.description, new.context, new.auto_tags, new.user_tags, new.content_snippet);
END;

CREATE TRIGGER IF NOT EXISTS bookmarks_ad AFTER DELETE ON bookmarks BEGIN
    INSERT INTO bookmarks_fts(bookmarks_fts, rowid, title, description, context, auto_tags, user_tags, content_snippet)
    VALUES ('delete', old.id, old.title, old.description, old.context, old.auto_tags, old.user_tags, old.content_snippet);
END;

CREATE TRIGGER IF NOT EXISTS bookmarks_au AFTER UPDATE ON bookmarks BEGIN
    INSERT INTO bookmarks_fts(bookmarks_fts, rowid, title, description, context, auto_tags, user_tags, content_snippet)
    VALUES ('delete', old.id, old.title, old.description, old.context, old.auto_tags, old.user_tags, old.content_snippet);
    INSERT INTO bookmarks_fts(rowid, title, description, context, auto_tags, user_tags, content_snippet)
    VALUES (new.id, new.title, new.description, new.context, new.auto_tags, new.user_tags, new.content_snippet);
END;
"""


def _row_to_dict(row: sqlite3.Row, columns: Optional[list[str]] = None) -> dict:
    """Convert a sqlite3.Row to a plain dict."""
    return dict(row)


@contextmanager
def get_db(db_path: Optional[str] = None):
    """Context manager yielding a sqlite3 Connection with sqlite_vec loaded and schema ensured."""
    path = Path(db_path) if db_path else DB_PATH
    path.parent.mkdir(parents=True, exist_ok=True)

    conn = sqlite3.connect(str(path))
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")

    conn.enable_load_extension(True)
    try:
        import sqlite_vec
        sqlite_vec.load(conn)
    except ImportError:
        raise RuntimeError("sqlite_vec is required for vector operations. Install it with 'pip install sqlite-vec'")
    conn.enable_load_extension(False)

    _ensure_schema(conn)

    try:
        yield conn
    finally:
        conn.close()


def _ensure_schema(conn: sqlite3.Connection) -> None:
    """Create tables and triggers if they do not exist."""
    for statement in SCHEMA.strip().split(";"):
        stmt = statement.strip()
        if stmt:
            conn.execute(stmt)

    for statement in FTS_SCHEMA.strip().split(";"):
        stmt = statement.strip()
        if stmt:
            conn.execute(stmt)

    for statement in VEC_SCHEMA.strip().split(";"):
        stmt = statement.strip()
        if stmt:
            conn.execute(stmt)

    for statement in FTS_TRIGGERS.strip().split(";"):
        stmt = statement.strip()
        if stmt:
            conn.execute(stmt)

    conn.commit()


def init_db(db_path: Optional[str] = None) -> None:
    """Explicitly initialise the database (creates all tables/triggers)."""
    with get_db(db_path) as conn:
        pass  # schema is ensured inside get_db


def add_bookmark(
    conn: sqlite3.Connection,
    url: str,
    title: str,
    description: str,
    context: str,
    folder_path: str,
    domain: str,
    auto_tags: list[str],
    user_tags: list[str],
    content_snippet: str,
    embedding_text: str,
    embedding_blob: bytes,
) -> int:
    """Insert a bookmark and its vector. Returns the new row id."""
    cursor = conn.execute(
        """
        INSERT INTO bookmarks
            (url, title, description, context, folder_path, domain,
             auto_tags, user_tags, content_snippet, embedding_text)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            url,
            title,
            description,
            context,
            folder_path,
            domain,
            json.dumps(auto_tags, ensure_ascii=False),
            json.dumps(user_tags, ensure_ascii=False),
            content_snippet,
            embedding_text,
        ),
    )
    new_id = cursor.lastrowid

    conn.execute(
        "INSERT INTO vec_bookmarks(id, embedding) VALUES (?, ?)",
        (new_id, embedding_blob),
    )

    return new_id


def get_bookmark(conn: sqlite3.Connection, id: int) -> Optional[dict]:
    """Fetch a single bookmark by id, or None if not found."""
    row = conn.execute("SELECT * FROM bookmarks WHERE id = ?", (id,)).fetchone()
    if row is None:
        return None
    bm = dict(row)
    bm["auto_tags"] = json.loads(bm.get("auto_tags") or "[]")
    bm["user_tags"] = json.loads(bm.get("user_tags") or "[]")
    return bm


def update_bookmark(conn: sqlite3.Connection, id: int, **kwargs) -> None:
    """Update specified fields on a bookmark. updated_at is set automatically."""
    if not kwargs:
        return

    # Define allowed fields for update
    ALLOWED_FIELDS = {
        "url", "title", "description", "context", "folder_path", "domain",
        "auto_tags", "user_tags", "content_snippet", "embedding_text",
        "created_at", "updated_at", "last_visited_at", "visit_count", "is_dead"
    }

    # Validate that all provided fields are allowed
    for key in kwargs:
        if key not in ALLOWED_FIELDS:
            raise ValueError(f"Unknown field '{key}' for bookmark update. Allowed fields: {ALLOWED_FIELDS}")

    # Serialise tag lists if provided
    for key in ("auto_tags", "user_tags"):
        if key in kwargs and isinstance(kwargs[key], list):
            kwargs[key] = json.dumps(kwargs[key], ensure_ascii=False)

    kwargs["updated_at"] = "CURRENT_TIMESTAMP"

    set_clauses = []
    values = []
    for key, value in kwargs.items():
        if key == "updated_at":
            set_clauses.append("updated_at = CURRENT_TIMESTAMP")
        else:
            set_clauses.append(f"{key} = ?")
            values.append(value)

    values.append(id)
    conn.execute(
        f"UPDATE bookmarks SET {', '.join(set_clauses)} WHERE id = ?",
        values,
    )


def delete_bookmark(conn: sqlite3.Connection, id: int) -> None:
    """Delete a bookmark and its vector entry."""
    conn.execute("DELETE FROM bookmarks WHERE id = ?", (id,))
    conn.execute("DELETE FROM vec_bookmarks WHERE id = ?", (id,))


def list_bookmarks(
    conn: sqlite3.Connection, page: int = 1, per_page: int = 20
) -> tuple[list[dict], int]:
    """Return a page of bookmarks and the total count."""
    total = conn.execute("SELECT COUNT(*) FROM bookmarks").fetchone()[0]
    offset = (page - 1) * per_page
    rows = conn.execute(
        "SELECT * FROM bookmarks ORDER BY created_at DESC LIMIT ? OFFSET ?",
        (per_page, offset),
    ).fetchall()

    bookmarks = []
    for row in rows:
        bm = dict(row)
        bm["auto_tags"] = json.loads(bm.get("auto_tags") or "[]")
        bm["user_tags"] = json.loads(bm.get("user_tags") or "[]")
        bookmarks.append(bm)

    return bookmarks, total


def get_all_tags(conn: sqlite3.Connection) -> list[tuple[str, int]]:
    """Return all unique tags with their occurrence counts, sorted by count desc."""
    rows = conn.execute("SELECT auto_tags, user_tags FROM bookmarks").fetchall()

    counts: dict[str, int] = {}
    for row in rows:
        for col in ("auto_tags", "user_tags"):
            tags = json.loads(row[col] or "[]")
            for tag in tags:
                if tag:
                    counts[tag] = counts.get(tag, 0) + 1

    return sorted(counts.items(), key=lambda x: -x[1])


def get_stats(conn: sqlite3.Connection) -> dict:
    """Return aggregate statistics."""
    total = conn.execute("SELECT COUNT(*) FROM bookmarks").fetchone()[0]
    dead_links = conn.execute(
        "SELECT COUNT(*) FROM bookmarks WHERE is_dead = TRUE"
    ).fetchone()[0]

    top_domains_rows = conn.execute(
        """
        SELECT domain, COUNT(*) as cnt
        FROM bookmarks
        WHERE domain != ''
        GROUP BY domain
        ORDER BY cnt DESC
        LIMIT 10
        """
    ).fetchall()
    top_domains = [{"domain": r["domain"], "count": r["cnt"]} for r in top_domains_rows]

    return {
        "total_bookmarks": total,
        "unique_domains": len(top_domains),
        "dead_links": dead_links,
        "top_domains": top_domains,
    }


def search_fts(
    conn: sqlite3.Connection, query: str, limit: int = 20
) -> list[tuple[int, float]]:
    """FTS5 full-text search. Returns list of (rowid, rank)."""
    try:
        rows = conn.execute(
            """
            SELECT rowid, rank
            FROM bookmarks_fts
            WHERE bookmarks_fts MATCH ?
            ORDER BY rank
            LIMIT ?
            """,
            (query, limit),
        ).fetchall()
        return [(row[0], row[1]) for row in rows]
    except sqlite3.OperationalError:
        # Invalid FTS query (e.g. bare special chars)
        return []


def search_vec(
    conn: sqlite3.Connection, query_blob: bytes, limit: int = 20
) -> list[tuple[int, float]]:
    """sqlite-vec KNN search. Returns list of (id, distance)."""
    rows = conn.execute(
        """
        SELECT id, distance
        FROM vec_bookmarks
        WHERE embedding MATCH ?
        ORDER BY distance
        LIMIT ?
        """,
        (query_blob, limit),
    ).fetchall()
    return [(row[0], row[1]) for row in rows]