Barclays-banner-builder/backend/app/api/conversations.py

from datetime import datetime, timezone
from uuid import UUID

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy import select, desc
from sqlalchemy.orm import selectinload
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.auth import get_current_user
from app.database import get_db
from app.models.brief import Brief
from app.models.conversation import Conversation, ConversationMessage
from app.models.job import Job, JobStatus, JobType
from app.models.user import User
from app.services.intent_router import classify_turn

router = APIRouter()


def _serialize_message(m: ConversationMessage) -> dict:
    return {
        "id": str(m.id),
        "role": m.role,
        "kind": m.kind,
        "content": m.content,
        "banner_set_id": str(m.banner_set_id) if m.banner_set_id else None,
        "job_id": str(m.job_id) if m.job_id else None,
        "created_at": m.created_at,
    }


def _serialize_conversation(c: Conversation) -> dict:
    return {
        "id": str(c.id),
        "brief_id": str(c.brief_id),
        "title": c.title,
        "created_at": c.created_at,
        "updated_at": c.updated_at,
        "messages": [_serialize_message(m) for m in c.messages],
    }


class CreateConversationPayload(BaseModel):
    first_message: str = ""


class SendMessagePayload(BaseModel):
    text: str


@router.post("", status_code=201)
async def create_conversation(
    payload: CreateConversationPayload,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    from app.workers.tasks import queue

    brief = Brief(user_id=current_user.id, text=payload.first_message or "")
    db.add(brief)
    await db.flush()

    title = (payload.first_message or "New conversation")[:60]
    conversation = Conversation(
        user_id=current_user.id,
        brief_id=brief.id,
        title=title,
    )
    db.add(conversation)
    await db.flush()

    first_message_id = None
    job_id = None

    if payload.first_message.strip():
        user_msg = ConversationMessage(
            conversation_id=conversation.id,
            role="user",
            kind="text",
            content=payload.first_message,
        )
        db.add(user_msg)
        await db.flush()
        first_message_id = str(user_msg.id)

        job = Job(
            type=JobType.CHAT_TURN,
            status=JobStatus.PENDING,
            payload={
                "conversation_id": str(conversation.id),
                "message_id": str(user_msg.id),
                "message_text": payload.first_message,
            },
        )
        db.add(job)
        await db.commit()
        await db.refresh(job)
        queue.enqueue("app.workers.tasks.chat_turn", str(job.id), str(conversation.id), job_timeout=180)
        job_id = str(job.id)
    else:
        await db.commit()

    return {
        "conversation_id": str(conversation.id),
        "brief_id": str(brief.id),
        "message_id": first_message_id,
        "job_id": job_id,
    }


@router.get("")
async def list_conversations(
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    result = await db.execute(
        select(Conversation)
        .where(Conversation.user_id == current_user.id)
        .order_by(desc(Conversation.updated_at))
        .limit(50)
    )
    convs = result.scalars().all()
    return {
        "conversations": [
            {"id": str(c.id), "title": c.title, "created_at": c.created_at, "updated_at": c.updated_at}
            for c in convs
        ]
    }


@router.get("/{conversation_id}")
async def get_conversation(
    conversation_id: UUID,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    result = await db.execute(
        select(Conversation)
        .where(Conversation.id == conversation_id, Conversation.user_id == current_user.id)
        .options(selectinload(Conversation.messages))
    )
    conv = result.scalar_one_or_none()
    if not conv:
        raise HTTPException(status_code=404, detail="Conversation not found")
    return _serialize_conversation(conv)


@router.post("/{conversation_id}/messages", status_code=202)
async def send_message(
    conversation_id: UUID,
    payload: SendMessagePayload,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    from app.workers.tasks import queue

    result = await db.execute(
        select(Conversation)
        .where(Conversation.id == conversation_id, Conversation.user_id == current_user.id)
        .options(selectinload(Conversation.messages))
    )
    conv = result.scalar_one_or_none()
    if not conv:
        raise HTTPException(status_code=404, detail="Conversation not found")

    brief_result = await db.execute(select(Brief).where(Brief.id == conv.brief_id))
    brief = brief_result.scalar_one()

    history = [
        {"role": m.role, "content": m.content}
        for m in conv.messages
        if m.role in ("user", "assistant")
    ]

    intent = await classify_turn(brief.text, history, payload.text)

    user_msg = ConversationMessage(
        conversation_id=conv.id,
        role="user",
        kind="text",
        content=payload.text,
    )
    db.add(user_msg)
    await db.flush()

    if intent.intent == "clarify":
        assistant_msg = ConversationMessage(
            conversation_id=conv.id,
            role="assistant",
            kind="text",
            content=intent.reply_text or "How can I help you with your banners?",
        )
        db.add(assistant_msg)
        conv.updated_at = datetime.now(timezone.utc)
        await db.commit()
        return {
            "message_id": str(user_msg.id),
            "job_id": None,
            "reply": intent.reply_text,
        }

    job = Job(
        type=JobType.CHAT_TURN,
        status=JobStatus.PENDING,
        payload={
            "conversation_id": str(conv.id),
            "message_id": str(user_msg.id),
            "message_text": payload.text,
        },
    )
    db.add(job)
    conv.updated_at = datetime.now(timezone.utc)
    await db.commit()
    await db.refresh(job)

    queue.enqueue("app.workers.tasks.chat_turn", str(job.id), str(conv.id), job_timeout=180)

    return {
        "message_id": str(user_msg.id),
        "job_id": str(job.id),
        "reply": None,
    }