grimlock/backend/api/messages.py

"""
Messages API - Send and receive messages with @mention support
"""

from fastapi import APIRouter, Depends, HTTPException, status, BackgroundTasks
from sqlalchemy.orm import Session
from pydantic import BaseModel
from typing import List, Optional
import re
from datetime import datetime

from core.database import get_db
from core.models import Message, Channel, User, ChannelType
from core.ai_client import AIClient
from core.context_manager import ContextManager
from api.auth import get_current_user
import main

router = APIRouter()

# Pydantic models
class MessageCreate(BaseModel):
    content: str
    reply_to: Optional[int] = None

class UserInfo(BaseModel):
    id: int
    name: str
    email: str
    role: str
    is_online: bool
    
    class Config:
        from_attributes = True

class MessageResponse(BaseModel):
    id: int
    content: str
    is_ai_message: bool
    user: Optional[UserInfo]
    reply_to_message_id: Optional[int]
    created_at: str
    edited_at: Optional[str]
    
    class Config:
        from_attributes = True

def detect_grimlock_mention(content: str) -> bool:
    """Detect if @grimlock is mentioned in message"""
    return bool(re.search(r'@grimlock\b', content, re.IGNORECASE))

def extract_mention_query(content: str) -> str:
    """Extract the query after @grimlock mention"""
    # Remove @grimlock and return the rest
    return re.sub(r'@grimlock\s*', '', content, flags=re.IGNORECASE).strip()

async def handle_grimlock_mention(
    message: Message,
    channel: Channel,
    db: Session,
    context_manager: ContextManager,
    ai_client: AIClient
):
    """Handle @grimlock mention - respond with AI"""
    
    try:
        # Extract query
        query = extract_mention_query(message.content)
        
        # Get channel history for context (last 10 messages)
        recent_messages = db.query(Message)\
            .filter(Message.channel_id == channel.id)\
            .filter(Message.id < message.id)\
            .order_by(Message.id.desc())\
            .limit(10)\
            .all()
        
        # Build conversation history
        conversation = []
        for msg in reversed(recent_messages):
            if msg.user:
                conversation.append({
                    "role": "user",
                    "content": f"{msg.user.name}: {msg.content}"
                })
            elif msg.is_ai_message:
                conversation.append({
                    "role": "assistant",
                    "content": msg.content
                })
        
        # Add current message
        if message.user:
            conversation.append({
                "role": "user",
                "content": f"{message.user.name}: {query}"
            })
        
        # Get context from context manager
        context = context_manager.get_context_for_query(query)
        system_prompt = context_manager.get_system_prompt()
        
        # Add channel context
        system_prompt += f"\n\nYou are responding in channel #{channel.name}."
        if context:
            system_prompt += f"\n\n# Company Context\n{context}"
        
        # Get AI response
        response = await ai_client.chat(
            messages=conversation,
            system_prompt=system_prompt
        )
        
        # Create AI message
        ai_message = Message(
            channel_id=channel.id,
            user_id=None,
            content=response,
            is_ai_message=True,
            reply_to_message_id=message.id
        )
        
        db.add(ai_message)
        db.commit()
        db.refresh(ai_message)
        
        return ai_message
        
    except Exception as e:
        # Create error message
        error_message = Message(
            channel_id=channel.id,
            user_id=None,
            content=f"Sorry, I encountered an error: {str(e)}",
            is_ai_message=True,
            reply_to_message_id=message.id
        )
        db.add(error_message)
        db.commit()
        return error_message

@router.post("/{channel_id}/messages", response_model=MessageResponse)
async def send_message(
    channel_id: int,
    message_data: MessageCreate,
    background_tasks: BackgroundTasks,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
    context_manager: ContextManager = Depends(main.get_context_manager),
    ai_client: AIClient = Depends(main.get_ai_client)
):
    """Send a message to a channel"""
    
    # Get channel
    channel = db.query(Channel).filter(Channel.id == channel_id).first()
    if not channel:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Channel not found"
        )
    
    # Check if user is member
    if current_user not in channel.members:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Not a member of this channel"
        )
    
    # Create message
    message = Message(
        channel_id=channel_id,
        user_id=current_user.id,
        content=message_data.content,
        reply_to_message_id=message_data.reply_to
    )
    
    db.add(message)
    db.commit()
    db.refresh(message)
    
    # Check for @grimlock mention
    if detect_grimlock_mention(message_data.content):
        # Handle in background to not block response
        background_tasks.add_task(
            handle_grimlock_mention,
            message,
            channel,
            db,
            context_manager,
            ai_client
        )
    
    return {
        "id": message.id,
        "content": message.content,
        "is_ai_message": message.is_ai_message,
        "user": {
            "id": current_user.id,
            "name": current_user.name,
            "email": current_user.email,
            "role": current_user.role.value,
            "is_online": current_user.is_online
        } if message.user else None,
        "reply_to_message_id": message.reply_to_message_id,
        "created_at": message.created_at.isoformat(),
        "edited_at": message.edited_at.isoformat() if message.edited_at else None
    }

@router.get("/{channel_id}/messages", response_model=List[MessageResponse])
async def get_messages(
    channel_id: int,
    limit: int = 50,
    before: Optional[int] = None,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Get messages from a channel"""
    
    # Get channel
    channel = db.query(Channel).filter(Channel.id == channel_id).first()
    if not channel:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Channel not found"
        )
    
    # Check if user is member
    if current_user not in channel.members:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Not a member of this channel"
        )
    
    # Query messages
    query = db.query(Message).filter(Message.channel_id == channel_id)
    
    if before:
        query = query.filter(Message.id < before)
    
    messages = query.order_by(Message.id.desc()).limit(limit).all()
    messages.reverse()  # Return in chronological order
    
    # Format response
    result = []
    for msg in messages:
        result.append({
            "id": msg.id,
            "content": msg.content,
            "is_ai_message": msg.is_ai_message,
            "user": {
                "id": msg.user.id,
                "name": msg.user.name,
                "email": msg.user.email,
                "role": msg.user.role.value,
                "is_online": msg.user.is_online
            } if msg.user else None,
            "reply_to_message_id": msg.reply_to_message_id,
            "created_at": msg.created_at.isoformat(),
            "edited_at": msg.edited_at.isoformat() if msg.edited_at else None
        })
    
    return result