backend/db/chat_store.py

"""SQLite chat session storage."""

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

from config import settings


class ChatStore:
    def __init__(self, db_path: Optional[str] = None):
        self.db_path = db_path or settings.CHAT_DB_PATH
        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn

    def _schema_sql(self) -> tuple[str, str, str]:
        sessions = """
                CREATE TABLE IF NOT EXISTS sessions (
                    id TEXT PRIMARY KEY,
                    title TEXT NOT NULL,
                    created_at TEXT NOT NULL
                )
                """
        messages = """
                CREATE TABLE IF NOT EXISTS messages (
                    id TEXT PRIMARY KEY,
                    session_id TEXT NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT NOT NULL,
                    sources_json TEXT,
                    confidence REAL,
                    created_at TEXT NOT NULL,
                    FOREIGN KEY (session_id) REFERENCES sessions(id)
                )
                """
        session_documents = """
                CREATE TABLE IF NOT EXISTS session_documents (
                    session_id TEXT NOT NULL,
                    document_id TEXT NOT NULL,
                    PRIMARY KEY (session_id, document_id),
                    FOREIGN KEY (session_id) REFERENCES sessions(id)
                )
                """
        return sessions, messages, session_documents

    def init_db(self) -> None:
        sessions, messages, session_documents = self._schema_sql()
        with sqlite3.connect(self.db_path) as db:
            db.execute(sessions)
            db.execute(messages)
            db.execute(session_documents)
            db.commit()

    def create_session(
        self, title: str = "New Chat", document_ids: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        session_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc).isoformat()
        with sqlite3.connect(self.db_path) as db:
            db.execute(
                "INSERT INTO sessions (id, title, created_at) VALUES (?, ?, ?)",
                (session_id, title, now),
            )
            if document_ids:
                for doc_id in document_ids:
                    db.execute(
                        "INSERT INTO session_documents (session_id, document_id) VALUES (?, ?)",
                        (session_id, doc_id),
                    )
            db.commit()
        return {"id": session_id, "title": title, "created_at": now}

    def list_sessions(self) -> List[Dict[str, Any]]:
        with self._connect() as db:
            cursor = db.execute(
                "SELECT id, title, created_at FROM sessions ORDER BY created_at DESC"
            )
            return [dict(row) for row in cursor.fetchall()]

    def get_messages(self, session_id: str) -> List[Dict[str, Any]]:
        with self._connect() as db:
            cursor = db.execute(
                """
                SELECT id, role, content, sources_json, confidence, created_at
                FROM messages WHERE session_id = ? ORDER BY created_at ASC
                """,
                (session_id,),
            )
            result = []
            for row in cursor.fetchall():
                msg = dict(row)
                if msg.get("sources_json"):
                    msg["sources"] = json.loads(msg["sources_json"])
                else:
                    msg["sources"] = []
                del msg["sources_json"]
                result.append(msg)
            return result

    def add_message(
        self,
        session_id: str,
        role: str,
        content: str,
        sources: Optional[List[Dict]] = None,
        confidence: Optional[float] = None,
    ) -> Dict[str, Any]:
        msg_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc).isoformat()
        sources_json = json.dumps(sources or [])
        with sqlite3.connect(self.db_path) as db:
            db.execute(
                """
                INSERT INTO messages (id, session_id, role, content, sources_json, confidence, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (msg_id, session_id, role, content, sources_json, confidence, now),
            )
            db.commit()
        return {
            "id": msg_id,
            "session_id": session_id,
            "role": role,
            "content": content,
            "sources": sources or [],
            "confidence": confidence,
            "created_at": now,
        }

    def export_session(self, session_id: str) -> str:
        with self._connect() as db:
            cursor = db.execute(
                "SELECT title, created_at FROM sessions WHERE id = ?",
                (session_id,),
            )
            session = cursor.fetchone()
            if not session:
                return ""

        messages = self.get_messages(session_id)
        lines = [
            "FinSight AI Chat Export",
            f"Session: {session['title']}",
            f"Created: {session['created_at']}",
            "=" * 60,
            "",
        ]
        for msg in messages:
            role_label = "You" if msg["role"] == "user" else "FinSight AI"
            lines.append(f"## {role_label} ({msg['created_at']})")
            lines.append(msg["content"])
            if msg.get("confidence"):
                lines.append(f"\n[Confidence: {msg['confidence']}/10]")
            if msg.get("sources"):
                lines.append("\nSources:")
                for i, src in enumerate(msg["sources"], 1):
                    lines.append(
                        f"  {i}. {src.get('document_name', '')} p.{src.get('page_number', '?')} "
                        f"({src.get('section', '')})"
                    )
            lines.append("")
        return "\n".join(lines)