db.py

"""
db.py
Dynamic schema SQLite. One prescription upload can create MULTIPLE
records (one per visit/date). All records from the same image share
the same source_filename. Columns are created automatically based on
whatever fields are extracted from each prescription type.
"""

import sqlite3
import json
import re
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Optional

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

BASE_SCHEMA = """
CREATE TABLE IF NOT EXISTS prescriptions (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    source_filename TEXT,
    document_type   TEXT,
    hospital_name   TEXT,
    patient_name    TEXT,
    hospital_no     TEXT,
    visit_date      TEXT,
    ocr_engine      TEXT,
    created_at      TEXT
);

CREATE TABLE IF NOT EXISTS medications (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    prescription_id INTEGER NOT NULL,
    drug_name       TEXT,
    dosage          TEXT,
    frequency       TEXT,
    route           TEXT,
    FOREIGN KEY (prescription_id) REFERENCES prescriptions(id)
);
"""

RESERVED_COLS = {
    "id", "source_filename", "document_type", "hospital_name",
    "patient_name", "hospital_no", "visit_date", "ocr_engine", "created_at"
}


def _safe_col(name: str) -> str:
    name = name.strip().lower()
    name = re.sub(r"[^a-z0-9_]", "_", name)
    name = re.sub(r"_+", "_", name).strip("_")
    return name or "field"


def get_conn() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    conn = get_conn()
    conn.executescript(BASE_SCHEMA)
    conn.commit()
    conn.close()


def _existing_columns(conn: sqlite3.Connection) -> List[str]:
    return [r["name"] for r in conn.execute("PRAGMA table_info(prescriptions)").fetchall()]


def _ensure_columns(conn: sqlite3.Connection, field_names: List[str]):
    existing = set(_existing_columns(conn))
    for name in field_names:
        col = _safe_col(name)
        if col and col not in existing and col not in RESERVED_COLS:
            conn.execute(f'ALTER TABLE prescriptions ADD COLUMN "{col}" TEXT')
            existing.add(col)
    conn.commit()


def save_record(
    visit_date:     Optional[str],
    fields:         Dict[str, Any],
    medications:    List[Dict],
    meta:           Dict,
    ocr_engine:     str,
    source_filename: str,
) -> int:
    """Save one visit record. Called once per visit detected in the image."""
    conn = get_conn()

    # Dynamic fields — exclude reserved and medications key
    dynamic = {
        _safe_col(k): str(v)
        for k, v in fields.items()
        if v and _safe_col(k) not in RESERVED_COLS and k != "medications"
    }
    _ensure_columns(conn, list(dynamic.keys()))

    row = {
        "source_filename": source_filename,
        "document_type":   meta.get("document_type"),
        "hospital_name":   meta.get("hospital_name"),
        "patient_name":    meta.get("patient_name"),
        "hospital_no":     meta.get("hospital_no"),
        "visit_date":      visit_date,
        "ocr_engine":      ocr_engine,
        "created_at":      datetime.utcnow().isoformat(),
        **dynamic,
    }
    # Remove None values
    row = {k: v for k, v in row.items() if v is not None}

    cols         = ", ".join(f'"{c}"' for c in row.keys())
    placeholders = ", ".join(["?"] * len(row))
    cur = conn.execute(
        f"INSERT INTO prescriptions ({cols}) VALUES ({placeholders})",
        list(row.values()),
    )
    pid = cur.lastrowid

    for med in (medications or []):
        if not isinstance(med, dict):
            continue
        conn.execute(
            "INSERT INTO medications (prescription_id, drug_name, dosage, frequency, route) VALUES (?,?,?,?,?)",
            (pid, med.get("drug_name"), med.get("dosage"), med.get("frequency"), med.get("route")),
        )

    conn.commit()
    conn.close()
    return pid


def fetch_prescriptions() -> List[Dict]:
    conn = get_conn()
    rows = conn.execute("SELECT * FROM prescriptions ORDER BY id DESC").fetchall()
    conn.close()
    return [dict(r) for r in rows]


def fetch_medications(prescription_id: int) -> List[Dict]:
    conn = get_conn()
    rows = conn.execute(
        "SELECT drug_name, dosage, frequency, route FROM medications WHERE prescription_id = ?",
        (prescription_id,)
    ).fetchall()
    conn.close()
    return [dict(r) for r in rows]


def fetch_all_flat() -> List[Dict]:
    conn = get_conn()
    rows = conn.execute(
        """
        SELECT p.*, m.drug_name, m.dosage, m.frequency, m.route
        FROM prescriptions p
        LEFT JOIN medications m ON m.prescription_id = p.id
        ORDER BY p.id DESC
        """
    ).fetchall()
    conn.close()
    return [dict(r) for r in rows]


def delete_prescription(prescription_id: int):
    conn = get_conn()
    conn.execute("DELETE FROM medications WHERE prescription_id = ?", (prescription_id,))
    conn.execute("DELETE FROM prescriptions WHERE id = ?", (prescription_id,))
    conn.commit()
    conn.close()