gitdataai/libs/transport/pagination.rs

use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessagePage {
    pub messages: Vec<MessageItem>,
    pub has_more: bool,
    pub next_cursor: Option<String>,
    pub prev_cursor: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageItem {
    pub id: Uuid,
    pub room_id: Uuid,
    pub seq: i64,
    pub content: String,
    pub sender_id: Option<Uuid>,
    pub sender_name: Option<String>,
    pub send_at: chrono::DateTime<chrono::Utc>,
    pub edited_at: Option<chrono::DateTime<chrono::Utc>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PaginationParams {
    pub room_id: Uuid,
    pub limit: u64,
    pub cursor: Option<String>,
    pub direction: PaginationDirection,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum PaginationDirection {
    Before,
    After,
}

pub struct MessagePagination {
    db: db::database::AppDatabase,
}

impl MessagePagination {
    pub fn new(db: db::database::AppDatabase) -> Self {
        Self { db }
    }

    pub async fn get_messages(
        &self,
        params: PaginationParams,
    ) -> Result<MessagePage, crate::error::AppTransportError> {
        use models::rooms::room_message;
        use sea_orm::*;

        let limit = std::cmp::Ord::min(params.limit, 100);
        let cursor_seq = if let Some(cursor) = params.cursor {
            cursor.parse::<i64>().ok()
        } else {
            None
        };

        let mut query =
            room_message::Entity::find().filter(room_message::Column::Room.eq(params.room_id));

        query = match (params.direction, cursor_seq) {
            (PaginationDirection::Before, Some(seq)) => query
                .filter(room_message::Column::Seq.lt(seq))
                .order_by_desc(room_message::Column::Seq),
            (PaginationDirection::After, Some(seq)) => query
                .filter(room_message::Column::Seq.gt(seq))
                .order_by_asc(room_message::Column::Seq),
            _ => query.order_by_desc(room_message::Column::Seq),
        };

        let messages = query
            .limit(limit + 1)
            .all(&self.db)
            .await
            .map_err(|_| crate::error::AppTransportError::Internal)?;

        let has_more = messages.len() > limit as usize;
        let messages: Vec<_> = messages.into_iter().take(limit as usize).collect();

        let next_cursor = if has_more {
            messages.last().map(|m| m.seq.to_string())
        } else {
            None
        };

        let prev_cursor = messages.first().map(|m| m.seq.to_string());

        let items: Vec<MessageItem> = messages
            .into_iter()
            .map(|m| MessageItem {
                id: m.id,
                room_id: m.room,
                seq: m.seq,
                content: m.content,
                sender_id: m.sender_id,
                sender_name: None,
                send_at: m.send_at,
                edited_at: m.edited_at,
            })
            .collect();

        Ok(MessagePage {
            messages: items,
            has_more,
            next_cursor,
            prev_cursor,
        })
    }

    pub async fn get_messages_around(
        &self,
        room_id: Uuid,
        message_id: Uuid,
        context_size: u64,
    ) -> Result<MessagePage, crate::error::AppTransportError> {
        use models::rooms::room_message;
        use sea_orm::*;

        let target = room_message::Entity::find_by_id(message_id)
            .one(&self.db)
            .await
            .map_err(|_| crate::error::AppTransportError::Internal)?
            .ok_or(crate::error::AppTransportError::Internal)?;

        let before = room_message::Entity::find()
            .filter(room_message::Column::Room.eq(room_id))
            .filter(room_message::Column::Seq.lt(target.seq))
            .order_by_desc(room_message::Column::Seq)
            .limit(context_size)
            .all(&self.db)
            .await
            .map_err(|_| crate::error::AppTransportError::Internal)?;

        let after = room_message::Entity::find()
            .filter(room_message::Column::Room.eq(room_id))
            .filter(room_message::Column::Seq.gt(target.seq))
            .order_by_asc(room_message::Column::Seq)
            .limit(context_size)
            .all(&self.db)
            .await
            .map_err(|_| crate::error::AppTransportError::Internal)?;

        let mut all_messages = before;
        all_messages.reverse();
        all_messages.push(target);
        all_messages.extend(after);

        let items: Vec<MessageItem> = all_messages
            .into_iter()
            .map(|m| MessageItem {
                id: m.id,
                room_id: m.room,
                seq: m.seq,
                content: m.content,
                sender_id: m.sender_id,
                sender_name: None,
                send_at: m.send_at,
                edited_at: m.edited_at,
            })
            .collect();

        Ok(MessagePage {
            messages: items,
            has_more: false,
            next_cursor: None,
            prev_cursor: None,
        })
    }
}