db.py

"""
db.py - SQLite storage layer for Local AI.

Uses Python's built-in sqlite3 module (zero extra dependencies, fully offline).
"""
import sqlite3
import json
import os
import time
import hashlib
from contextlib import contextmanager
from typing import Optional

DB_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "local_ai.db")


SCHEMA = """
CREATE TABLE IF NOT EXISTS documents (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    filename TEXT NOT NULL,
    file_hash TEXT,
    doc_type TEXT,
    raw_text TEXT,
    extracted_json TEXT,
    confidence REAL,
    processing_time REAL,
    created_at REAL,
    pipeline_mode TEXT
);
CREATE INDEX IF NOT EXISTS idx_doc_type ON documents(doc_type);
CREATE INDEX IF NOT EXISTS idx_file_hash ON documents(file_hash);
"""


@contextmanager
def get_conn():
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
        conn.commit()
    finally:
        conn.close()


def init_db():
    with get_conn() as conn:
        conn.executescript(SCHEMA)


def file_hash_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def find_by_hash(file_hash: str) -> Optional[sqlite3.Row]:
    with get_conn() as conn:
        cur = conn.execute("SELECT * FROM documents WHERE file_hash = ? LIMIT 1", (file_hash,))
        return cur.fetchone()


def insert_document(filename, file_hash, doc_type, raw_text, extracted: dict,
                     confidence: float, processing_time: float, pipeline_mode: str) -> int:
    with get_conn() as conn:
        cur = conn.execute(
            """INSERT INTO documents
               (filename, file_hash, doc_type, raw_text, extracted_json,
                confidence, processing_time, created_at, pipeline_mode)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (filename, file_hash, doc_type, raw_text, json.dumps(extracted),
             confidence, processing_time, time.time(), pipeline_mode),
        )
        return cur.lastrowid


def list_documents(limit: int = 200):
    with get_conn() as conn:
        cur = conn.execute(
            "SELECT * FROM documents ORDER BY created_at DESC LIMIT ?", (limit,)
        )
        return cur.fetchall()


def get_document(doc_id: int):
    with get_conn() as conn:
        cur = conn.execute("SELECT * FROM documents WHERE id = ?", (doc_id,))
        return cur.fetchone()


def delete_document(doc_id: int):
    with get_conn() as conn:
        conn.execute("DELETE FROM documents WHERE id = ?", (doc_id,))


def search_documents(query: str):
    q = f"%{query.lower()}%"
    with get_conn() as conn:
        cur = conn.execute(
            """SELECT * FROM documents
               WHERE lower(filename) LIKE ?
                  OR lower(raw_text) LIKE ?
                  OR lower(extracted_json) LIKE ?
               ORDER BY created_at DESC""",
            (q, q, q),
        )
        return cur.fetchall()


def stats():
    with get_conn() as conn:
        total = conn.execute("SELECT COUNT(*) c FROM documents").fetchone()["c"]
        avg_conf = conn.execute("SELECT AVG(confidence) c FROM documents").fetchone()["c"] or 0
        avg_time = conn.execute("SELECT AVG(processing_time) c FROM documents").fetchone()["c"] or 0
        by_type = conn.execute(
            "SELECT doc_type, COUNT(*) c FROM documents GROUP BY doc_type"
        ).fetchall()
        return {
            "total_documents": total,
            "avg_confidence": round(avg_conf, 3),
            "avg_processing_time": round(avg_time, 3),
            "by_type": {r["doc_type"]: r["c"] for r in by_type},
        }