app/core/persistent_db.py

"""
persistent_db.py
────────────────
Unified persistence layer supporting both PostgreSQL and SQLite.
If `DATABASE_URL` or `POSTGRES_URL` is set, PostgreSQL is used for durability across
restarts/deployments. Otherwise, it gracefully falls back to local SQLite.
"""

import os
import json
import sqlite3
import threading
import time
import uuid
from contextlib import contextmanager
from typing import Any, Dict, List, Optional
from utils.logger import app_logger

# ── paths ──────────────────────────────────────────────────────────────────
_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
_APP_DIR = os.path.dirname(_THIS_DIR)
DATA_DIR = "/data" if (os.path.exists("/data") and os.access("/data", os.W_OK)) else os.path.join(_APP_DIR, "data")
DB_PATH = os.path.join(DATA_DIR, "memory.db")
os.makedirs(DATA_DIR, exist_ok=True)

# ── SQLite schema ───────────────────────────────────────────────────────────
_SQLITE_DDL = """
PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;

CREATE TABLE IF NOT EXISTS chats (
    chat_id     TEXT PRIMARY KEY,
    title       TEXT,
    doc_id      TEXT,
    doc_name    TEXT,
    created_at  REAL NOT NULL,
    updated_at  REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS messages (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    chat_id     TEXT NOT NULL REFERENCES chats(chat_id) ON DELETE CASCADE,
    role        TEXT NOT NULL CHECK(role IN ('user','assistant','system')),
    content     TEXT NOT NULL,
    created_at  REAL NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_messages_chat ON messages(chat_id, created_at);

CREATE TABLE IF NOT EXISTS doc_meta (
    doc_id          TEXT PRIMARY KEY,
    filename        TEXT,
    status          TEXT,
    page_count      INTEGER DEFAULT 0,
    section_count   INTEGER DEFAULT 0,
    progress        INTEGER DEFAULT 0,
    current_stage   TEXT DEFAULT 'uploaded',
    indexing_method TEXT,
    duration        REAL,
    error           TEXT,
    content_hash    TEXT,
    file_path       TEXT,
    created_at      REAL NOT NULL,
    updated_at      REAL NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_doc_meta_hash ON doc_meta(content_hash);

CREATE TABLE IF NOT EXISTS doc_indices (
    doc_id          TEXT PRIMARY KEY REFERENCES doc_meta(doc_id) ON DELETE CASCADE,
    content_hash    TEXT,
    tree            TEXT,
    node_map        TEXT,
    flat_sections   TEXT,
    file_base64     TEXT,
    created_at      REAL NOT NULL,
    updated_at      REAL NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_doc_indices_hash ON doc_indices(content_hash);

CREATE TABLE IF NOT EXISTS chat_documents (
    chat_id     TEXT NOT NULL REFERENCES chats(chat_id) ON DELETE CASCADE,
    doc_id      TEXT NOT NULL REFERENCES doc_meta(doc_id) ON DELETE CASCADE,
    doc_name    TEXT,
    attached_at REAL NOT NULL,
    PRIMARY KEY (chat_id, doc_id)
);
CREATE INDEX IF NOT EXISTS idx_chat_docs ON chat_documents(chat_id);

CREATE TABLE IF NOT EXISTS chat_summaries (
    chat_id     TEXT PRIMARY KEY REFERENCES chats(chat_id) ON DELETE CASCADE,
    summary     TEXT NOT NULL,
    turn_count  INTEGER DEFAULT 0,
    updated_at  REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS knowledge_memory (
    doc_id            TEXT PRIMARY KEY REFERENCES doc_meta(doc_id) ON DELETE CASCADE,
    filename          TEXT NOT NULL,
    executive_summary TEXT,
    key_entities      TEXT,
    key_findings      TEXT,
    topics            TEXT,
    updated_at        REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS global_knowledge_summary (
    id                  TEXT PRIMARY KEY,
    collective_summary  TEXT NOT NULL,
    total_documents     INTEGER DEFAULT 0,
    key_insights        TEXT,
    updated_at          REAL NOT NULL
);
"""


class SQLitePersistentDB:
    """Thread-safe SQLite wrapper for long-term chat & document memory."""

    def __init__(self, db_path: str = DB_PATH) -> None:
        self._path = db_path
        self._local = threading.local()
        self._init_schema()

    # ── connection management ───────────────────────────────────────────────
    def _conn(self) -> sqlite3.Connection:
        if not getattr(self._local, "conn", None):
            conn = sqlite3.connect(self._path, check_same_thread=False)
            conn.row_factory = sqlite3.Row
            self._local.conn = conn
        return self._local.conn

    @contextmanager
    def _tx(self):
        conn = self._conn()
        try:
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise

    def _init_schema(self) -> None:
        with self._tx() as conn:
            conn.executescript(_SQLITE_DDL)
            # Ensure newly added columns exist in older DB instances
            cursor = conn.execute("PRAGMA table_info(doc_meta)")
            existing_cols = {row["name"] for row in cursor.fetchall()}
            for col, col_type, default_val in [
                ("section_count", "INTEGER", "0"),
                ("progress", "INTEGER", "0"),
                ("current_stage", "TEXT", "'uploaded'"),
                ("indexing_method", "TEXT", "NULL"),
                ("duration", "REAL", "NULL"),
                ("error", "TEXT", "NULL"),
            ]:
                if col not in existing_cols:
                    conn.execute(f"ALTER TABLE doc_meta ADD COLUMN {col} {col_type} DEFAULT {default_val}")

            # Migrate existing chats.doc_id into chat_documents table
            conn.execute("""
                INSERT OR IGNORE INTO chat_documents(chat_id, doc_id, doc_name, attached_at)
                SELECT chat_id, doc_id, doc_name, created_at FROM chats
                WHERE doc_id IS NOT NULL AND doc_id != ''
            """)

    # ── chat CRUD ───────────────────────────────────────────────────────────
    def create_chat(self, chat_id: Optional[str] = None, title: Optional[str] = None) -> str:
        cid = chat_id or str(uuid.uuid4())
        now = time.time()
        with self._tx() as conn:
            conn.execute(
                "INSERT OR IGNORE INTO chats(chat_id, title, created_at, updated_at) VALUES(?,?,?,?)",
                (cid, title, now, now),
            )
        return cid

    def get_chat(self, chat_id: str) -> Optional[Dict[str, Any]]:
        row = self._conn().execute(
            "SELECT * FROM chats WHERE chat_id=?", (chat_id,)
        ).fetchone()
        return dict(row) if row else None

    def list_chats(self) -> List[Dict[str, Any]]:
        rows = self._conn().execute(
            "SELECT * FROM chats ORDER BY updated_at DESC"
        ).fetchall()
        return [dict(r) for r in rows]

    def update_chat_title(self, chat_id: str, title: str) -> None:
        with self._tx() as conn:
            conn.execute(
                "UPDATE chats SET title=?, updated_at=? WHERE chat_id=?",
                (title, time.time(), chat_id),
            )

    def set_chat_doc(self, chat_id: str, doc_id: str, doc_name: Optional[str] = None) -> None:
        with self._tx() as conn:
            conn.execute(
                "UPDATE chats SET doc_id=?, doc_name=?, updated_at=? WHERE chat_id=?",
                (doc_id, doc_name, time.time(), chat_id),
            )

    def remove_chat_doc(self, chat_id: str) -> None:
        with self._tx() as conn:
            conn.execute(
                "UPDATE chats SET doc_id=NULL, doc_name=NULL, updated_at=? WHERE chat_id=?",
                (time.time(), chat_id),
            )

    def delete_chat(self, chat_id: str) -> None:
        with self._tx() as conn:
            conn.execute("DELETE FROM chats WHERE chat_id=?", (chat_id,))

    # ── message CRUD ────────────────────────────────────────────────────────
    def add_message(self, chat_id: str, role: str, content: str) -> None:
        now = time.time()
        with self._tx() as conn:
            conn.execute(
                "INSERT OR IGNORE INTO chats(chat_id, title, created_at, updated_at) VALUES(?, 'New research', ?, ?)",
                (chat_id, now, now),
            )
            conn.execute(
                "INSERT INTO messages(chat_id, role, content, created_at) VALUES(?,?,?,?)",
                (chat_id, role, content, now),
            )
            conn.execute(
                "UPDATE chats SET updated_at=? WHERE chat_id=?", (now, chat_id)
            )

    def get_messages(self, chat_id: str, limit: int = 0) -> List[Dict[str, str]]:
        q = "SELECT role, content FROM messages WHERE chat_id=? ORDER BY created_at ASC"
        if limit > 0:
            rows = self._conn().execute(
                "SELECT role, content FROM messages WHERE chat_id=? ORDER BY created_at DESC LIMIT ?",
                (chat_id, limit),
            ).fetchall()
            return [{"role": r["role"], "content": r["content"]} for r in reversed(rows)]
        rows = self._conn().execute(q, (chat_id,)).fetchall()
        return [{"role": r["role"], "content": r["content"]} for r in rows]

    def clear_messages(self, chat_id: str) -> None:
        with self._tx() as conn:
            conn.execute("DELETE FROM messages WHERE chat_id=?", (chat_id,))

    # ── document metadata ────────────────────────────────────────────────────
    def upsert_doc_meta(
        self,
        doc_id: str,
        filename: str,
        status: str,
        page_count: int = 0,
        section_count: int = 0,
        progress: int = 0,
        current_stage: str = "uploaded",
        indexing_method: Optional[str] = None,
        duration: Optional[float] = None,
        error: Optional[str] = None,
        content_hash: Optional[str] = None,
        file_path: Optional[str] = None,
    ) -> None:
        now = time.time()
        with self._tx() as conn:
            conn.execute(
                """
                INSERT INTO doc_meta(
                    doc_id, filename, status, page_count, section_count,
                    progress, current_stage, indexing_method, duration,
                    error, content_hash, file_path, created_at, updated_at
                )
                VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)
                ON CONFLICT(doc_id) DO UPDATE SET
                    filename=excluded.filename,
                    status=excluded.status,
                    page_count=excluded.page_count,
                    section_count=excluded.section_count,
                    progress=excluded.progress,
                    current_stage=excluded.current_stage,
                    indexing_method=COALESCE(excluded.indexing_method, doc_meta.indexing_method),
                    duration=COALESCE(excluded.duration, doc_meta.duration),
                    error=excluded.error,
                    content_hash=COALESCE(excluded.content_hash, doc_meta.content_hash),
                    file_path=COALESCE(excluded.file_path, doc_meta.file_path),
                    updated_at=excluded.updated_at
                """,
                (
                    doc_id, filename, status, page_count, section_count,
                    progress, current_stage, indexing_method, duration,
                    error, content_hash, file_path, now, now
                ),
            )

    def get_doc_meta(self, doc_id: str) -> Optional[Dict[str, Any]]:
        row = self._conn().execute(
            "SELECT * FROM doc_meta WHERE doc_id=?", (doc_id,)
        ).fetchone()
        return dict(row) if row else None

    def list_docs(self) -> List[Dict[str, Any]]:
        rows = self._conn().execute(
            "SELECT * FROM doc_meta ORDER BY created_at DESC"
        ).fetchall()
        return [dict(r) for r in rows]

    def find_doc_by_hash(self, content_hash: str) -> Optional[Dict[str, Any]]:
        row = self._conn().execute(
            "SELECT * FROM doc_meta WHERE content_hash=? ORDER BY updated_at DESC LIMIT 1",
            (content_hash,),
        ).fetchone()
        return dict(row) if row else None

    def delete_doc(self, doc_id: str) -> bool:
        with self._tx() as conn:
            cur = conn.execute("DELETE FROM doc_meta WHERE doc_id=?", (doc_id,))
            conn.execute("UPDATE chats SET doc_id=NULL, doc_name=NULL WHERE doc_id=?", (doc_id,))
            conn.execute("DELETE FROM chat_documents WHERE doc_id=?", (doc_id,))
            conn.execute("DELETE FROM knowledge_memory WHERE doc_id=?", (doc_id,))
            conn.execute("DELETE FROM doc_indices WHERE doc_id=?", (doc_id,))
            return cur.rowcount > 0

    # ── document heavy indices store ─────────────────────────────────────────
    def save_doc_index(
        self,
        doc_id: str,
        content_hash: Optional[str],
        tree: Any,
        node_map: Optional[Dict[str, Any]] = None,
        flat_sections: Optional[List[Dict[str, Any]]] = None,
        file_base64: Optional[str] = None,
    ) -> None:
        now = time.time()
        with self._tx() as conn:
            conn.execute(
                """
                INSERT INTO doc_indices(
                    doc_id, content_hash, tree, node_map, flat_sections, file_base64, created_at, updated_at
                )
                VALUES(?,?,?,?,?,?,?,?)
                ON CONFLICT(doc_id) DO UPDATE SET
                    content_hash=COALESCE(excluded.content_hash, doc_indices.content_hash),
                    tree=excluded.tree,
                    node_map=excluded.node_map,
                    flat_sections=excluded.flat_sections,
                    file_base64=COALESCE(excluded.file_base64, doc_indices.file_base64),
                    updated_at=excluded.updated_at
                """,
                (
                    doc_id,
                    content_hash,
                    json.dumps(tree) if tree is not None else None,
                    json.dumps(node_map) if node_map is not None else None,
                    json.dumps(flat_sections) if flat_sections is not None else None,
                    file_base64,
                    now,
                    now,
                ),
            )

    def get_doc_index(self, doc_id: str) -> Optional[Dict[str, Any]]:
        row = self._conn().execute(
            "SELECT * FROM doc_indices WHERE doc_id=?", (doc_id,)
        ).fetchone()
        if not row:
            return None
        res = dict(row)
        for k in ("tree", "node_map", "flat_sections"):
            if res.get(k) and isinstance(res[k], str):
                try:
                    res[k] = json.loads(res[k])
                except Exception:
                    pass
        return res

    def get_doc_index_by_hash(self, content_hash: str) -> Optional[Dict[str, Any]]:
        row = self._conn().execute(
            "SELECT * FROM doc_indices WHERE content_hash=? ORDER BY updated_at DESC LIMIT 1",
            (content_hash,),
        ).fetchone()
        if not row:
            return None
        res = dict(row)
        for k in ("tree", "node_map", "flat_sections"):
            if res.get(k) and isinstance(res[k], str):
                try:
                    res[k] = json.loads(res[k])
                except Exception:
                    pass
        return res

    def delete_doc_index(self, doc_id: str) -> None:
        with self._tx() as conn:
            conn.execute("DELETE FROM doc_indices WHERE doc_id=?", (doc_id,))

    # ── multi-doc session associations ──────────────────────────────────────
    def add_chat_document(self, chat_id: str, doc_id: str, doc_name: Optional[str] = None) -> None:
        now = time.time()
        with self._tx() as conn:
            conn.execute(
                """
                INSERT INTO chat_documents(chat_id, doc_id, doc_name, attached_at)
                VALUES(?,?,?,?)
                ON CONFLICT(chat_id, doc_id) DO UPDATE SET
                    doc_name=COALESCE(excluded.doc_name, chat_documents.doc_name),
                    attached_at=excluded.attached_at
                """,
                (chat_id, doc_id, doc_name, now),
            )
            conn.execute(
                "UPDATE chats SET doc_id=?, doc_name=?, updated_at=? WHERE chat_id=?",
                (doc_id, doc_name, now, chat_id),
            )

    def remove_chat_document(self, chat_id: str, doc_id: str) -> None:
        now = time.time()
        with self._tx() as conn:
            conn.execute("DELETE FROM chat_documents WHERE chat_id=? AND doc_id=?", (chat_id, doc_id))
            rem = conn.execute(
                "SELECT doc_id, doc_name FROM chat_documents WHERE chat_id=? ORDER BY attached_at DESC LIMIT 1",
                (chat_id,),
            ).fetchone()
            if rem:
                conn.execute(
                    "UPDATE chats SET doc_id=?, doc_name=?, updated_at=? WHERE chat_id=?",
                    (rem["doc_id"], rem["doc_name"], now, chat_id),
                )
            else:
                conn.execute(
                    "UPDATE chats SET doc_id=NULL, doc_name=NULL, updated_at=? WHERE chat_id=?",
                    (now, chat_id),
                )

    def get_chat_documents(self, chat_id: str) -> List[Dict[str, Any]]:
        rows = self._conn().execute(
            """
            SELECT cd.chat_id, cd.doc_id, cd.doc_name, cd.attached_at,
                   dm.status, dm.page_count, dm.section_count, dm.filename
            FROM chat_documents cd
            LEFT JOIN doc_meta dm ON cd.doc_id = dm.doc_id
            WHERE cd.chat_id=?
            ORDER BY cd.attached_at ASC
            """,
            (chat_id,),
        ).fetchall()
        return [dict(r) for r in rows]

    def clear_chat_documents(self, chat_id: str) -> None:
        now = time.time()
        with self._tx() as conn:
            conn.execute("DELETE FROM chat_documents WHERE chat_id=?", (chat_id,))
            conn.execute("UPDATE chats SET doc_id=NULL, doc_name=NULL, updated_at=? WHERE chat_id=?", (now, chat_id))

    # ── session long memory & rolling summaries ─────────────────────────────
    def set_chat_summary(self, chat_id: str, summary: str, turn_count: int = 0) -> None:
        now = time.time()
        with self._tx() as conn:
            conn.execute(
                """
                INSERT INTO chat_summaries(chat_id, summary, turn_count, updated_at)
                VALUES(?,?,?,?)
                ON CONFLICT(chat_id) DO UPDATE SET
                    summary=excluded.summary,
                    turn_count=excluded.turn_count,
                    updated_at=excluded.updated_at
                """,
                (chat_id, summary, turn_count, now),
            )

    def get_chat_summary(self, chat_id: str) -> Optional[Dict[str, Any]]:
        row = self._conn().execute(
            "SELECT * FROM chat_summaries WHERE chat_id=?", (chat_id,)
        ).fetchone()
        return dict(row) if row else None

    # ── global knowledge memory & document dossiers ────────────────────────
    def save_knowledge_dossier(
        self,
        doc_id: str,
        filename: str,
        executive_summary: str,
        key_entities: Optional[List[str]] = None,
        key_findings: Optional[List[str]] = None,
        topics: Optional[List[str]] = None,
    ) -> None:
        now = time.time()
        with self._tx() as conn:
            conn.execute(
                """
                INSERT INTO knowledge_memory(
                    doc_id, filename, executive_summary, key_entities, key_findings, topics, updated_at
                )
                VALUES(?,?,?,?,?,?,?)
                ON CONFLICT(doc_id) DO UPDATE SET
                    filename=excluded.filename,
                    executive_summary=excluded.executive_summary,
                    key_entities=excluded.key_entities,
                    key_findings=excluded.key_findings,
                    topics=excluded.topics,
                    updated_at=excluded.updated_at
                """,
                (
                    doc_id,
                    filename,
                    executive_summary,
                    json.dumps(key_entities or []),
                    json.dumps(key_findings or []),
                    json.dumps(topics or []),
                    now,
                ),
            )

    def get_knowledge_dossier(self, doc_id: str) -> Optional[Dict[str, Any]]:
        row = self._conn().execute(
            "SELECT * FROM knowledge_memory WHERE doc_id=?", (doc_id,)
        ).fetchone()
        if not row:
            return None
        res = dict(row)
        for k in ("key_entities", "key_findings", "topics"):
            if res.get(k) and isinstance(res[k], str):
                try:
                    res[k] = json.loads(res[k])
                except Exception:
                    res[k] = []
        return res

    def list_knowledge_dossiers(self) -> List[Dict[str, Any]]:
        rows = self._conn().execute(
            "SELECT * FROM knowledge_memory ORDER BY updated_at DESC"
        ).fetchall()
        result = []
        for r in rows:
            item = dict(r)
            for k in ("key_entities", "key_findings", "topics"):
                if item.get(k) and isinstance(item[k], str):
                    try:
                        item[k] = json.loads(item[k])
                    except Exception:
                        item[k] = []
            result.append(item)
        return result

    def delete_knowledge_dossier(self, doc_id: str) -> None:
        with self._tx() as conn:
            conn.execute("DELETE FROM knowledge_memory WHERE doc_id=?", (doc_id,))

    def save_global_knowledge_summary(
        self,
        collective_summary: str,
        total_documents: int,
        key_insights: Optional[List[str]] = None,
    ) -> None:
        now = time.time()
        gid = "global"
        with self._tx() as conn:
            conn.execute(
                """
                INSERT INTO global_knowledge_summary(
                    id, collective_summary, total_documents, key_insights, updated_at
                )
                VALUES(?,?,?,?,?)
                ON CONFLICT(id) DO UPDATE SET
                    collective_summary=excluded.collective_summary,
                    total_documents=excluded.total_documents,
                    key_insights=excluded.key_insights,
                    updated_at=excluded.updated_at
                """,
                (gid, collective_summary, total_documents, json.dumps(key_insights or []), now),
            )

    def get_global_knowledge_summary(self) -> Optional[Dict[str, Any]]:
        row = self._conn().execute(
            "SELECT * FROM global_knowledge_summary WHERE id='global'"
        ).fetchone()
        if not row:
            return None
        res = dict(row)
        if res.get("key_insights") and isinstance(res["key_insights"], str):
            try:
                res["key_insights"] = json.loads(res["key_insights"])
            except Exception:
                res["key_insights"] = []
        return res


def _migrate_sqlite_to_postgres_if_needed(pg_db: Any) -> None:
    """If local SQLite contains chats/docs and PostgreSQL is newly initialized, migrate data smoothly."""
    if not os.path.exists(DB_PATH):
        return

    # Check if Postgres already has documents or chats
    if pg_db.list_docs() or pg_db.list_chats():
        return

    sqlite_db = SQLitePersistentDB(DB_PATH)
    sqlite_chats = sqlite_db.list_chats()
    sqlite_docs = sqlite_db.list_docs()

    if not sqlite_chats and not sqlite_docs:
        return

    app_logger.info(f"[Migration] Migrating {len(sqlite_docs)} docs and {len(sqlite_chats)} chats from SQLite to PostgreSQL...")

    for d in sqlite_docs:
        try:
            pg_db.upsert_doc_meta(
                doc_id=d["doc_id"],
                filename=d.get("filename", "document.pdf"),
                status=d.get("status", "ready"),
                page_count=d.get("page_count", 0),
                section_count=d.get("section_count", 0),
                progress=d.get("progress", 100),
                current_stage=d.get("current_stage", "Ready"),
                indexing_method=d.get("indexing_method"),
                duration=d.get("duration"),
                error=d.get("error"),
                content_hash=d.get("content_hash"),
                file_path=d.get("file_path"),
            )
            # Check if index exists in SQLite
            s_idx = sqlite_db.get_doc_index(d["doc_id"])
            if s_idx:
                pg_db.save_doc_index(
                    doc_id=d["doc_id"],
                    content_hash=d.get("content_hash"),
                    tree=s_idx.get("tree"),
                    node_map=s_idx.get("node_map"),
                    flat_sections=s_idx.get("flat_sections"),
                    file_base64=s_idx.get("file_base64"),
                )
        except Exception as dex:
            app_logger.warning(f"[Migration] Error migrating doc {d.get('doc_id')}: {dex}")

    for c in sqlite_chats:
        try:
            cid = c["chat_id"]
            pg_db.create_chat(chat_id=cid, title=c.get("title"))
            if c.get("doc_id"):
                pg_db.set_chat_doc(cid, c["doc_id"], c.get("doc_name"))
            # Migrate messages
            msgs = sqlite_db.get_messages(cid)
            for m in msgs:
                pg_db.add_message(cid, m["role"], m["content"])
            # Migrate attached documents
            att = sqlite_db.get_chat_documents(cid)
            for a in att:
                pg_db.add_chat_document(cid, a["doc_id"], a.get("doc_name"))
        except Exception as cex:
            app_logger.warning(f"[Migration] Error migrating chat {c.get('chat_id')}: {cex}")

    app_logger.info("[Migration] SQLite to PostgreSQL migration completed successfully.")


DB_STATUS = {
    "provider": "SQLite",
    "connected": True,
    "error": None,
    "url_configured": False,
}


def create_persistent_db() -> Any:
    """
    Factory function for persistent database.
    Checks environment for DATABASE_URL / POSTGRES_URL.
    If provided and accessible, uses PostgresDB. Otherwise falls back to SQLite.
    """
    global DB_STATUS
    pg_url = os.getenv("DATABASE_URL") or os.getenv("POSTGRES_URL") or os.getenv("POSTGRESQL_URL")
    if pg_url:
        pg_url = pg_url.strip().strip("'").strip('"')
        DB_STATUS["url_configured"] = True
        try:
            from core.postgres_db import PostgresDB
            pg_db = PostgresDB(pg_url)
            safe_url = pg_url.split("@")[-1] if "@" in pg_url else "local"
            app_logger.info(f"[PersistentDB] Connected to PostgreSQL: {safe_url}")
            DB_STATUS["provider"] = "PostgreSQL"
            DB_STATUS["connected"] = True
            DB_STATUS["error"] = None
            try:
                _migrate_sqlite_to_postgres_if_needed(pg_db)
            except Exception as mex:
                app_logger.warning(f"[PersistentDB] Auto-migration skipped: {mex}")
            return pg_db
        except Exception as e:
            err_msg = str(e)
            # Sanitize password from error if any
            if "@" in err_msg and ":" in err_msg:
                err_msg = err_msg.split("@")[-1]
            DB_STATUS["provider"] = "SQLite (fallback)"
            DB_STATUS["connected"] = False
            DB_STATUS["error"] = err_msg
            app_logger.error(f"[PersistentDB] Failed to connect to PostgreSQL ({e}). Falling back to SQLite.")

    app_logger.info(f"[PersistentDB] Using SQLite persistence at {DB_PATH}")
    DB_STATUS["provider"] = "SQLite"
    DB_STATUS["connected"] = True
    return SQLitePersistentDB(DB_PATH)


# Global singleton instance
persistent_db = create_persistent_db()