database.py

"""
database.py
-----------
SQLite database for Radiology AI Report System.
Stores all generated reports with patient info, findings, and metadata.

Tables:
  patients — patient demographics
  reports  — generated radiology reports
  images   — uploaded image metadata
"""

import json
import logging
import sqlite3
from datetime import datetime
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

DB_PATH = Path("data/radiology.db")


def get_connection() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def init_db():
    """Create all tables if they don't exist."""
    conn = get_connection()
    try:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS patients (
                id              INTEGER PRIMARY KEY AUTOINCREMENT,
                patient_name    TEXT    NOT NULL,
                age             INTEGER NOT NULL,
                gender          TEXT    NOT NULL,
                created_at      TEXT    DEFAULT (datetime('now')),
                updated_at      TEXT    DEFAULT (datetime('now'))
            );

            CREATE TABLE IF NOT EXISTS reports (
                id                   INTEGER PRIMARY KEY AUTOINCREMENT,
                report_uuid          TEXT    NOT NULL UNIQUE,
                patient_id           INTEGER NOT NULL REFERENCES patients(id),
                patient_name         TEXT    NOT NULL,
                age                  INTEGER NOT NULL,
                gender               TEXT    NOT NULL,
                symptoms             TEXT,
                clinical_history     TEXT,
                modality             TEXT,
                body_region          TEXT,
                view                 TEXT,
                examination_details  TEXT,
                clinical_information TEXT,
                findings             TEXT,
                impression           TEXT,
                recommendations      TEXT,
                confidence_score     REAL,
                generation_time_sec  REAL,
                image_caption        TEXT,
                validation_json      TEXT,
                medical_coding_json  TEXT,
                created_at           TEXT    DEFAULT (datetime('now'))
            );

            CREATE INDEX IF NOT EXISTS idx_reports_uuid
                ON reports(report_uuid);
            CREATE INDEX IF NOT EXISTS idx_reports_patient
                ON reports(patient_id);
            CREATE INDEX IF NOT EXISTS idx_reports_created
                ON reports(created_at);
            CREATE INDEX IF NOT EXISTS idx_patients_name
                ON patients(patient_name);
        """)
        conn.commit()
        logger.info(f"Database initialised at {DB_PATH}")
    finally:
        conn.close()


# ═══════════════════════════════════════════════════════════════════════════
# Patient operations
# ═══════════════════════════════════════════════════════════════════════════

def create_or_get_patient(name: str, age: int, gender: str) -> int:
    """
    Creates a new patient record or returns existing ID
    if same name+age+gender exists.
    """
    conn = get_connection()
    try:
        row = conn.execute(
            "SELECT id FROM patients WHERE patient_name=? AND age=? AND gender=?",
            (name.strip(), age, gender)
        ).fetchone()
        if row:
            return row["id"]
        cur = conn.execute(
            "INSERT INTO patients (patient_name, age, gender) VALUES (?,?,?)",
            (name.strip(), age, gender)
        )
        conn.commit()
        return cur.lastrowid
    finally:
        conn.close()


def search_patients(query: str) -> list:
    """Search patients by name."""
    conn = get_connection()
    try:
        rows = conn.execute(
            "SELECT DISTINCT patient_name, age, gender, COUNT(*) as report_count "
            "FROM reports WHERE patient_name LIKE ? "
            "GROUP BY patient_name, age, gender "
            "ORDER BY created_at DESC LIMIT 20",
            (f"%{query}%",)
        ).fetchall()
        return [dict(r) for r in rows]
    finally:
        conn.close()


# ═══════════════════════════════════════════════════════════════════════════
# Report operations
# ═══════════════════════════════════════════════════════════════════════════

def save_report(payload: dict, patient_name: str) -> int:
    """Save a complete report payload to the database. Returns row ID."""
    patient = payload.get("patient", {})
    imaging = payload.get("imaging", {})
    report  = payload.get("report",  {})
    valid   = payload.get("validation", {})
    coding  = payload.get("medical_coding", {})

    age    = patient.get("age", 0)
    gender = patient.get("gender", "Unknown")

    patient_id = create_or_get_patient(patient_name, age, gender)

    conn = get_connection()
    try:
        cur = conn.execute("""
            INSERT INTO reports (
                report_uuid, patient_id, patient_name, age, gender,
                symptoms, clinical_history,
                modality, body_region, view,
                examination_details, clinical_information,
                findings, impression, recommendations,
                confidence_score, generation_time_sec, image_caption,
                validation_json, medical_coding_json
            ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
        """, (
            payload.get("report_id", ""),
            patient_id,
            patient_name.strip(),
            age,
            gender,
            patient.get("symptoms", ""),
            patient.get("clinical_history", ""),
            imaging.get("modality", ""),
            imaging.get("body_region", ""),
            imaging.get("view", ""),
            report.get("examination_details", ""),
            report.get("clinical_information", ""),
            report.get("findings", ""),
            report.get("impression", ""),
            report.get("recommendations", ""),
            report.get("confidence_score", 0.0),
            report.get("generation_time_seconds", 0.0),
            payload.get("image_caption", ""),
            json.dumps(valid),
            json.dumps(coding),
        ))
        conn.commit()
        logger.info(f"Report saved — DB row: {cur.lastrowid} | patient: {patient_name}")
        return cur.lastrowid
    finally:
        conn.close()


def get_all_reports(limit: int = 50, offset: int = 0) -> list:
    """Get all reports ordered by newest first."""
    conn = get_connection()
    try:
        rows = conn.execute("""
            SELECT report_uuid, patient_name, age, gender,
                   modality, body_region, confidence_score,
                   created_at
            FROM reports
            ORDER BY created_at DESC
            LIMIT ? OFFSET ?
        """, (limit, offset)).fetchall()
        return [dict(r) for r in rows]
    finally:
        conn.close()


def get_report_by_uuid(uuid: str) -> Optional[dict]:
    """Get full report by UUID."""
    conn = get_connection()
    try:
        row = conn.execute(
            "SELECT * FROM reports WHERE report_uuid=?", (uuid,)
        ).fetchone()
        if not row:
            return None
        result = dict(row)
        # Parse JSON fields
        for field in ("validation_json", "medical_coding_json"):
            try:
                result[field] = json.loads(result[field] or "{}")
            except Exception:
                result[field] = {}
        return result
    finally:
        conn.close()


def get_reports_by_patient(patient_name: str) -> list:
    """Get all reports for a specific patient."""
    conn = get_connection()
    try:
        rows = conn.execute("""
            SELECT report_uuid, patient_name, age, gender,
                   modality, body_region, confidence_score,
                   findings, impression, created_at
            FROM reports
            WHERE patient_name LIKE ?
            ORDER BY created_at DESC
        """, (f"%{patient_name}%",)).fetchall()
        return [dict(r) for r in rows]
    finally:
        conn.close()


def get_db_stats() -> dict:
    """Get database statistics."""
    conn = get_connection()
    try:
        total_reports   = conn.execute("SELECT COUNT(*) FROM reports").fetchone()[0]
        total_patients  = conn.execute("SELECT COUNT(DISTINCT patient_name) FROM reports").fetchone()[0]
        avg_confidence  = conn.execute("SELECT AVG(confidence_score) FROM reports").fetchone()[0]
        modality_counts = conn.execute(
            "SELECT modality, COUNT(*) as cnt FROM reports GROUP BY modality ORDER BY cnt DESC"
        ).fetchall()
        return {
            "total_reports":   total_reports,
            "total_patients":  total_patients,
            "avg_confidence":  round(avg_confidence or 0, 2),
            "modality_counts": {r["modality"]: r["cnt"] for r in modality_counts},
        }
    finally:
        conn.close()


def delete_report(uuid: str) -> bool:
    """Delete a report by UUID."""
    conn = get_connection()
    try:
        conn.execute("DELETE FROM reports WHERE report_uuid=?", (uuid,))
        conn.commit()
        return True
    except Exception as e:
        logger.error(f"Delete failed: {e}")
        return False
    finally:
        conn.close()


# Auto-initialise on import
init_db()