app/core/postgres_db.py

"""
postgres_db.py
──────────────
PostgreSQL persistence layer for enterprise-grade durability.
Stores chat sessions, messages, document metadata, and heavy document index trees (JSONB).
Even if the container or app restarts or sleeps, all documents, chats, and indices remain intact.
"""

import os
import json
import time
import uuid
import base64
from contextlib import contextmanager
from typing import Any, Dict, List, Optional
import psycopg2
import psycopg2.extras
import psycopg2.pool
from utils.logger import app_logger

# Register JSON/JSONB handlers so JSON columns are transparently converted to/from dict/list
psycopg2.extras.register_default_jsonb()

_PG_DDL = """
CREATE TABLE IF NOT EXISTS chats (
    chat_id     TEXT PRIMARY KEY,
    title       TEXT,
    doc_id      TEXT,
    doc_name    TEXT,
    created_at  DOUBLE PRECISION NOT NULL,
    updated_at  DOUBLE PRECISION NOT NULL
);

CREATE TABLE IF NOT EXISTS messages (
    id          BIGSERIAL PRIMARY KEY,
    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  DOUBLE PRECISION NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_pg_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        DOUBLE PRECISION,
    error           TEXT,
    content_hash    TEXT,
    file_path       TEXT,
    created_at      DOUBLE PRECISION NOT NULL,
    updated_at      DOUBLE PRECISION NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_pg_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            JSONB,
    node_map        JSONB,
    flat_sections   JSONB,
    file_base64     TEXT,
    created_at      DOUBLE PRECISION NOT NULL,
    updated_at      DOUBLE PRECISION NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_pg_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 DOUBLE PRECISION NOT NULL,
    PRIMARY KEY (chat_id, doc_id)
);

CREATE INDEX IF NOT EXISTS idx_pg_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  DOUBLE PRECISION 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      JSONB,
    key_findings      JSONB,
    topics            JSONB,
    updated_at        DOUBLE PRECISION 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        JSONB,
    updated_at          DOUBLE PRECISION NOT NULL
);
"""


class PostgresDB:
    """Robust, thread-safe PostgreSQL database client with connection pool and auto-reconnect."""

    def __init__(self, dsn: str):
        # Normalize postgres:// to postgresql:// if needed
        if dsn.startswith("postgres://"):
            dsn = "postgresql://" + dsn[len("postgres://"):]
        self._dsn = dsn
        self._pool: Optional[psycopg2.pool.ThreadedConnectionPool] = None
        self._init_pool()
        self._init_schema()

    def _init_pool(self) -> None:
        try:
            self._pool = psycopg2.pool.ThreadedConnectionPool(
                minconn=1,
                maxconn=20,
                dsn=self._dsn,
            )
            app_logger.info("[PostgresDB] Connection pool created successfully.")
        except Exception as e:
            app_logger.error(f"[PostgresDB] Failed to initialize connection pool: {e}")
            raise

    @contextmanager
    def _conn(self):
        """Context manager that borrows a healthy connection from the pool and handles commits/rollbacks."""
        conn = None
        try:
            conn = self._pool.getconn()
            # Test connection liveness
            if conn.closed:
                conn = psycopg2.connect(self._dsn)
            else:
                try:
                    with conn.cursor() as cur:
                        cur.execute("SELECT 1;")
                except Exception:
                    # Connection is dead, recreate
                    try:
                        self._pool.putconn(conn, close=True)
                    except Exception:
                        pass
                    conn = psycopg2.connect(self._dsn)

            yield conn
            conn.commit()
        except Exception:
            if conn and not conn.closed:
                conn.rollback()
            raise
        finally:
            if conn and self._pool:
                try:
                    self._pool.putconn(conn)
                except Exception:
                    pass

    def _init_schema(self) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(_PG_DDL)
        app_logger.info("[PostgresDB] Schema initialized / verified successfully.")

    # ── 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._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO chats(chat_id, title, created_at, updated_at)
                    VALUES (%s, %s, %s, %s)
                    ON CONFLICT (chat_id) DO NOTHING
                    """,
                    (cid, title, now, now),
                )
        return cid

    def get_chat(self, chat_id: str) -> Optional[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM chats WHERE chat_id = %s", (chat_id,))
                row = cur.fetchone()
                return dict(row) if row else None

    def list_chats(self) -> List[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM chats ORDER BY updated_at DESC")
                rows = cur.fetchall()
                return [dict(r) for r in rows]

    def update_chat_title(self, chat_id: str, title: str) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE chats SET title = %s, updated_at = %s WHERE chat_id = %s",
                    (title, time.time(), chat_id),
                )

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

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

    def delete_chat(self, chat_id: str) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM chats WHERE chat_id = %s", (chat_id,))

    # ── message CRUD ────────────────────────────────────────────────────────
    def add_message(self, chat_id: str, role: str, content: str) -> None:
        now = time.time()
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO chats(chat_id, title, created_at, updated_at)
                    VALUES (%s, 'New research', %s, %s)
                    ON CONFLICT (chat_id) DO NOTHING
                    """,
                    (chat_id, now, now),
                )
                cur.execute(
                    "INSERT INTO messages(chat_id, role, content, created_at) VALUES (%s, %s, %s, %s)",
                    (chat_id, role, content, now),
                )
                cur.execute(
                    "UPDATE chats SET updated_at = %s WHERE chat_id = %s",
                    (now, chat_id),
                )

    def get_messages(self, chat_id: str, limit: int = 0) -> List[Dict[str, str]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                if limit > 0:
                    cur.execute(
                        "SELECT role, content FROM messages WHERE chat_id = %s ORDER BY created_at DESC LIMIT %s",
                        (chat_id, limit),
                    )
                    rows = cur.fetchall()
                    return [{"role": r["role"], "content": r["content"]} for r in reversed(rows)]
                cur.execute(
                    "SELECT role, content FROM messages WHERE chat_id = %s ORDER BY created_at ASC",
                    (chat_id,),
                )
                rows = cur.fetchall()
                return [{"role": r["role"], "content": r["content"]} for r in rows]

    def clear_messages(self, chat_id: str) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM messages WHERE chat_id = %s", (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._conn() as conn:
            with conn.cursor() as cur:
                cur.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 (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    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]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM doc_meta WHERE doc_id = %s", (doc_id,))
                row = cur.fetchone()
                return dict(row) if row else None

    def list_docs(self) -> List[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM doc_meta ORDER BY created_at DESC")
                rows = cur.fetchall()
                return [dict(r) for r in rows]

    def find_doc_by_hash(self, content_hash: str) -> Optional[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM doc_meta WHERE content_hash = %s ORDER BY updated_at DESC LIMIT 1",
                    (content_hash,),
                )
                row = cur.fetchone()
                return dict(row) if row else None

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

    # ── document heavy indices store (JSONB) ─────────────────────────────────
    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._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO doc_indices (
                        doc_id, content_hash, tree, node_map, flat_sections, file_base64, created_at, updated_at
                    )
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                    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,
                        psycopg2.extras.Json(tree) if tree is not None else None,
                        psycopg2.extras.Json(node_map) if node_map is not None else None,
                        psycopg2.extras.Json(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]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM doc_indices WHERE doc_id = %s", (doc_id,))
                row = cur.fetchone()
                return dict(row) if row else None

    def get_doc_index_by_hash(self, content_hash: str) -> Optional[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM doc_indices WHERE content_hash = %s ORDER BY updated_at DESC LIMIT 1",
                    (content_hash,),
                )
                row = cur.fetchone()
                return dict(row) if row else None

    def delete_doc_index(self, doc_id: str) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM doc_indices WHERE doc_id = %s", (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._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO chat_documents(chat_id, doc_id, doc_name, attached_at)
                    VALUES (%s, %s, %s, %s)
                    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),
                )
                # Maintain backward compatibility on chats table
                cur.execute(
                    "UPDATE chats SET doc_id = %s, doc_name = %s, updated_at = %s WHERE chat_id = %s",
                    (doc_id, doc_name, now, chat_id),
                )

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

    def get_chat_documents(self, chat_id: str) -> List[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.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 = %s
                    ORDER BY cd.attached_at ASC
                    """,
                    (chat_id,),
                )
                rows = cur.fetchall()
                return [dict(r) for r in rows]

    def clear_chat_documents(self, chat_id: str) -> None:
        now = time.time()
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM chat_documents WHERE chat_id = %s", (chat_id,))
                cur.execute(
                    "UPDATE chats SET doc_id = NULL, doc_name = NULL, updated_at = %s WHERE chat_id = %s",
                    (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._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO chat_summaries(chat_id, summary, turn_count, updated_at)
                    VALUES (%s, %s, %s, %s)
                    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]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM chat_summaries WHERE chat_id = %s", (chat_id,))
                row = cur.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._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO knowledge_memory(
                        doc_id, filename, executive_summary, key_entities, key_findings, topics, updated_at
                    )
                    VALUES (%s, %s, %s, %s, %s, %s, %s)
                    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,
                        psycopg2.extras.Json(key_entities or []),
                        psycopg2.extras.Json(key_findings or []),
                        psycopg2.extras.Json(topics or []),
                        now,
                    ),
                )

    def get_knowledge_dossier(self, doc_id: str) -> Optional[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM knowledge_memory WHERE doc_id = %s", (doc_id,))
                row = cur.fetchone()
                if not row:
                    return None
                res = dict(row)
                for k in ("key_entities", "key_findings", "topics"):
                    if res.get(k) is None:
                        res[k] = []
                return res

    def list_knowledge_dossiers(self) -> List[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM knowledge_memory ORDER BY updated_at DESC")
                rows = cur.fetchall()
                result = []
                for r in rows:
                    item = dict(r)
                    for k in ("key_entities", "key_findings", "topics"):
                        if item.get(k) is None:
                            item[k] = []
                    result.append(item)
                return result

    def delete_knowledge_dossier(self, doc_id: str) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM knowledge_memory WHERE doc_id = %s", (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._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO global_knowledge_summary(
                        id, collective_summary, total_documents, key_insights, updated_at
                    )
                    VALUES (%s, %s, %s, %s, %s)
                    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,
                        psycopg2.extras.Json(key_insights or []),
                        now,
                    ),
                )

    def get_global_knowledge_summary(self) -> Optional[Dict[str, Any]]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM global_knowledge_summary WHERE id = 'global'")
                row = cur.fetchone()
                if not row:
                    return None
                res = dict(row)
                if res.get("key_insights") is None:
                    res["key_insights"] = []
                return res