import uuid
from datetime import datetime, timezone
from pymongo.asynchronous.database import AsyncDatabase

async def create_message(
    db: AsyncDatabase,
    conversation_id: str,
    session_id: str,
    role: str,
    content: str,
    agent_id: str | None = None
) -> dict:
    msg_id = f"msg_{uuid.uuid4().hex[:12]}"
    now = datetime.now(timezone.utc)
    
    msg_doc = {
        "messageId": msg_id,
        "conversationId": conversation_id,
        "sessionId": session_id,
        "role": role,
        "content": content,
        "createdAt": now
    }

    if agent_id is not None:
        msg_doc["agentId"] = agent_id
    
    await db["messages"].insert_one(msg_doc)
    
    # Update conversation's updatedAt timestamp
    await db["conversations"].update_one(
        {"conversationId": conversation_id},
        {"$set": {"updatedAt": now}}
    )
    
    return msg_doc

async def get_messages_by_conversation_id(db: AsyncDatabase, conversation_id: str) -> list[dict]:
    cursor = db["messages"].find({"conversationId": conversation_id}).sort("createdAt", 1)
    return await cursor.to_list(length=1000)