app/db/chat_db.py

"""
Per-user chat database models.

All conversation and message data are stored in the main PostgreSQL database
and separated strictly using the user_id column.

Schema:
  conversations  – chat sessions (like ChatGPT sidebar items)
  messages       – individual messages within a conversation
"""
import os
import uuid
from datetime import datetime, timezone
from typing import List, Optional

from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import String, DateTime, Text, ForeignKey, select, update
from loguru import logger

from app.db.database import engine, AsyncSessionLocal


class ChatBase(DeclarativeBase):
    pass


class Conversation(ChatBase):
    __tablename__ = "conversations"

    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    user_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True)
    title: Mapped[str] = mapped_column(String(512), default="New Chat")
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        onupdate=lambda: datetime.now(timezone.utc),
    )


class Message(ChatBase):
    __tablename__ = "messages"

    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    conversation_id: Mapped[str] = mapped_column(
        String(36), ForeignKey("conversations.id"), nullable=False, index=True
    )
    role: Mapped[str] = mapped_column(String(20), nullable=False)   # "user" | "assistant"
    content: Mapped[str] = mapped_column(Text, nullable=False)
    sources_json: Mapped[str | None] = mapped_column(Text)          # JSON-serialised SourceChunk list
    meta_json: Mapped[str | None] = mapped_column(Text)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
    )


# ─── Connection Sharing ────────────────────────────────────────────────────────

def _get_user_engine(user_id: str):
    return engine


async def ensure_user_db(user_id: str) -> None:
    # No-op since tables are created on the main connection during startup
    pass


def _session_maker(user_id: str) -> async_sessionmaker:
    return AsyncSessionLocal


# ─── CRUD helpers ─────────────────────────────────────────────────────────────

async def create_conversation(user_id: str, title: str = "New Chat") -> Conversation:
    async with _session_maker(user_id)() as db:
        conv = Conversation(id=str(uuid.uuid4()), user_id=user_id, title=title)
        db.add(conv)
        await db.commit()
        await db.refresh(conv)
        return conv


async def list_conversations(user_id: str) -> List[Conversation]:
    async with _session_maker(user_id)() as db:
        result = await db.execute(
            select(Conversation)
            .where(Conversation.user_id == user_id)
            .order_by(Conversation.updated_at.desc())
        )
        return result.scalars().all()


async def get_conversation(user_id: str, conv_id: str) -> Optional[Conversation]:
    async with _session_maker(user_id)() as db:
        result = await db.execute(
            select(Conversation).where(
                Conversation.id == conv_id, Conversation.user_id == user_id
            )
        )
        return result.scalar_one_or_none()


async def rename_conversation(user_id: str, conv_id: str, new_title: str) -> None:
    async with _session_maker(user_id)() as db:
        await db.execute(
            update(Conversation)
            .where(Conversation.id == conv_id, Conversation.user_id == user_id)
            .values(title=new_title, updated_at=datetime.now(timezone.utc))
        )
        await db.commit()


async def delete_conversation(user_id: str, conv_id: str) -> None:
    async with _session_maker(user_id)() as db:
        result = await db.execute(
            select(Conversation).where(
                Conversation.id == conv_id, Conversation.user_id == user_id
            )
        )
        conv = result.scalar_one_or_none()
        if conv:
            # cascade delete messages first
            msgs = await db.execute(select(Message).where(Message.conversation_id == conv_id))
            for msg in msgs.scalars().all():
                await db.delete(msg)
            await db.delete(conv)
            await db.commit()


async def add_message(
    user_id: str,
    conv_id: str,
    role: str,
    content: str,
    sources_json: str | None = None,
    meta_json: str | None = None,
) -> Message:
    async with _session_maker(user_id)() as db:
        # Verify ownership
        conv_check = await db.execute(
            select(Conversation).where(
                Conversation.id == conv_id, Conversation.user_id == user_id
            )
        )
        if not conv_check.scalar_one_or_none():
            raise ValueError("Conversation not found or access denied")

        msg = Message(
            id=str(uuid.uuid4()),
            conversation_id=conv_id,
            role=role,
            content=content,
            sources_json=sources_json,
            meta_json=meta_json,
        )
        db.add(msg)
        # Touch conversation updated_at so it floats to top
        await db.execute(
            update(Conversation)
            .where(Conversation.id == conv_id)
            .values(updated_at=datetime.now(timezone.utc))
        )
        await db.commit()
        await db.refresh(msg)
        return msg


async def get_messages(user_id: str, conv_id: str) -> List[Message]:
    async with _session_maker(user_id)() as db:
        # Verify ownership
        conv_check = await db.execute(
            select(Conversation).where(
                Conversation.id == conv_id, Conversation.user_id == user_id
            )
        )
        if not conv_check.scalar_one_or_none():
            return []

        result = await db.execute(
            select(Message)
            .where(Message.conversation_id == conv_id)
            .order_by(Message.created_at.asc())
        )
        return result.scalars().all()


async def auto_title_conversation(user_id: str, conv_id: str, first_question: str) -> None:
    """Set a smart title from the first user message (truncated)."""
    title = first_question[:60] + ("…" if len(first_question) > 60 else "")
    await rename_conversation(user_id, conv_id, title)