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

import dev.sonpx.loyalty.agent.core.config.AgentProperties;
import dev.sonpx.loyalty.agent.exception.UnrecoverableToolExecutionException;
import dev.sonpx.loyalty.agent.service.AgentEventPublisher;
import dev.sonpx.loyalty.agent.service.ToolResultPresentationProcessor;
import io.modelcontextprotocol.client.McpSyncClient;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.beans.factory.ObjectProvider;

import java.io.EOFException;
import java.util.List;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.Mockito.*;

@ExtendWith(MockitoExtension.class)
public class ToolInterceptorTest {

    @Mock
    private ToolCallbackProvider toolCallbackProvider;

    @Mock
    private AgentEventPublisher eventPublisher;

    @Mock
    private ToolResultPresentationProcessor toolResultProcessor;

    @Mock
    private AgentProperties agentProperties;

    @Mock
    private ObjectProvider<List<McpSyncClient>> mcpSyncClientsProvider;

    @Mock
    private McpSyncClient mcpSyncClient;

    @Mock
    private ToolCallback mockTool;

    @Mock
    private ToolDefinition mockToolDefinition;

    private ToolInterceptor toolInterceptor;

    @BeforeEach
    void setUp() {
        lenient().when(mockToolDefinition.name()).thenReturn("test_tool");
        lenient().when(mockTool.getToolDefinition()).thenReturn(mockToolDefinition);
        lenient().when(toolCallbackProvider.getToolCallbacks()).thenReturn(new ToolCallback[]{mockTool});

        toolInterceptor = new ToolInterceptor(
                List.of(toolCallbackProvider),
                eventPublisher,
                toolResultProcessor,
                agentProperties,
                mcpSyncClientsProvider
        );
    }

    @Test
    void testSuccessfulToolExecution() {
        when(mockTool.call("input")).thenReturn("{\"status\":\"ok\"}");
        when(toolResultProcessor.process("{\"status\":\"ok\"}")).thenReturn("{\"status\":\"ok\"}");
        when(agentProperties.getMaxToolResultChars()).thenReturn(1000);

        List<ToolCallback> wrapped = toolInterceptor.getWrappedTools("session1", "msg1");
        assertThat(wrapped).hasSize(1);

        String response = wrapped.get(0).call("input");
        assertThat(response).isEqualTo("{\"status\":\"ok\"}");

        verify(eventPublisher).toolRunning("session1", "msg1", "test_tool");
        verify(eventPublisher).toolDone("session1", "msg1", "test_tool");
    }

    @Test
    void testAutoReconnectAndRetryOnConnectionDrop() {
        when(mcpSyncClientsProvider.getIfAvailable()).thenReturn(List.of(mcpSyncClient));
        when(mockTool.call("input"))
                .thenThrow(new RuntimeException("404 Session Not Found"))
                .thenReturn("{\"recovered\":true}");

        when(toolResultProcessor.process("{\"recovered\":true}")).thenReturn("{\"recovered\":true}");
        when(agentProperties.getMaxToolResultChars()).thenReturn(1000);

        List<ToolCallback> wrapped = toolInterceptor.getWrappedTools("session1", "msg1");
        String response = wrapped.get(0).call("input");

        assertThat(response).isEqualTo("{\"recovered\":true}");
        verify(mcpSyncClient, times(1)).initialize();
        verify(mockTool, times(2)).call("input");
    }

    @Test
    void testUnrecoverableErrorThrowsCircuitBreakerException() {
        when(mcpSyncClientsProvider.getIfAvailable()).thenReturn(List.of(mcpSyncClient));
        when(mockTool.call("input")).thenThrow(new RuntimeException(new EOFException("Connection reset by peer")));

        List<ToolCallback> wrapped = toolInterceptor.getWrappedTools("session1", "msg1");

        assertThatThrownBy(() -> wrapped.get(0).call("input"))
                .isInstanceOf(UnrecoverableToolExecutionException.class)
                .hasMessageContaining("Không thể khôi phục kết nối MCP khi gọi tool test_tool");

        verify(eventPublisher).toolDone("session1", "msg1", "test_tool");
    }

    @Test
    void testIsConnectionOrSessionErrorDetection() {
        assertThat(toolInterceptor.isConnectionOrSessionError(new EOFException("EOF"))).isTrue();
        assertThat(toolInterceptor.isConnectionOrSessionError(new RuntimeException("HTTP 404 Not Found"))).isTrue();
        assertThat(toolInterceptor.isConnectionOrSessionError(new RuntimeException("Session not found"))).isTrue();
        assertThat(toolInterceptor.isConnectionOrSessionError(new IllegalArgumentException("Invalid input"))).isFalse();
    }
}
