app/memory/database.py

"""
SQLite database setup — async aiosqlite/libsql with table creation on startup.
"""

from contextlib import asynccontextmanager
from pathlib import Path
import aiosqlite
from app.config import get_settings


class LibsqlRow(dict):
    """aiosqlite.Row compatible row wrapper for libsql-client."""
    def __init__(self, colnames, row_data):
        super().__init__(zip(colnames, row_data))
        self._list = list(row_data)

    def __getitem__(self, key):
        if isinstance(key, int):
            return self._list[key]
        return super().__getitem__(key)


class LibsqlCursorWrapper:
    """Cursor-like wrapper for remote libsql execution. Supports both awaiting and async with."""
    def __init__(self, connection, client, sql, parameters):
        self.connection = connection
        self.client = client
        self.sql = sql
        self.parameters = parameters
        self.result_set = None
        self.row_factory = connection.row_factory
        self._index = 0

    def __await__(self):
        return self._execute().__await__()

    async def _execute(self):
        if self.result_set is None:
            self.result_set = await self.client.execute(self.sql, self.parameters)
            if self.row_factory == aiosqlite.Row:
                self.row_factory = lambda cursor, row_data: LibsqlRow(self.result_set.columns, row_data)
        return self

    @property
    def rowcount(self) -> int:
        return getattr(self.result_set, "rows_affected", 0) if self.result_set is not None else 0

    async def fetchone(self):
        await self._execute()
        if self._index >= len(self.result_set.rows):
            return None
        row = self.result_set.rows[self._index]
        self._index += 1
        if self.row_factory:
            return self.row_factory(self, row)
        return row

    async def fetchall(self):
        await self._execute()
        rows = self.result_set.rows[self._index:]
        self._index = len(self.result_set.rows)
        if self.row_factory:
            return [self.row_factory(self, r) for r in rows]
        return rows

    async def __aenter__(self):
        await self._execute()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass


class LibsqlConnectionWrapper:
    """Connection-like wrapper for remote libsql client."""
    def __init__(self, client):
        self.client = client
        self._row_factory = None

    @property
    def row_factory(self):
        return self._row_factory

    @row_factory.setter
    def row_factory(self, factory):
        self._row_factory = factory

    def execute(self, sql: str, parameters: tuple = ()):
        # Return a wrapper that is both awaitable and an async context manager
        return LibsqlCursorWrapper(self, self.client, sql, parameters)

    async def executescript(self, script: str):
        # Split statements by semicolon and execute each
        statements = [stmt.strip() for stmt in script.split(";") if stmt.strip()]
        for stmt in statements:
            await self.client.execute(stmt)

    async def commit(self):
        # libsql auto-commits
        pass

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass


_global_client = None


async def close_db():
    """Close the global database client if active."""
    global _global_client
    if _global_client is not None:
        await _global_client.close()
        _global_client = None


@asynccontextmanager
async def connect_db():
    """Yield a database connection wrapper based on active settings (Turso vs. local sqlite)."""
    global _global_client
    settings = get_settings()
    if settings.TURSO_DATABASE_URL and settings.TURSO_AUTH_TOKEN:
        if _global_client is None:
            url = settings.TURSO_DATABASE_URL
            if url.startswith("libsql://"):
                url = "https://" + url[len("libsql://"):]
            import libsql_client
            _global_client = libsql_client.create_client(
                url=url,
                auth_token=settings.TURSO_AUTH_TOKEN
            )
        yield LibsqlConnectionWrapper(_global_client)
    else:
        db_path = settings.DB_PATH
        Path(db_path).parent.mkdir(parents=True, exist_ok=True)
        async with aiosqlite.connect(db_path) as db:
            yield db


async def get_db_path() -> str:
    settings = get_settings()
    Path(settings.DB_PATH).parent.mkdir(parents=True, exist_ok=True)
    return settings.DB_PATH


async def init_db() -> None:
    """Create all tables if they don't exist."""
    async with connect_db() as db:
        await db.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                username TEXT PRIMARY KEY,
                password_hash TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS user_sessions (
                token TEXT PRIMARY KEY,
                username TEXT NOT NULL REFERENCES users(username) ON DELETE CASCADE,
                expires_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS sessions (
                id TEXT PRIMARY KEY,
                user_id TEXT NOT NULL,
                title TEXT,
                created_at TEXT NOT NULL,
                last_active_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS messages (
                id TEXT PRIMARY KEY,
                session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                intent TEXT,
                confidence TEXT,
                response_id TEXT,
                citations TEXT,
                pipeline_trace TEXT,
                created_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS feedback (
                id TEXT PRIMARY KEY,
                response_id TEXT NOT NULL,
                session_id TEXT NOT NULL,
                rating INTEGER NOT NULL,
                comment TEXT,
                created_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS user_documents (
                id TEXT PRIMARY KEY,
                user_id TEXT NOT NULL,
                filename TEXT NOT NULL,
                file_type TEXT,
                collection_name TEXT NOT NULL,
                chunk_count INTEGER DEFAULT 0,
                file_size_bytes INTEGER DEFAULT 0,
                uploaded_at TEXT NOT NULL
            );

            CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id);
            CREATE INDEX IF NOT EXISTS idx_sessions_user ON sessions(user_id);
            CREATE INDEX IF NOT EXISTS idx_user_docs_user ON user_documents(user_id);
            CREATE INDEX IF NOT EXISTS idx_feedback_response ON feedback(response_id);
        """)
        
        # Run-time schema migration check
        db.row_factory = aiosqlite.Row
        async with db.execute("PRAGMA table_info(messages)") as cur:
            columns = [row["name"] for row in await cur.fetchall()]
        
        if "citations" not in columns:
            await db.execute("ALTER TABLE messages ADD COLUMN citations TEXT")
        if "pipeline_trace" not in columns:
            await db.execute("ALTER TABLE messages ADD COLUMN pipeline_trace TEXT")
            
        await db.commit()