package com.infoepoch.pms.agent.platform.chat.application;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.action.InterruptionMetadata;
import com.alibaba.cloud.ai.graph.agent.Agent;
import com.alibaba.cloud.ai.graph.checkpoint.BaseCheckpointSaver;
import com.infoepoch.pms.agent.platform.agent.application.AgentExecutionService;
import com.infoepoch.pms.agent.platform.agent.application.AgentQueryService;
import com.infoepoch.pms.agent.platform.chat.application.dto.ApprovalDecisionRequest;
import com.infoepoch.pms.agent.platform.chat.application.dto.CreateSessionRequest;
import com.infoepoch.pms.agent.platform.chat.application.dto.MessageResponse;
import com.infoepoch.pms.agent.platform.chat.application.dto.PendingApprovalResponse;
import com.infoepoch.pms.agent.platform.chat.application.dto.SendMessageRequest;
import com.infoepoch.pms.agent.platform.chat.application.dto.SessionResponse;
import com.infoepoch.pms.agent.platform.chat.checkpoint.DatabaseCheckpointStore;
import com.infoepoch.pms.agent.platform.chat.domain.Conversation;
import com.infoepoch.pms.agent.platform.chat.domain.Session;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.ConversationCriteria;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.SessionCriteria;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.IConversationRepository;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.ISessionRepository;
import com.infoepoch.pms.agent.platform.chatbot.application.ChatbotQueryService;
import com.infoepoch.pms.agent.platform.chatbot.domain.Chatbot;
import com.infoepoch.pms.agent.platform.shared.exception.BusinessException;
import com.infoepoch.pms.agent.platform.shared.observation.ToolExecutionTraceBuffer;
import com.infoepoch.pms.agent.platform.shared.response.ResultCode;
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 reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class ChatServiceTest {

    @Mock
    private ISessionRepository sessionRepository;

    @Mock
    private IConversationRepository conversationRepository;

    @Mock
    private ChatbotQueryService chatbotQueryService;

    @Mock
    private AgentQueryService agentQueryService;

    @Mock
    private AgentExecutionService agentExecutionService;

    @Mock
    private DatabaseCheckpointStore checkpointStore;

    @Mock
    private SessionTitleService sessionTitleService;

    @Mock
    private ToolExecutionTraceBuffer toolExecutionTraceBuffer;

    @Mock
    private ApprovalStateService approvalStateService;

    @Mock
    private ApprovalIntentClassifier approvalIntentClassifier;

    @Mock
    private Agent mainAgent;

    @Test
    void shouldExposeInitialActiveThreadStateWhenCreatingSession() {
        Chatbot chatbot = chatbot();
        CreateSessionRequest req = new CreateSessionRequest();
        req.setUserId("user-1");
        req.setChatbotId("chatbot-1");

        when(chatbotQueryService.findEnabledById("chatbot-1")).thenReturn(Optional.of(chatbot));

        ChatService service = service();

        SessionResponse response = service.createSession(req);

        assertEquals(response.getSessionId(), response.getActiveThreadId());
        assertFalse(response.isContextResetPending());
        verify(sessionRepository).insert(any(Session.class));
    }

    @Test
    void shouldRejectMarkingContextResetWhenApprovalIsPending() {
        Session session = session("session-1", "session-1", false);
        when(sessionRepository.selectById("session-1")).thenReturn(session);
        when(approvalStateService.getPendingApproval("session-1", "session-1"))
                .thenReturn(PendingApprovalResponse.pending("interrupt-1", "session-1", null, null, List.of(), List.of()));

        ChatService service = service();

        BusinessException ex = assertThrows(BusinessException.class,
                () -> service.markContextResetPending("session-1"));

        assertEquals(ResultCode.PARAM_ERROR, ex.getResultCode());
    }

    @Test
    void shouldSwitchToNewThreadWhenSendingMessageAfterPendingReset() throws Exception {
        Session session = session("session-1", "session-1", true);
        Chatbot chatbot = chatbot();
        SendMessageRequest req = new SendMessageRequest();
        req.setContent("你好");
        List<Conversation> savedMessages = new ArrayList<>();

        when(sessionRepository.selectById("session-1")).thenReturn(session);
        when(chatbotQueryService.findDetailById("chatbot-1")).thenReturn(Optional.of(chatbot));
        when(agentExecutionService.buildChatbotAgent(any(Chatbot.class), any(BaseCheckpointSaver.class), anyString()))
                .thenReturn(mainAgent);
        when(mainAgent.invoke(anyString(), any(RunnableConfig.class))).thenReturn(Optional.of(
                new OverAllState(java.util.Map.of("messages", List.of(new AssistantMessage("收到"))))
        ));
        when(conversationRepository.selectCountByCriteria(argThat(criteria ->
                criteria != null && "session-1".equals(criteria.getSessionId())))).thenReturn(2);
        when(conversationRepository.insert(any(Conversation.class))).thenAnswer(invocation -> {
            Conversation conversation = invocation.getArgument(0);
            savedMessages.add(conversation);
            return true;
        });

        ChatService service = service();

        MessageResponse response = service.sendMessage("session-1", req);

        assertEquals("收到", response.getContent());
        assertEquals(2, savedMessages.size());
        String newThreadId = savedMessages.get(0).getThreadId();
        assertTrue(newThreadId.startsWith("session-1_thread_"));
        assertNotEquals("session-1", newThreadId);
        assertEquals(newThreadId, savedMessages.get(1).getThreadId());
        assertEquals(newThreadId, session.getActiveThreadId());
        assertFalse(session.isContextResetPending());
        verify(sessionRepository).update(session);

        ArgumentCaptor<RunnableConfig> configCaptor = ArgumentCaptor.forClass(RunnableConfig.class);
        verify(mainAgent).invoke(anyString(), configCaptor.capture());
        assertEquals(newThreadId, configCaptor.getValue().threadId().orElseThrow());
    }

    @Test
    void shouldQueryPendingApprovalWithActiveThreadIdWhenStreamingMessage() {
        Session session = session("session-1", "session-1_thread_active", false);
        SendMessageRequest req = new SendMessageRequest();
        req.setContent("继续");
        when(sessionRepository.selectById("session-1")).thenReturn(session);
        when(approvalStateService.getPendingApproval("session-1", "session-1_thread_active"))
                .thenReturn(PendingApprovalResponse.pending("interrupt-1", "session-1", null, "审批Agent", List.of(), List.of()));
        when(conversationRepository.insert(any(Conversation.class))).thenReturn(true);
        when(approvalIntentClassifier.classify("继续"))
                .thenReturn(new ApprovalIntentClassifier.Result(ApprovalIntent.NONE, 0.25D));

        ChatService service = service();

        List<?> events = service.streamMessage("session-1", req).collectList().block();

        assertEquals(1, events.size());
        verify(approvalStateService).getPendingApproval("session-1", "session-1_thread_active");
    }

    @Test
    void shouldCapturePendingApprovalWithResolvedActiveThreadWhenStreaming() throws Exception {
        Session session = session("session-1", "session-1_thread_active", false);
        Chatbot chatbot = chatbot();
        SendMessageRequest req = new SendMessageRequest();
        req.setContent("提交申请");
        InterruptionMetadata interruptionMetadata = InterruptionMetadata.builder("node-1", new OverAllState()).build();

        when(sessionRepository.selectById("session-1")).thenReturn(session);
        when(approvalStateService.getPendingApproval("session-1", "session-1_thread_active"))
                .thenReturn(PendingApprovalResponse.none());
        when(conversationRepository.insert(any(Conversation.class))).thenReturn(true);
        when(chatbotQueryService.findDetailById("chatbot-1")).thenReturn(Optional.of(chatbot));
        when(agentExecutionService.buildChatbotAgent(any(Chatbot.class), any(BaseCheckpointSaver.class), anyString()))
                .thenReturn(mainAgent);
        when(mainAgent.stream(anyString(), any(RunnableConfig.class))).thenReturn(Flux.just(interruptionMetadata));
        when(approvalStateService.capturePendingApproval(anyString(), eq("session-1"), eq("session-1_thread_active"),
                any(InterruptionMetadata.class), any(Chatbot.class)))
                .thenReturn(PendingApprovalResponse.pending("interrupt-1", "session-1", null, "审批Agent", List.of(), List.of()));

        ChatService service = service();

        List<?> events = service.streamMessage("session-1", req).collectList().block();

        assertEquals(1, events.size());
        verify(approvalStateService).capturePendingApproval(anyString(), eq("session-1"), eq("session-1_thread_active"),
                any(InterruptionMetadata.class), any(Chatbot.class));
    }

    @Test
    void shouldResumeApprovalDecisionWithActiveThreadId() throws Exception {
        Session session = session("session-1", "session-1_thread_active", false);
        ApprovalDecisionRequest req = new ApprovalDecisionRequest();
        req.setInterruptionId("interrupt-1");
        req.setDecision("approve");
        ApprovalStateService.ApprovalDecisionContext decisionContext =
                new ApprovalStateService.ApprovalDecisionContext(
                        PendingApprovalResponse.pending("interrupt-1", "session-1", null, "审批Agent", List.of(), List.of()),
                        "session-1_thread_active_subgraph_agent",
                        "interrupt-1",
                        InterruptionMetadata.builder("node-1", new OverAllState()).build()
                );
        when(sessionRepository.selectById("session-1")).thenReturn(session);
        when(approvalStateService.streamApprovalDecision("session-1", "session-1_thread_active",
                "interrupt-1", "approve", null))
                .thenReturn(decisionContext);
        when(chatbotQueryService.findDetailById("chatbot-1")).thenReturn(Optional.of(chatbot()));
        when(agentExecutionService.buildChatbotAgent(any(Chatbot.class), any(BaseCheckpointSaver.class), anyString()))
                .thenReturn(mainAgent);
        when(mainAgent.stream(any(Map.class), any(RunnableConfig.class))).thenReturn(Flux.empty());

        ChatService service = service();

        service.streamApprovalDecision("session-1", req).collectList().block();

        verify(approvalStateService).streamApprovalDecision("session-1", "session-1_thread_active",
                "interrupt-1", "approve", null);
    }

    @Test
    void shouldClearDerivedThreadGroupsWhenDeletingSession() {
        Session session = session("session-1", "session-1_thread_active", false);
        Conversation rootConversation = Conversation.restore(
                "conv-1", session, Conversation.ROLE_USER, "hello", null, null, null, LocalDateTime.now(), "session-1");
        Conversation activeThreadConversation = Conversation.restore(
                "conv-2", session, Conversation.ROLE_ASSISTANT, "world", null, null, null, LocalDateTime.now(),
                "session-1_thread_active");
        Conversation subgraphConversation = Conversation.restore(
                "conv-3", session, Conversation.ROLE_ASSISTANT, "sub", null, null, null, LocalDateTime.now(),
                "session-1_thread_active_subgraph_agent");
        when(sessionRepository.selectById("session-1")).thenReturn(session);
        when(conversationRepository.selectByCriteria(argThat(criteria ->
                criteria != null && "session-1".equals(criteria.getSessionId()))))
                .thenReturn(List.of(rootConversation, activeThreadConversation, subgraphConversation));

        ChatService service = service();

        service.deleteSession("session-1");

        ArgumentCaptor<List<String>> rootThreadCaptor = ArgumentCaptor.forClass(List.class);
        verify(approvalStateService).clearSessionState(eq("session-1"), rootThreadCaptor.capture());
        assertEquals(List.of("session-1", "session-1_thread_active"), rootThreadCaptor.getValue());
    }

    private ChatService service() {
        return new ChatService(
                sessionRepository,
                conversationRepository,
                chatbotQueryService,
                agentQueryService,
                agentExecutionService,
                checkpointStore,
                sessionTitleService,
                toolExecutionTraceBuffer,
                approvalStateService,
                approvalIntentClassifier
        );
    }

    private Session session(String id, String activeThreadId, boolean contextResetPending) {
        return Session.restore(
                id,
                "user-1",
                chatbot(),
                Session.STATUS_ACTIVE,
                null,
                activeThreadId,
                contextResetPending,
                LocalDateTime.now(),
                LocalDateTime.now()
        );
    }

    private Chatbot chatbot() {
        LocalDateTime now = LocalDateTime.now();
        return Chatbot.restore(
                "chatbot-1",
                "客服助手",
                "desc",
                null,
                null,
                List.of(),
                20,
                true,
                now,
                now,
                Set.of()
        );
    }
}
