gitdataai/libs/room/src/message.rs

use crate::error::RoomError;
use crate::service::RoomService;
use crate::ws_context::WsUserContext;
use models::rooms::{room_attachment, room_message, room_message_reaction};
use models::users::user as user_model;
use sea_orm::*;
use uuid::Uuid;

impl RoomService {
    pub async fn room_message_list(
        &self,
        room_id: Uuid,
        before_seq: Option<i64>,
        after_seq: Option<i64>,
        limit: Option<u64>,
        ctx: &WsUserContext,
    ) -> Result<super::RoomMessageListResponse, RoomError> {
        let user_id = ctx.user_id;
        self.require_room_access(room_id, user_id).await?;

        let use_asc = after_seq.is_some();

        let mut query = room_message::Entity::find().filter(room_message::Column::Room.eq(room_id));
        if let Some(bs) = before_seq {
            query = query.filter(room_message::Column::Seq.lt(bs));
        }
        if let Some(a_s) = after_seq {
            query = query.filter(room_message::Column::Seq.gt(a_s));
        }

        let total = query.clone().count(&self.db).await? as i64;
        let models = {
            let mut q = query;
            if use_asc {
                q = q.order_by_asc(room_message::Column::Seq);
            } else {
                q = q.order_by_desc(room_message::Column::Seq);
            }
            q.limit(limit.unwrap_or(50)).all(&self.db).await?
        };

        let user_ids: Vec<Uuid> = models
            .iter()
            .filter(|m| m.sender_type.to_string() == "user")
            .filter_map(|m| m.sender_id)
            .collect();
        let ai_model_ids: Vec<Uuid> = models
            .iter()
            .filter(|m| m.sender_type.to_string() == "ai")
            .filter_map(|m| m.model_id)
            .collect();

        let users: std::collections::HashMap<Uuid, String> = if !user_ids.is_empty() {
            user_model::Entity::find()
                .filter(user_model::Column::Uid.is_in(user_ids))
                .all(&self.db)
                .await?
                .into_iter()
                .map(|u| (u.uid, u.display_name.unwrap_or(u.username)))
                .collect()
        } else {
            std::collections::HashMap::new()
        };

        let ai_names: std::collections::HashMap<Uuid, String> = if !ai_model_ids.is_empty() {
            models::agents::model::Entity::find()
                .filter(models::agents::model::Column::Id.is_in(ai_model_ids))
                .all(&self.db)
                .await?
                .into_iter()
                .map(|m| (m.id, m.name))
                .collect()
        } else {
            std::collections::HashMap::new()
        };

        let mut messages: Vec<super::RoomMessageResponse> = models
            .into_iter()
            .map(|msg| {
                let sender_type = msg.sender_type.to_string();
                let display_name = match sender_type.as_str() {
                    "ai" => msg.model_id.and_then(|id| {
                        ai_names
                            .get(&id)
                            .cloned()
                            .or_else(|| Some(format!("AI({})", &id.to_string()[..8])))
                    }),
                    _ => msg.sender_id.and_then(|id| users.get(&id).cloned()),
                }.or_else(|| msg.sender_id.map(|id| id.to_string()));
                let chunked = super::RoomMessageResponse::detect_chunked(&msg.thinking_content);
                super::RoomMessageResponse {
                    id: msg.id,
                    seq: msg.seq,
                    room: msg.room,
                    sender_type,
                    sender_id: msg.sender_id,
                    display_name,
                    thread: msg.thread,
                    in_reply_to: msg.in_reply_to,
                    content: msg.content,
                    content_type: msg.content_type.to_string(),
                    thinking_content: msg.thinking_content,
                    thinking_is_chunked: chunked,
                    edited_at: msg.edited_at,
                    send_at: msg.send_at,
                    revoked: msg.revoked,
                    revoked_by: msg.revoked_by,
                    highlighted_content: None,
                    attachment_ids: Vec::new(),
                    reactions: Vec::new(),
                }
            })
            .collect();
        if !use_asc {
            messages.reverse();
        }

        if !messages.is_empty() {
            let msg_ids: Vec<Uuid> = messages.iter().map(|m| m.id).collect();
            let attachments = room_attachment::Entity::find()
                .filter(room_attachment::Column::Message.is_in(msg_ids.clone()))
                .all(&self.db)
                .await
                .unwrap_or_default();

            let mut attachment_map: std::collections::HashMap<Uuid, Vec<Uuid>> =
                std::collections::HashMap::new();
            for att in attachments {
                attachment_map.entry(att.message).or_default().push(att.id);
            }

            for msg in &mut messages {
                if let Some(ids) = attachment_map.remove(&msg.id) {
                    msg.attachment_ids = ids;
                }
            }

            // Load reactions for all messages
            let reactions = room_message_reaction::Entity::find()
                .filter(room_message_reaction::Column::Message.is_in(msg_ids))
                .all(&self.db)
                .await
                .unwrap_or_default();

            let mut reaction_map: std::collections::HashMap<Uuid, Vec<room_message_reaction::Model>> =
                std::collections::HashMap::new();
            for r in reactions {
                reaction_map.entry(r.message).or_default().push(r);
            }

            for msg in &mut messages {
                if let Some(reaction_models) = reaction_map.remove(&msg.id) {
                    msg.reactions = self.build_reaction_groups(reaction_models, Some(user_id));
                }
            }
        }

        Ok(super::RoomMessageListResponse { messages, total })
    }

    pub async fn room_message_get(
        &self,
        message_id: Uuid,
        ctx: &WsUserContext,
    ) -> Result<super::RoomMessageResponse, RoomError> {
        let user_id = ctx.user_id;
        let model = room_message::Entity::find_by_id(message_id)
            .one(&self.db)
            .await?
            .ok_or_else(|| RoomError::NotFound("Message not found".to_string()))?;
        let room_id = model.room;
        self.require_room_access(room_id, user_id).await?;
        Ok(self.resolve_display_name(model, room_id).await)
    }
}