app/sqlite_store.py

"""SQLite FTS5 fee-code store for the no-OpenSearch pilot backend."""
from __future__ import annotations

import json
import logging
import re
import sqlite3
from pathlib import Path

from .config import settings

logger = logging.getLogger(__name__)

_FTS_SAFE = re.compile(r"[^A-Za-z0-9]+")


def db_path() -> Path:
    return Path(settings.sqlite_index_path)


def connect() -> sqlite3.Connection:
    path = db_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(path), check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def ensure_schema(conn: sqlite3.Connection | None = None) -> None:
    own = conn is None
    conn = conn or connect()
    try:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS fee_codes (
              billing_code TEXT PRIMARY KEY,
              description_text TEXT NOT NULL DEFAULT '',
              rules_and_constraints TEXT NOT NULL DEFAULT '',
              parent_section TEXT NOT NULL DEFAULT '',
              base_fee_cad REAL NOT NULL DEFAULT 0,
              reference TEXT,
              differentiators_json TEXT,
              fee_components_json TEXT,
              effective_date TEXT,
              termination_date TEXT,
              in_current_schedule INTEGER NOT NULL DEFAULT 0,
              has_description INTEGER NOT NULL DEFAULT 0
            );

            CREATE VIRTUAL TABLE IF NOT EXISTS fee_codes_fts USING fts5(
              billing_code,
              description_text,
              rules_and_constraints,
              parent_section,
              tokenize = 'porter'
            );
            """
        )
        conn.commit()
    finally:
        if own:
            conn.close()


def _row_to_doc(row: sqlite3.Row | dict) -> dict:
    data = dict(row)
    diffs = None
    fees = None
    if data.get("differentiators_json"):
        try:
            diffs = json.loads(data["differentiators_json"])
        except json.JSONDecodeError:
            diffs = None
    if data.get("fee_components_json"):
        try:
            fees = json.loads(data["fee_components_json"])
        except json.JSONDecodeError:
            fees = None
    return {
        "billing_code": data["billing_code"],
        "description_text": data.get("description_text") or "",
        "rules_and_constraints": data.get("rules_and_constraints") or "",
        "parent_section": data.get("parent_section") or "",
        "base_fee_cad": float(data.get("base_fee_cad") or 0),
        "reference": data.get("reference"),
        "differentiators": diffs,
        "fee_components": fees,
        "effective_date": data.get("effective_date"),
        "termination_date": data.get("termination_date"),
        "in_current_schedule": bool(data.get("in_current_schedule")),
        "has_description": bool(data.get("has_description")),
        "score": data.get("score"),
    }


def replace_all(docs: list[dict]) -> int:
    """Rebuild the SQLite index from schedule documents."""
    conn = connect()
    try:
        conn.execute("DROP TABLE IF EXISTS fee_codes_fts")
        conn.execute("DROP TABLE IF EXISTS fee_codes")
        ensure_schema(conn)
        rows = []
        fts_rows = []
        for doc in docs:
            rows.append(
                (
                    doc["billing_code"],
                    doc.get("description_text") or "",
                    doc.get("rules_and_constraints") or "",
                    doc.get("parent_section") or "",
                    float(doc.get("base_fee_cad") or 0),
                    doc.get("reference"),
                    json.dumps(doc.get("differentiators"))
                    if doc.get("differentiators")
                    else None,
                    json.dumps(doc.get("fee_components"))
                    if doc.get("fee_components")
                    else None,
                    doc.get("effective_date"),
                    doc.get("termination_date"),
                    1 if doc.get("in_current_schedule") else 0,
                    1 if doc.get("has_description") else 0,
                )
            )
            fts_rows.append(
                (
                    doc["billing_code"],
                    doc.get("description_text") or "",
                    doc.get("rules_and_constraints") or "",
                    doc.get("parent_section") or "",
                )
            )
        conn.executemany(
            """
            INSERT INTO fee_codes (
              billing_code, description_text, rules_and_constraints, parent_section,
              base_fee_cad, reference, differentiators_json, fee_components_json,
              effective_date, termination_date, in_current_schedule, has_description
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            rows,
        )
        conn.executemany(
            """
            INSERT INTO fee_codes_fts (
              billing_code, description_text, rules_and_constraints, parent_section
            ) VALUES (?, ?, ?, ?)
            """,
            fts_rows,
        )
        conn.commit()
        n = conn.execute("SELECT COUNT(*) FROM fee_codes").fetchone()[0]
        logger.info("SQLite fee index rebuilt with %d codes at %s", n, db_path())
        return int(n)
    finally:
        conn.close()


def count_codes() -> int:
    if not db_path().exists():
        return 0
    conn = connect()
    try:
        ensure_schema(conn)
        return int(conn.execute("SELECT COUNT(*) FROM fee_codes").fetchone()[0])
    finally:
        conn.close()


def get_code(code: str) -> dict | None:
    if not code:
        return None
    conn = connect()
    try:
        ensure_schema(conn)
        row = conn.execute(
            "SELECT *, NULL AS score FROM fee_codes WHERE billing_code = ?",
            (code.upper(),),
        ).fetchone()
        return _row_to_doc(row) if row else None
    finally:
        conn.close()


def all_current_docs() -> list[dict]:
    conn = connect()
    try:
        ensure_schema(conn)
        rows = conn.execute(
            "SELECT *, NULL AS score FROM fee_codes WHERE in_current_schedule = 1"
        ).fetchall()
        return [_row_to_doc(r) for r in rows]
    finally:
        conn.close()


def _fts_query(text: str) -> str:
    tokens = [t for t in _FTS_SAFE.split(text.upper()) if len(t) >= 2]
    tokens = tokens[:24]
    if not tokens:
        return "code OR visit OR assessment"
    return " OR ".join(tokens)


def fts_search(query_text: str, *, top_k: int = 40) -> list[dict]:
    """Lexical search over the fee schedule (FTS5 BM25 rank)."""
    conn = connect()
    try:
        ensure_schema(conn)
        q = _fts_query(query_text)
        rows = conn.execute(
            """
            SELECT c.*, bm25(fee_codes_fts) AS score
            FROM fee_codes_fts
            JOIN fee_codes c ON c.billing_code = fee_codes_fts.billing_code
            WHERE fee_codes_fts MATCH ?
              AND c.in_current_schedule = 1
            ORDER BY score
            LIMIT ?
            """,
            (q, top_k),
        ).fetchall()
        docs = []
        for r in rows:
            doc = _row_to_doc(r)
            raw = doc.get("score")
            doc["score"] = float(-raw) if raw is not None else 0.0
            docs.append(doc)
        return docs
    except sqlite3.OperationalError as exc:
        logger.warning("FTS search failed (%s); returning empty", exc)
        return []
    finally:
        conn.close()