AppFlowy/frontend/rust-lib/event-integration-test/tests/sql_test/chat_message_test.rs

use event_integration_test::user_event::use_localhost_af_cloud;
use event_integration_test::EventIntegrationTest;
use flowy_ai_pub::cloud::MessageCursor;
use flowy_ai_pub::persistence::{
  select_answer_where_match_reply_message_id, select_chat_messages, select_message,
  select_message_content, total_message_count, upsert_chat_messages, ChatMessageTable,
};
use uuid::Uuid;

#[tokio::test]
async fn chat_message_table_insert_select_test() {
  use_localhost_af_cloud().await;
  let test = EventIntegrationTest::new().await;
  test.sign_up_as_anon().await;

  let uid = test.user_manager.get_anon_user().await.unwrap().id;
  let db_conn = test.user_manager.db_connection(uid).unwrap();

  let chat_id = Uuid::new_v4().to_string();
  let message_id_1 = 1000;
  let message_id_2 = 2000;

  // Create test messages
  let messages = vec![
    ChatMessageTable {
      message_id: message_id_1,
      chat_id: chat_id.clone(),
      content: "Hello, this is a test message".to_string(),
      created_at: 1625097600, // 2021-07-01
      author_type: 1,         // User
      author_id: "user_1".to_string(),
      reply_message_id: None,
      metadata: None,
      is_sync: false,
    },
    ChatMessageTable {
      message_id: message_id_2,
      chat_id: chat_id.clone(),
      content: "This is a reply to the test message".to_string(),
      created_at: 1625097700, // 2021-07-01, 100 seconds later
      author_type: 0,         // AI
      author_id: "ai".to_string(),
      reply_message_id: Some(message_id_1),
      metadata: Some(r#"{"source": "test"}"#.to_string()),
      is_sync: false,
    },
  ];

  // Test insert_chat_messages
  let result = upsert_chat_messages(db_conn, &messages);
  assert!(
    result.is_ok(),
    "Failed to insert chat messages: {:?}",
    result
  );

  // Test select_chat_messages
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let messages_result =
    select_chat_messages(db_conn, &chat_id, 10, MessageCursor::Offset(0)).unwrap();

  assert_eq!(messages_result.messages.len(), 2);
  assert_eq!(messages_result.total_count, 2);
  assert!(!messages_result.has_more);

  // Verify the content of the returned messages
  let first_message = messages_result
    .messages
    .iter()
    .find(|m| m.message_id == message_id_1)
    .unwrap();
  assert_eq!(first_message.content, "Hello, this is a test message");
  assert_eq!(first_message.author_type, 1);

  let second_message = messages_result
    .messages
    .iter()
    .find(|m| m.message_id == message_id_2)
    .unwrap();
  assert_eq!(
    second_message.content,
    "This is a reply to the test message"
  );
  assert_eq!(second_message.reply_message_id, Some(message_id_1));
}

#[tokio::test]
async fn chat_message_table_cursor_test() {
  use_localhost_af_cloud().await;
  let test = EventIntegrationTest::new().await;
  test.sign_up_as_anon().await;

  let uid = test.user_manager.get_anon_user().await.unwrap().id;
  let db_conn = test.user_manager.db_connection(uid).unwrap();

  let chat_id = Uuid::new_v4().to_string();

  // Create multiple test messages with sequential IDs
  let mut messages = Vec::new();
  for i in 1..6 {
    messages.push(ChatMessageTable {
      message_id: i * 1000,
      chat_id: chat_id.clone(),
      content: format!("Message {}", i),
      created_at: 1625097600 + (i * 100), // Increasing timestamps
      author_type: 1,                     // User
      author_id: "user_1".to_string(),
      reply_message_id: None,
      metadata: None,
      is_sync: false,
    });
  }

  // Insert messages
  upsert_chat_messages(db_conn, &messages).unwrap();

  // Test MessageCursor::Offset
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result_offset = select_chat_messages(
    db_conn,
    &chat_id,
    2, // Limit to 2 messages
    MessageCursor::Offset(0),
  )
  .unwrap();

  assert_eq!(result_offset.messages.len(), 2);
  assert!(result_offset.has_more);

  // Test MessageCursor::AfterMessageId
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result_after = select_chat_messages(
    db_conn,
    &chat_id,
    3, // Limit to 3 messages
    MessageCursor::AfterMessageId(2000),
  )
  .unwrap();

  assert_eq!(result_after.messages.len(), 3); // Should get message IDs 3000, 4000, 5000
  assert!(result_after.messages.iter().all(|m| m.message_id > 2000));

  // Test MessageCursor::BeforeMessageId
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result_before = select_chat_messages(
    db_conn,
    &chat_id,
    2, // Limit to 2 messages
    MessageCursor::BeforeMessageId(4000),
  )
  .unwrap();

  assert_eq!(result_before.messages.len(), 2); // Should get message IDs 1000, 2000, 3000
  assert!(result_before.messages.iter().all(|m| m.message_id < 4000));
}

#[tokio::test]
async fn chat_message_total_count_test() {
  use_localhost_af_cloud().await;
  let test = EventIntegrationTest::new().await;
  test.sign_up_as_anon().await;

  let uid = test.user_manager.get_anon_user().await.unwrap().id;
  let db_conn = test.user_manager.db_connection(uid).unwrap();

  let chat_id = Uuid::new_v4().to_string();

  // Create test messages
  let messages = vec![
    ChatMessageTable {
      message_id: 1001,
      chat_id: chat_id.clone(),
      content: "Message 1".to_string(),
      created_at: 1625097600,
      author_type: 1,
      author_id: "user_1".to_string(),
      reply_message_id: None,
      metadata: None,
      is_sync: false,
    },
    ChatMessageTable {
      message_id: 1002,
      chat_id: chat_id.clone(),
      content: "Message 2".to_string(),
      created_at: 1625097700,
      author_type: 0,
      author_id: "ai".to_string(),
      reply_message_id: None,
      metadata: None,
      is_sync: false,
    },
  ];

  // Insert messages
  upsert_chat_messages(db_conn, &messages).unwrap();

  // Test total_message_count
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let count = total_message_count(db_conn, &chat_id).unwrap();
  assert_eq!(count, 2);

  // Add one more message
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let additional_message = ChatMessageTable {
    message_id: 1003,
    chat_id: chat_id.clone(),
    content: "Message 3".to_string(),
    created_at: 1625097800,
    author_type: 1,
    author_id: "user_1".to_string(),
    reply_message_id: None,
    metadata: None,
    is_sync: false,
  };

  upsert_chat_messages(db_conn, &[additional_message]).unwrap();

  // Verify count increased
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let updated_count = total_message_count(db_conn, &chat_id).unwrap();
  assert_eq!(updated_count, 3);

  // Test count for non-existent chat
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let empty_count = total_message_count(db_conn, "non_existent_chat").unwrap();
  assert_eq!(empty_count, 0);
}

#[tokio::test]
async fn chat_message_select_message_test() {
  use_localhost_af_cloud().await;
  let test = EventIntegrationTest::new().await;
  test.sign_up_as_anon().await;

  let uid = test.user_manager.get_anon_user().await.unwrap().id;
  let db_conn = test.user_manager.db_connection(uid).unwrap();

  let chat_id = Uuid::new_v4().to_string();
  let message_id = 2001;

  // Create test message
  let message = ChatMessageTable {
    message_id,
    chat_id: chat_id.clone(),
    content: "This is a test message for select_message".to_string(),
    created_at: 1625097600,
    author_type: 1,
    author_id: "user_1".to_string(),
    reply_message_id: None,
    metadata: Some(r#"{"test_key": "test_value"}"#.to_string()),
    is_sync: false,
  };

  // Insert message
  upsert_chat_messages(db_conn, &[message]).unwrap();

  // Test select_message
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result = select_message(db_conn, message_id).unwrap();
  assert!(result.is_some());

  let retrieved_message = result.unwrap();
  assert_eq!(retrieved_message.message_id, message_id);
  assert_eq!(retrieved_message.chat_id, chat_id);
  assert_eq!(
    retrieved_message.content,
    "This is a test message for select_message"
  );
  assert_eq!(retrieved_message.author_id, "user_1");
  assert_eq!(
    retrieved_message.metadata,
    Some(r#"{"test_key": "test_value"}"#.to_string())
  );

  // Test select_message with non-existent ID
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let non_existent = select_message(db_conn, 9999).unwrap();
  assert!(non_existent.is_none());
}

#[tokio::test]
async fn chat_message_select_content_test() {
  use_localhost_af_cloud().await;
  let test = EventIntegrationTest::new().await;
  test.sign_up_as_anon().await;

  let uid = test.user_manager.get_anon_user().await.unwrap().id;
  let db_conn = test.user_manager.db_connection(uid).unwrap();

  let chat_id = Uuid::new_v4().to_string();
  let message_id = 3001;
  let message_content = "This is the content to retrieve";

  // Create test message
  let message = ChatMessageTable {
    message_id,
    chat_id: chat_id.clone(),
    content: message_content.to_string(),
    created_at: 1625097600,
    author_type: 1,
    author_id: "user_1".to_string(),
    reply_message_id: None,
    metadata: None,
    is_sync: false,
  };

  // Insert message
  upsert_chat_messages(db_conn, &[message]).unwrap();

  // Test select_message_content
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let content = select_message_content(db_conn, message_id).unwrap();
  assert!(content.is_some());
  assert_eq!(content.unwrap(), message_content);

  // Test with non-existent message
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let no_content = select_message_content(db_conn, 9999).unwrap();
  assert!(no_content.is_none());
}

#[tokio::test]
async fn chat_message_reply_test() {
  use_localhost_af_cloud().await;
  let test = EventIntegrationTest::new().await;
  test.sign_up_as_anon().await;

  let uid = test.user_manager.get_anon_user().await.unwrap().id;
  let db_conn = test.user_manager.db_connection(uid).unwrap();

  let chat_id = Uuid::new_v4().to_string();
  let question_id = 4001;
  let answer_id = 4002;

  // Create question and answer messages
  let question = ChatMessageTable {
    message_id: question_id,
    chat_id: chat_id.clone(),
    content: "What is the question?".to_string(),
    created_at: 1625097600,
    author_type: 1, // User
    author_id: "user_1".to_string(),
    reply_message_id: None,
    metadata: None,
    is_sync: false,
  };

  let answer = ChatMessageTable {
    message_id: answer_id,
    chat_id: chat_id.clone(),
    content: "This is the answer".to_string(),
    created_at: 1625097700,
    author_type: 0, // AI
    author_id: "ai".to_string(),
    reply_message_id: Some(question_id), // Link to question
    metadata: None,
    is_sync: false,
  };

  // Insert messages
  upsert_chat_messages(db_conn, &[question, answer]).unwrap();

  // Test select_message_where_match_reply_message_id
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result = select_answer_where_match_reply_message_id(db_conn, &chat_id, question_id).unwrap();

  assert!(result.is_some());
  let reply = result.unwrap();
  assert_eq!(reply.message_id, answer_id);
  assert_eq!(reply.content, "This is the answer");
  assert_eq!(reply.reply_message_id, Some(question_id));

  // Test with non-existent reply relation
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let no_reply = select_answer_where_match_reply_message_id(
    db_conn, &chat_id, 9999, // Non-existent question ID
  )
  .unwrap();

  assert!(no_reply.is_none());

  // Test with wrong chat_id
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let wrong_chat =
    select_answer_where_match_reply_message_id(db_conn, "wrong_chat_id", question_id).unwrap();

  assert!(wrong_chat.is_none());
}

#[tokio::test]
async fn chat_message_upsert_test() {
  use_localhost_af_cloud().await;
  let test = EventIntegrationTest::new().await;
  test.sign_up_as_anon().await;

  let uid = test.user_manager.get_anon_user().await.unwrap().id;
  let db_conn = test.user_manager.db_connection(uid).unwrap();

  let chat_id = Uuid::new_v4().to_string();
  let message_id = 5001;

  // Create initial message
  let message = ChatMessageTable {
    message_id,
    chat_id: chat_id.clone(),
    content: "Original content".to_string(),
    created_at: 1625097600,
    author_type: 1,
    author_id: "user_1".to_string(),
    reply_message_id: None,
    metadata: None,
    is_sync: false,
  };

  // Insert message
  upsert_chat_messages(db_conn, &[message]).unwrap();

  // Check original content
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let original = select_message(db_conn, message_id).unwrap().unwrap();
  assert_eq!(original.content, "Original content");

  // Create updated message with same ID but different content
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let updated_message = ChatMessageTable {
    message_id, // Same ID
    chat_id: chat_id.clone(),
    content: "Updated content".to_string(), // New content
    created_at: 1625097700,                 // Updated timestamp
    author_type: 1,
    author_id: "user_1".to_string(),
    reply_message_id: Some(1000), // Added reply ID
    metadata: Some(r#"{"updated": true}"#.to_string()),
    is_sync: false,
  };

  // Upsert message
  upsert_chat_messages(db_conn, &[updated_message]).unwrap();

  // Verify update
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result = select_message(db_conn, message_id).unwrap().unwrap();
  assert_eq!(result.content, "Updated content");
  assert_eq!(result.created_at, 1625097700);
  assert_eq!(result.reply_message_id, Some(1000));
  assert_eq!(result.metadata, Some(r#"{"updated": true}"#.to_string()));

  // Count should still be 1 (update, not insert)
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let count = total_message_count(db_conn, &chat_id).unwrap();
  assert_eq!(count, 1);
}

#[tokio::test]
async fn chat_message_select_with_large_dataset() {
  use_localhost_af_cloud().await;
  let test = EventIntegrationTest::new().await;
  test.sign_up_as_anon().await;

  let uid = test.user_manager.get_anon_user().await.unwrap().id;
  let db_conn = test.user_manager.db_connection(uid).unwrap();

  let chat_id = Uuid::new_v4().to_string();

  // Create 100 test messages with sequential IDs
  let mut messages = Vec::new();
  for i in 1..=100 {
    messages.push(ChatMessageTable {
      message_id: i * 100,
      chat_id: chat_id.clone(),
      content: format!("Message {}", i),
      created_at: 1625097600 + (i * 10), // Increasing timestamps
      author_type: if i % 2 == 0 { 0 } else { 1 }, // Alternate between AI and User
      author_id: if i % 2 == 0 {
        "ai".to_string()
      } else {
        "user_1".to_string()
      },
      reply_message_id: if i > 1 && i % 2 == 0 {
        Some((i - 1) * 100)
      } else {
        None
      }, // Even messages reply to previous message
      metadata: if i % 5 == 0 {
        Some(format!(r#"{{"index": {}}}"#, i))
      } else {
        None
      },
      is_sync: false,
    });
  }

  // Insert all 100 messages
  upsert_chat_messages(db_conn, &messages).unwrap();

  // Verify total count
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let count = total_message_count(db_conn, &chat_id).unwrap();
  assert_eq!(count, 100, "Should have 100 messages in the database");

  // Test 1: MessageCursor::Offset with small page size
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let page_size = 10;
  let result_offset =
    select_chat_messages(db_conn, &chat_id, page_size, MessageCursor::Offset(0)).unwrap();

  assert_eq!(
    result_offset.messages.len(),
    page_size as usize,
    "Should return exactly {page_size} messages"
  );
  assert!(
    result_offset.has_more,
    "Should have more messages available"
  );
  assert_eq!(result_offset.total_count, 100, "Total count should be 100");

  // Test 2: Pagination with offset
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result_page2 = select_chat_messages(
    db_conn,
    &chat_id,
    page_size,
    MessageCursor::Offset(page_size),
  )
  .unwrap();

  assert_eq!(result_page2.messages.len(), page_size as usize);
  assert!(
    result_page2.messages[0].message_id != result_offset.messages[0].message_id,
    "Second page should have different messages than first page"
  );

  // Test 3: MessageCursor::AfterMessageId (forward pagination)
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let middle_message_id = 5000; // Message ID from the middle
  let result_after = select_chat_messages(
    db_conn,
    &chat_id,
    page_size,
    MessageCursor::AfterMessageId(middle_message_id),
  )
  .unwrap();

  assert_eq!(result_after.messages.len(), page_size as usize);
  assert!(
    result_after
      .messages
      .iter()
      .all(|m| m.message_id > middle_message_id),
    "All messages should have ID greater than the cursor"
  );

  // Test 4: MessageCursor::BeforeMessageId (backward pagination)
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result_before = select_chat_messages(
    db_conn,
    &chat_id,
    page_size,
    MessageCursor::BeforeMessageId(middle_message_id),
  )
  .unwrap();

  assert_eq!(result_before.messages.len(), page_size as usize);
  assert!(
    result_before
      .messages
      .iter()
      .all(|m| m.message_id < middle_message_id),
    "All messages should have ID less than the cursor"
  );

  // Test 5: Large page size (retrieve all)
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result_all = select_chat_messages(
    db_conn,
    &chat_id,
    200, // More than we have
    MessageCursor::Offset(0),
  )
  .unwrap();

  assert_eq!(
    result_all.messages.len(),
    100,
    "Should return all 100 messages"
  );
  assert!(!result_all.has_more, "Should not have more messages");

  // Test 6: Empty result when using out of range cursor
  let db_conn = test.user_manager.db_connection(uid).unwrap();
  let result_out_of_range = select_chat_messages(
    db_conn,
    &chat_id,
    page_size,
    MessageCursor::AfterMessageId(10000), // After the last message
  )
  .unwrap();

  assert_eq!(
    result_out_of_range.messages.len(),
    0,
    "Should return no messages"
  );
  assert!(
    !result_out_of_range.has_more,
    "Should not have more messages"
  );
}