package com.vincent.rsf.server.ai.service.impl.chat;

import com.vincent.rsf.server.ai.dto.AiChatMemoryDto;
import com.vincent.rsf.server.ai.dto.AiChatMessageDto;
import com.vincent.rsf.server.ai.entity.AiPrompt;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;

import java.util.List;

import static org.assertj.core.api.Assertions.assertThat;

class AiPromptMessageBuilderTest {

    private final AiPromptMessageBuilder aiPromptMessageBuilder = new AiPromptMessageBuilder();

    @Test
    void shouldMergeAllSystemContextIntoSingleLeadingSystemMessage() {
        AiChatMemoryDto memory = AiChatMemoryDto.builder()
                .memorySummary("这是摘要")
                .memoryFacts("这是事实")
                .build();
        AiPrompt prompt = new AiPrompt()
                .setSystemPrompt("你是助手")
                .setUserPromptTemplate("请回答：{{input}}");

        List<Message> messages = aiPromptMessageBuilder.buildPromptMessages(
                memory,
                List.of(
                        message("user", "第一问"),
                        message("assistant", "第一答"),
                        message("user", "第二问")
                ),
                prompt,
                null
        );

        assertThat(messages).hasSize(4);
        assertThat(messages.get(0)).isInstanceOf(SystemMessage.class);
        assertThat(messages.get(1)).isInstanceOf(UserMessage.class);
        assertThat(messages.get(2)).isInstanceOf(AssistantMessage.class);
        assertThat(messages.get(3)).isInstanceOf(UserMessage.class);
        assertThat(((SystemMessage) messages.get(0)).getText())
                .contains("你是助手")
                .contains("历史摘要:\n这是摘要")
                .contains("关键事实:\n这是事实");
        assertThat(((UserMessage) messages.get(3)).getText()).isEqualTo("请回答：第二问");
    }

    @Test
    void shouldMergePersistedAndMemoryMessages() {
        List<AiChatMessageDto> merged = aiPromptMessageBuilder.mergeMessages(
                List.of(message("user", "persisted")),
                List.of(message("assistant", "memory"))
        );

        assertThat(merged).hasSize(2);
        assertThat(merged.get(0).getContent()).isEqualTo("persisted");
        assertThat(merged.get(1).getContent()).isEqualTo("memory");
    }

    private AiChatMessageDto message(String role, String content) {
        AiChatMessageDto dto = new AiChatMessageDto();
        dto.setRole(role);
        dto.setContent(content);
        return dto;
    }
}