/*
 * Copyright 1999-2025 Alibaba Group Holding Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.alibaba.nacos.client.ai.remote;

import com.alibaba.nacos.api.PropertyKeyConst;
import com.alibaba.nacos.api.ability.constant.AbilityKey;
import com.alibaba.nacos.api.ability.constant.AbilityStatus;
import com.alibaba.nacos.api.ai.model.mcp.McpServerBasicInfo;
import com.alibaba.nacos.api.ai.model.mcp.McpServerDetailInfo;
import com.alibaba.nacos.api.ai.model.mcp.McpToolSpecification;
import com.alibaba.nacos.api.ai.model.mcp.registry.ServerVersionDetail;
import com.alibaba.nacos.api.ai.model.prompt.Prompt;
import com.alibaba.nacos.api.ai.remote.AiRemoteConstants;
import com.alibaba.nacos.api.ai.remote.request.McpServerEndpointRequest;
import com.alibaba.nacos.api.ai.remote.request.QueryMcpServerRequest;
import com.alibaba.nacos.api.ai.remote.request.QueryPromptRequest;
import com.alibaba.nacos.api.ai.remote.request.ReleaseMcpServerRequest;
import com.alibaba.nacos.api.ai.remote.response.McpServerEndpointResponse;
import com.alibaba.nacos.api.ai.remote.response.QueryMcpServerResponse;
import com.alibaba.nacos.api.ai.remote.response.QueryPromptResponse;
import com.alibaba.nacos.api.ai.remote.response.ReleaseMcpServerResponse;
import com.alibaba.nacos.api.exception.NacosException;
import com.alibaba.nacos.api.exception.runtime.NacosRuntimeException;
import com.alibaba.nacos.api.naming.remote.request.InstanceRequest;
import com.alibaba.nacos.api.naming.remote.response.InstanceResponse;
import com.alibaba.nacos.api.remote.request.Request;
import com.alibaba.nacos.api.remote.response.ErrorResponse;
import com.alibaba.nacos.api.remote.response.Response;
import com.alibaba.nacos.client.address.AbstractServerListManager;
import com.alibaba.nacos.client.ai.cache.NacosAgentCardCacheHolder;
import com.alibaba.nacos.client.ai.cache.NacosMcpServerCacheHolder;
import com.alibaba.nacos.client.ai.remote.redo.AiGrpcRedoService;
import com.alibaba.nacos.client.env.NacosClientProperties;
import com.alibaba.nacos.client.security.SecurityProxy;
import com.alibaba.nacos.common.remote.client.RpcClient;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
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 java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.UUID;

import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class AiGrpcClientTest {
    
    @Mock
    private RpcClient rpcClient;
    
    @Mock
    private AbstractServerListManager serverListManager;
    
    @Mock
    private AiGrpcRedoService redoService;
    
    @Mock
    private SecurityProxy securityProxy;
    
    @Mock
    private NacosMcpServerCacheHolder mcpServerCacheHolder;
    
    @Mock
    private NacosAgentCardCacheHolder agentCardCacheHolder;
    
    AiGrpcClient aiGrpcClient;
    
    @BeforeEach
    void setUp() {
        Properties properties = new Properties();
        properties.put(PropertyKeyConst.SERVER_ADDR, "127.0.0.1");
        NacosClientProperties clientProperties = NacosClientProperties.PROTOTYPE.derive(properties);
        aiGrpcClient = new AiGrpcClient("test", clientProperties);
    }
    
    @AfterEach
    void tearDown() throws NacosException {
        aiGrpcClient.shutdown();
    }
    
    @Test
    void start() throws NacosException {
        assertDoesNotThrow(() -> aiGrpcClient.start(mcpServerCacheHolder, agentCardCacheHolder));
    }
    
    @Test
    void queryMcpServer() throws NacosException, NoSuchFieldException, IllegalAccessException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        McpServerDetailInfo mcpServerDetailInfo = new McpServerDetailInfo();
        QueryMcpServerResponse response = new QueryMcpServerResponse();
        response.setMcpServerDetailInfo(mcpServerDetailInfo);
        when(rpcClient.request(any(QueryMcpServerRequest.class))).thenReturn(response);
        McpServerDetailInfo actual = aiGrpcClient.queryMcpServer("test", "1.0.0");
        assertEquals(mcpServerDetailInfo, actual);
    }
    
    @Test
    void queryMcpServerWithErrorCode() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        Response response = ErrorResponse.build(NacosException.INVALID_PARAM, "test");
        when(rpcClient.request(any(QueryMcpServerRequest.class))).thenReturn(response);
        assertThrows(NacosException.class, () -> aiGrpcClient.queryMcpServer("test", "1.0.0"));
    }
    
    @Test
    void queryMcpServerWithNoRight() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        Response response = ErrorResponse.build(NacosException.NO_RIGHT, "test");
        when(rpcClient.request(any(QueryMcpServerRequest.class))).thenReturn(response);
        assertThrows(NacosException.class, () -> aiGrpcClient.queryMcpServer("test", "1.0.0"));
        verify(securityProxy).reLogin();
    }
    
    @Test
    void queryMcpServerWithUnExpectedResponse() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        ReleaseMcpServerResponse response = new ReleaseMcpServerResponse();
        when(rpcClient.request(any(QueryMcpServerRequest.class))).thenReturn(response);
        assertThrows(NacosException.class, () -> aiGrpcClient.queryMcpServer("test", "1.0.0"));
    }
    
    @Test
    void queryMcpServerWithException() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        when(rpcClient.request(any(QueryMcpServerRequest.class))).thenThrow(new RuntimeException("test"));
        assertThrows(NacosException.class, () -> aiGrpcClient.queryMcpServer("test", "1.0.0"));
    }
    
    @Test
    void queryPromptShouldBuildRequestWithVersionLabelMd5() throws Exception {
        injectMock();
        when(securityProxy.getIdentityContext(any())).thenReturn(new HashMap<>());
        QueryPromptResponse response = new QueryPromptResponse();
        Prompt prompt = new Prompt();
        prompt.setPromptKey("p1");
        prompt.setVersion("1.0.0");
        response.setPromptInfo(prompt);
        when(rpcClient.request(any(QueryPromptRequest.class))).thenReturn(response);
        
        Prompt actual = aiGrpcClient.queryPrompt("p1", "1.0.0", "prod", "m1");
        
        assertEquals("p1", actual.getPromptKey());
        ArgumentCaptor<QueryPromptRequest> reqCaptor = ArgumentCaptor.forClass(QueryPromptRequest.class);
        verify(rpcClient).request(reqCaptor.capture());
        QueryPromptRequest captured = reqCaptor.getValue();
        assertEquals("test", captured.getNamespaceId());
        assertEquals("p1", captured.getPromptKey());
        assertEquals("1.0.0", captured.getVersion());
        assertEquals("prod", captured.getLabel());
        assertEquals("m1", captured.getMd5());
    }
    
    @Test
    void requestToServerShouldInjectSecurityHeaderForPromptRequest() throws Exception {
        injectMock();
        Map<String, String> headers = new HashMap<>();
        headers.put("k", "v");
        when(securityProxy.getIdentityContext(any())).thenReturn(headers);
        QueryPromptResponse response = new QueryPromptResponse();
        response.setPromptInfo(new Prompt("p1", "1.0.0", "hello"));
        when(rpcClient.request(any(QueryPromptRequest.class))).thenReturn(response);
        
        aiGrpcClient.queryPrompt("p1", "1.0.0", null, null);
        
        verify(securityProxy).getIdentityContext(any());
        ArgumentCaptor<QueryPromptRequest> reqCaptor = ArgumentCaptor.forClass(QueryPromptRequest.class);
        verify(rpcClient).request(reqCaptor.capture());
        assertEquals("v", reqCaptor.getValue().getHeader("k"));
    }
    
    @Test
    void queryPromptShouldReloginWhen403() throws Exception {
        injectMock();
        when(securityProxy.getIdentityContext(any())).thenReturn(new HashMap<>());
        Response response = ErrorResponse.build(NacosException.NO_RIGHT, "no right");
        when(rpcClient.request(any(QueryPromptRequest.class))).thenReturn(response);
        
        assertThrows(NacosException.class, () -> aiGrpcClient.queryPrompt("p1", null, null, null));
        verify(securityProxy).reLogin();
    }
    
    @Test
    void queryPromptShouldThrowWhenResponseTypeMismatch() throws Exception {
        injectMock();
        when(securityProxy.getIdentityContext(any())).thenReturn(new HashMap<>());
        QueryMcpServerResponse response = new QueryMcpServerResponse();
        when(rpcClient.request(any(QueryPromptRequest.class))).thenReturn(response);
        
        assertThrows(NacosException.class, () -> aiGrpcClient.queryPrompt("p1", null, null, null));
    }
    
    @Test
    void releaseMcpServer() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        McpServerBasicInfo serverSpec = new McpServerBasicInfo();
        serverSpec.setName("test");
        serverSpec.setVersionDetail(new ServerVersionDetail());
        serverSpec.getVersionDetail().setVersion("1.0.0");
        String id = UUID.randomUUID().toString();
        ReleaseMcpServerResponse response = new ReleaseMcpServerResponse();
        response.setMcpId(id);
        when(rpcClient.request(any(ReleaseMcpServerRequest.class))).thenReturn(response);
        assertEquals(id, aiGrpcClient.releaseMcpServer(serverSpec, new McpToolSpecification(), null));
    }
    
    @Test
    void registerMcpServerEndpoint() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        McpServerEndpointResponse response = new McpServerEndpointResponse();
        response.setType(AiRemoteConstants.REGISTER_ENDPOINT);
        when(rpcClient.request(any(McpServerEndpointRequest.class))).thenReturn(response);
        aiGrpcClient.registerMcpServerEndpoint("test", "127.0.0.1", 8080, "1.0.0");
        verify(redoService).cachedMcpServerEndpointForRedo("test", "127.0.0.1", 8080, "1.0.0");
        verify(redoService).mcpServerEndpointRegistered("test");
    }
    
    @Test
    void deregisterMcpServerEndpoint() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        McpServerEndpointResponse response = new McpServerEndpointResponse();
        response.setType(AiRemoteConstants.DE_REGISTER_ENDPOINT);
        when(rpcClient.request(any(McpServerEndpointRequest.class))).thenReturn(response);
        aiGrpcClient.deregisterMcpServerEndpoint("test", "127.0.0.1", 8080);
        verify(redoService).mcpServerEndpointDeregister("test");
        verify(redoService).mcpServerEndpointDeregistered("test");
    }
    
    @Test
    void subscribeMcpServer() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        McpServerDetailInfo mcpServerDetailInfo = new McpServerDetailInfo();
        QueryMcpServerResponse response = new QueryMcpServerResponse();
        response.setMcpServerDetailInfo(mcpServerDetailInfo);
        when(rpcClient.request(any(QueryMcpServerRequest.class))).thenReturn(response);
        assertEquals(mcpServerDetailInfo, aiGrpcClient.subscribeMcpServer("test", null));
        verify(mcpServerCacheHolder).processMcpServerDetailInfo(mcpServerDetailInfo);
        verify(mcpServerCacheHolder).addMcpServerUpdateTask("test", null);
    }
    
    @Test
    void subscribeMcpServerAlreadySubscribed() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        McpServerDetailInfo mcpServerDetailInfo = new McpServerDetailInfo();
        when(mcpServerCacheHolder.getMcpServer("test", null)).thenReturn(mcpServerDetailInfo);
        assertEquals(mcpServerDetailInfo, aiGrpcClient.subscribeMcpServer("test", null));
        verify(rpcClient, never()).request(any(QueryMcpServerRequest.class));
        verify(mcpServerCacheHolder, never()).processMcpServerDetailInfo(mcpServerDetailInfo);
        verify(mcpServerCacheHolder, never()).addMcpServerUpdateTask("test", null);
    }
    
    @Test
    void unsubscribeMcpServer() throws NoSuchFieldException, IllegalAccessException, NacosException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.SUPPORTED);
        aiGrpcClient.unsubscribeMcpServer("test", null);
        verify(mcpServerCacheHolder).removeMcpServerUpdateTask("test", null);
    }
    
    @Test
    void queryMcpServerWithFeatureDisabled() throws NoSuchFieldException, IllegalAccessException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.NOT_SUPPORTED);
        assertThrows(NacosRuntimeException.class, () -> aiGrpcClient.queryMcpServer("test", "1.0.0"));
    }
    
    @Test
    void releaseMcpServerWithFeatureDisabled() throws NoSuchFieldException, IllegalAccessException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.NOT_SUPPORTED);
        McpServerBasicInfo serverSpec = new McpServerBasicInfo();
        serverSpec.setName("test");
        serverSpec.setVersionDetail(new ServerVersionDetail());
        serverSpec.getVersionDetail().setVersion("1.0.0");
        assertThrows(NacosRuntimeException.class, () -> aiGrpcClient.releaseMcpServer(serverSpec, null, null));
    }
    
    @Test
    void registerMcpServerEndpointWithFeatureDisabled() throws NoSuchFieldException, IllegalAccessException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.NOT_SUPPORTED);
        assertThrows(NacosRuntimeException.class,
                () -> aiGrpcClient.registerMcpServerEndpoint("test", "127.0.0.1", 8080, "1.0.0"));
    }
    
    @Test
    void deregisterMcpServerEndpointWithFeatureDisabled() throws NoSuchFieldException, IllegalAccessException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.NOT_SUPPORTED);
        assertThrows(NacosRuntimeException.class,
                () -> aiGrpcClient.deregisterMcpServerEndpoint("test", "127.0.0.1", 8080));
    }
    
    @Test
    void subscribeMcpServerWithFeatureDisabled() throws NoSuchFieldException, IllegalAccessException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.NOT_SUPPORTED);
        assertThrows(NacosRuntimeException.class, () -> aiGrpcClient.subscribeMcpServer("test", null));
    }
    
    @Test
    void unsubscribeMcpServerWithFeatureDisabled() throws NoSuchFieldException, IllegalAccessException {
        injectMock();
        when(rpcClient.getConnectionAbility(AbilityKey.SERVER_MCP_REGISTRY)).thenReturn(AbilityStatus.NOT_SUPPORTED);
        assertThrows(NacosRuntimeException.class, () -> aiGrpcClient.unsubscribeMcpServer("test", null));
    }
    
    @Test
    void isEnable() throws NoSuchFieldException, IllegalAccessException {
        injectMock();
        assertFalse(aiGrpcClient.isEnable());
        when(rpcClient.isRunning()).thenReturn(true);
        assertTrue(aiGrpcClient.isEnable());
    }
    
    @Test
    void requestToServerWithoutMcpRequest() throws NoSuchMethodException, NoSuchFieldException, IllegalAccessException {
        Method method = AiGrpcClient.class.getDeclaredMethod("requestToServer", Request.class, Class.class);
        method.setAccessible(true);
        injectMock();
        try {
            method.invoke(aiGrpcClient, new InstanceRequest(), InstanceResponse.class);
        } catch (InvocationTargetException e) {
            Throwable targetException = e.getTargetException();
            assertInstanceOf(NacosException.class, targetException);
            assertEquals(400, ((NacosException) targetException).getErrCode());
        }
    }
    
    private void injectMock() throws NoSuchFieldException, IllegalAccessException {
        Field field = AiGrpcClient.class.getDeclaredField("rpcClient");
        field.setAccessible(true);
        RpcClient autoRpcClient = (RpcClient) field.get(aiGrpcClient);
        field.set(aiGrpcClient, rpcClient);
        
        field = AiGrpcClient.class.getDeclaredField("serverListManager");
        field.setAccessible(true);
        AbstractServerListManager autoServerListManager = (AbstractServerListManager) field.get(aiGrpcClient);
        field.set(aiGrpcClient, serverListManager);
        
        field = AiGrpcClient.class.getDeclaredField("redoService");
        field.setAccessible(true);
        AiGrpcRedoService autoRedoService = (AiGrpcRedoService) field.get(aiGrpcClient);
        field.set(aiGrpcClient, redoService);
        
        field = AiGrpcClient.class.getDeclaredField("securityProxy");
        field.setAccessible(true);
        field.set(aiGrpcClient, securityProxy);
        field = AiGrpcClient.class.getDeclaredField("mcpServerCacheHolder");
        field.setAccessible(true);
        field.set(aiGrpcClient, mcpServerCacheHolder);
        
        try {
            autoRpcClient.shutdown();
            autoServerListManager.shutdown();
            autoRedoService.shutdown();
        } catch (NacosException ignored) {
        }
    }
}