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

import com.vincent.rsf.server.ai.dto.AiChatRequest;
import com.vincent.rsf.server.ai.dto.AiChatMessageDto;
import com.vincent.rsf.server.ai.dto.AiResolvedConfig;
import com.vincent.rsf.server.ai.entity.AiParam;
import com.vincent.rsf.server.ai.entity.AiPrompt;
import com.vincent.rsf.server.ai.service.AiCallLogService;
import com.vincent.rsf.server.ai.service.AiChatMemoryService;
import com.vincent.rsf.server.ai.service.AiConfigResolverService;
import com.vincent.rsf.server.ai.service.AiParamService;
import com.vincent.rsf.server.ai.service.McpMountRuntimeFactory;
import com.vincent.rsf.server.ai.store.AiChatRateLimiter;
import com.vincent.rsf.server.ai.store.AiStreamStateStore;
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.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.List;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class AiChatOrchestratorTest {

    @Mock
    private AiConfigResolverService aiConfigResolverService;
    @Mock
    private AiChatMemoryService aiChatMemoryService;
    @Mock
    private AiParamService aiParamService;
    @Mock
    private McpMountRuntimeFactory mcpMountRuntimeFactory;
    @Mock
    private AiCallLogService aiCallLogService;
    @Mock
    private AiChatRateLimiter aiChatRateLimiter;
    @Mock
    private AiStreamStateStore aiStreamStateStore;
    @Mock
    private AiChatRuntimeAssembler aiChatRuntimeAssembler;
    @Mock
    private AiPromptMessageBuilder aiPromptMessageBuilder;
    @Mock
    private AiOpenAiChatModelFactory aiOpenAiChatModelFactory;
    @Mock
    private AiToolObservationService aiToolObservationService;
    @Mock
    private AiSseEventPublisher aiSseEventPublisher;
    @Mock
    private AiChatFailureHandler aiChatFailureHandler;

    @Test
    void shouldShortCircuitWhenRateLimited() {
        AiChatOrchestrator aiChatOrchestrator = new AiChatOrchestrator(
                aiConfigResolverService,
                aiChatMemoryService,
                aiParamService,
                mcpMountRuntimeFactory,
                aiCallLogService,
                aiChatRateLimiter,
                aiStreamStateStore,
                aiChatRuntimeAssembler,
                aiPromptMessageBuilder,
                aiOpenAiChatModelFactory,
                aiToolObservationService,
                aiSseEventPublisher,
                aiChatFailureHandler
        );
        AiChatRequest request = new AiChatRequest();
        request.setRequestId("req-1");
        request.setPromptCode("home.default");
        request.setMessages(List.of(message("user", "hello")));
        AiResolvedConfig config = AiResolvedConfig.builder()
                .promptCode("home.default")
                .aiParam(new AiParam().setModel("gpt-test"))
                .prompt(new AiPrompt().setName("default"))
                .mcpMounts(List.of())
                .build();
        when(aiConfigResolverService.resolve("home.default", 1L, null)).thenReturn(config);
        when(aiParamService.listChatModelOptions(1L)).thenReturn(List.of());
        when(aiChatRateLimiter.allowChatRequest(1L, 2L, "home.default")).thenReturn(false);
        when(aiChatFailureHandler.buildAiException(eq("AI_RATE_LIMITED"), any(), eq("RATE_LIMIT"), any(), eq(null)))
                .thenCallRealMethod();

        aiChatOrchestrator.executeStream(request, 2L, 1L, new SseEmitter(1000L));

        verify(aiChatFailureHandler).handleStreamFailure(any(), eq("req-1"), eq(null), eq(null), any(Long.class),
                eq(null), any(), eq(null), eq(0L), eq(0L), eq(null), eq(1L), eq(2L), eq("home.default"));
        verify(aiCallLogService, never()).startCallLog(any(), any(), any(), any(), any(), any(), any(), any(), any(), any());
    }

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