package dev.sonpx.loyalty.agent.core.memory;

import com.fasterxml.jackson.databind.ObjectMapper;
import dev.sonpx.loyalty.agent.core.config.AgentProperties;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;

import java.util.List;

import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

@ExtendWith(MockitoExtension.class)
class JdbcChatMemoryTest {

    @Mock
    private JdbcTemplate jdbcTemplate;

    @Mock
    private AgentProperties agentProperties;

    private final ObjectMapper objectMapper = new ObjectMapper();

    private JdbcChatMemory jdbcChatMemory;

    @BeforeEach
    void setUp() {
        lenient().when(agentProperties.getMaxMemoryTokens()).thenReturn(1000);
        lenient().when(agentProperties.getCharsPerToken()).thenReturn(2);
        jdbcChatMemory = new JdbcChatMemory(jdbcTemplate, agentProperties, objectMapper);
    }

    @Test
    @DisplayName("Should save UserMessage and SystemMessage to database")
    void testAddUserAndSystemMessage() {
        when(jdbcTemplate.queryForObject(anyString(), eq(Long.class), eq("conv-1"))).thenReturn(0L);

        List<Message> messages = List.of(
                new SystemMessage("System prompt"),
                new UserMessage("Hello AI")
        );

        jdbcChatMemory.add("conv-1", messages);

        verify(jdbcTemplate, times(2)).update(
                contains("INSERT INTO chat_memory_messages"),
                anyString(), eq("conv-1"), anyLong(), anyString(), any(), any(), any()
        );
    }

    @Test
    @DisplayName("Should serialize AssistantMessage tool calls into metadata JSON")
    void testAddAssistantMessageWithToolCalls() {
        when(jdbcTemplate.queryForObject(anyString(), eq(Long.class), eq("conv-2"))).thenReturn(0L);

        AssistantMessage.ToolCall toolCall = new AssistantMessage.ToolCall("call-123", "function", "get_campaign", "{\"id\":1}");
        AssistantMessage assistantMessage = AssistantMessage.builder()
                .content("Checking campaign...")
                .toolCalls(List.of(toolCall))
                .build();

        jdbcChatMemory.add("conv-2", List.of(assistantMessage));

        ArgumentCaptor<String> metadataCaptor = ArgumentCaptor.forClass(String.class);
        verify(jdbcTemplate).update(
                contains("INSERT INTO chat_memory_messages"),
                anyString(), eq("conv-2"), eq(1L), eq("ASSISTANT"), eq("Checking campaign..."), metadataCaptor.capture(), any()
        );

        String metadataJson = metadataCaptor.getValue();
        assertNotNull(metadataJson);
        assertTrue(metadataJson.contains("call-123"));
        assertTrue(metadataJson.contains("get_campaign"));
    }

    @Test
    @DisplayName("Should truncate ToolResponseMessage when response data exceeds 4000 chars")
    void testAddToolResponseMessageTruncation() {
        when(jdbcTemplate.queryForObject(anyString(), eq(Long.class), eq("conv-3"))).thenReturn(0L);

        String longResponseData = "A".repeat(5000);
        ToolResponseMessage.ToolResponse toolResponse = new ToolResponseMessage.ToolResponse("call-123", "get_campaign", longResponseData);
        ToolResponseMessage toolResponseMessage = ToolResponseMessage.builder()
                .responses(List.of(toolResponse))
                .build();

        jdbcChatMemory.add("conv-3", List.of(toolResponseMessage));

        ArgumentCaptor<String> metadataCaptor = ArgumentCaptor.forClass(String.class);
        verify(jdbcTemplate).update(
                contains("INSERT INTO chat_memory_messages"),
                anyString(), eq("conv-3"), eq(1L), eq("TOOL"), eq(""), metadataCaptor.capture(), any()
        );

        String metadataJson = metadataCaptor.getValue();
        assertNotNull(metadataJson);
        assertTrue(metadataJson.contains("... [truncated]"));
        assertFalse(metadataJson.contains(longResponseData));
    }

    @Test
    @DisplayName("Should deserialize AssistantMessage with tool calls correctly on get()")
    @SuppressWarnings("unchecked")
    void testGetAssistantMessageWithToolCalls() {
        String metadataJson = "{\"toolCalls\":[{\"id\":\"call-99\",\"type\":\"function\",\"name\":\"check_rule\",\"arguments\":\"{\\\"ruleId\\\":5}\"}]}";

        when(jdbcTemplate.query(anyString(), any(RowMapper.class), eq("conv-4")))
                .thenAnswer(invocation -> {
                    RowMapper<?> mapper = invocation.getArgument(1);
                    java.sql.ResultSet rs = mock(java.sql.ResultSet.class);
                    when(rs.getString("id")).thenReturn("msg-1");
                    when(rs.getString("message_type")).thenReturn("ASSISTANT");
                    when(rs.getString("content")).thenReturn("Found rule");
                    when(rs.getString("metadata")).thenReturn(metadataJson);

                    Object wrapper = mapper.mapRow(rs, 0);
                    return List.of(wrapper);
                });

        List<Message> messages = jdbcChatMemory.get("conv-4");

        assertEquals(1, messages.size());
        assertTrue(messages.get(0) instanceof AssistantMessage);
        AssistantMessage assistantMsg = (AssistantMessage) messages.get(0);
        assertEquals("Found rule", assistantMsg.getText());
        assertTrue(assistantMsg.hasToolCalls());
        assertEquals(1, assistantMsg.getToolCalls().size());
        assertEquals("call-99", assistantMsg.getToolCalls().get(0).id());
        assertEquals("check_rule", assistantMsg.getToolCalls().get(0).name());
    }

    @Test
    @DisplayName("Should deserialize ToolResponseMessage with responses correctly on get()")
    @SuppressWarnings("unchecked")
    void testGetToolResponseMessage() {
        String metadataJson = "{\"responses\":[{\"id\":\"call-99\",\"name\":\"check_rule\",\"responseData\":\"{\\\"success\\\":true}\"}]}";

        when(jdbcTemplate.query(anyString(), any(RowMapper.class), eq("conv-5")))
                .thenAnswer(invocation -> {
                    RowMapper<?> mapper = invocation.getArgument(1);
                    java.sql.ResultSet rs = mock(java.sql.ResultSet.class);
                    when(rs.getString("id")).thenReturn("msg-2");
                    when(rs.getString("message_type")).thenReturn("TOOL");
                    when(rs.getString("content")).thenReturn("");
                    when(rs.getString("metadata")).thenReturn(metadataJson);

                    Object wrapper = mapper.mapRow(rs, 0);
                    return List.of(wrapper);
                });

        List<Message> messages = jdbcChatMemory.get("conv-5");

        assertEquals(1, messages.size());
        assertTrue(messages.get(0) instanceof ToolResponseMessage);
        ToolResponseMessage toolMsg = (ToolResponseMessage) messages.get(0);
        assertEquals(1, toolMsg.getResponses().size());
        assertEquals("call-99", toolMsg.getResponses().get(0).id());
        assertEquals("{\"success\":true}", toolMsg.getResponses().get(0).responseData());
    }

    @Test
    @DisplayName("Should clear chat memory messages for a conversation")
    void testClear() {
        jdbcChatMemory.clear("conv-6");

        verify(jdbcTemplate).update("DELETE FROM chat_memory_messages WHERE conversation_id = ?", "conv-6");
    }
}
