/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.flink.agents.plan.compatibility;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.PythonFunction;
import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
import org.apache.flink.agents.plan.resourceprovider.PythonSerializableResourceProvider;
import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertTrue;

/** Utils for testing agent plan compatibility between Python and Java. */
public class CreateJavaAgentPlanFromJson {

    /**
     * Test creating Java AgentPlan from Python AgentPlan json.
     *
     * <p>The python agent plan json to be checked is generated by
     * flink-agents/python/flink_agents/plan/tests/compatibility/generate_agent_plan_json.py
     */
    public static void main(String[] args) throws IOException {
        String agentJsonFile = args[0];
        String json = Files.readString(Paths.get(agentJsonFile));
        AgentPlan agentPlan = new ObjectMapper().readValue(json, AgentPlan.class);
        assertEquals(5, agentPlan.getActions().size());

        String myEvent = "_my_event";
        String inputEvent = "_input_event";

        // Check the first action
        String testModule =
                "flink_agents.plan.tests.compatibility.python_agent_plan_compatibility_test_agent";
        assertTrue(agentPlan.getActions().containsKey("first_action"));
        Action firstAction = agentPlan.getActions().get("first_action");
        assertInstanceOf(PythonFunction.class, firstAction.getExec());
        PythonFunction firstActionFunction = (PythonFunction) firstAction.getExec();
        assertEquals(testModule, firstActionFunction.getModule());
        assertEquals(
                "PythonAgentPlanCompatibilityTestAgent.first_action",
                firstActionFunction.getQualName());
        assertEquals(List.of(inputEvent), firstAction.getListenEventTypes());

        // Check the second action
        assertTrue(agentPlan.getActions().containsKey("second_action"));
        Action secondAction = agentPlan.getActions().get("second_action");
        assertInstanceOf(PythonFunction.class, secondAction.getExec());
        PythonFunction secondActionFunc = (PythonFunction) secondAction.getExec();
        assertEquals(testModule, secondActionFunc.getModule());
        assertEquals(
                "PythonAgentPlanCompatibilityTestAgent.second_action",
                secondActionFunc.getQualName());
        assertEquals(List.of(inputEvent, myEvent), secondAction.getListenEventTypes());

        // Check the built-in actions
        assertTrue(agentPlan.getActions().containsKey("chat_model_action"));
        Action chatModelAction = agentPlan.getActions().get("chat_model_action");
        assertInstanceOf(PythonFunction.class, chatModelAction.getExec());
        PythonFunction processChatRequestFunc = (PythonFunction) chatModelAction.getExec();
        assertEquals(
                "flink_agents.plan.actions.chat_model_action", processChatRequestFunc.getModule());
        assertEquals("process_chat_request_or_tool_response", processChatRequestFunc.getQualName());
        String chatRequestEvent = "_chat_request_event";
        String toolResponseEvent = "_tool_response_event";
        assertEquals(
                List.of(chatRequestEvent, toolResponseEvent),
                chatModelAction.getListenEventTypes());

        assertTrue(agentPlan.getActions().containsKey("tool_call_action"));
        Action toolCallAction = agentPlan.getActions().get("tool_call_action");
        assertInstanceOf(PythonFunction.class, toolCallAction.getExec());
        PythonFunction processToolRequestFunc = (PythonFunction) toolCallAction.getExec();
        assertEquals(
                "flink_agents.plan.actions.tool_call_action", processToolRequestFunc.getModule());
        assertEquals("process_tool_request", processToolRequestFunc.getQualName());
        String toolRequestEvent = "_tool_request_event";
        assertEquals(List.of(toolRequestEvent), toolCallAction.getListenEventTypes());

        assertTrue(agentPlan.getActions().containsKey("context_retrieval_action"));
        Action contextRetrievalAction = agentPlan.getActions().get("context_retrieval_action");
        assertInstanceOf(PythonFunction.class, contextRetrievalAction.getExec());
        PythonFunction processContextRetrievalRequestFunc =
                (PythonFunction) contextRetrievalAction.getExec();
        assertEquals(
                "flink_agents.plan.actions.context_retrieval_action",
                processContextRetrievalRequestFunc.getModule());
        assertEquals(
                "process_context_retrieval_request",
                processContextRetrievalRequestFunc.getQualName());
        String contextRetrievalRequestEvent = "_context_retrieval_request_event";
        assertEquals(
                List.of(contextRetrievalRequestEvent),
                contextRetrievalAction.getListenEventTypes());

        // Check event trigger actions
        Map<String, List<Action>> actionsByEvent = agentPlan.getActionsByEvent();
        assertEquals(6, actionsByEvent.size());
        assertTrue(actionsByEvent.containsKey(inputEvent));
        assertTrue(actionsByEvent.containsKey(myEvent));
        assertTrue(actionsByEvent.containsKey(chatRequestEvent));
        assertTrue(actionsByEvent.containsKey(toolRequestEvent));
        assertTrue(actionsByEvent.containsKey(toolResponseEvent));
        assertTrue(actionsByEvent.containsKey(contextRetrievalRequestEvent));
        assertEquals(
                List.of(firstAction, secondAction), agentPlan.getActionsByEvent().get(inputEvent));
        assertEquals(List.of(secondAction), agentPlan.getActionsByEvent().get(myEvent));
        assertEquals(List.of(chatModelAction), actionsByEvent.get(chatRequestEvent));
        assertEquals(List.of(toolCallAction), actionsByEvent.get(toolRequestEvent));
        assertEquals(List.of(chatModelAction), actionsByEvent.get(toolResponseEvent));
        assertEquals(
                List.of(contextRetrievalAction), actionsByEvent.get(contextRetrievalRequestEvent));

        // Check resource providers
        Map<String, Object> kwargs = new HashMap<>();
        kwargs.put("name", "chat_model");
        kwargs.put("prompt", "prompt");
        kwargs.put("tools", List.of("add"));
        kwargs.put("connection", "mock_connection");
        kwargs.put("model", "mock-model");
        ResourceDescriptor chatModelDescriptor =
                new ResourceDescriptor(
                        "flink_agents.plan.tests.compatibility.python_agent_plan_compatibility_test_agent",
                        "MockChatModel",
                        kwargs);
        PythonResourceProvider resourceProvider =
                new PythonResourceProvider(
                        "chat_model", ResourceType.CHAT_MODEL, chatModelDescriptor);

        Map<String, Object> serialized = new HashMap<>();

        // construct arguments schema
        Map<String, String> a = new HashMap<>();
        a.put("description", "The first operand");
        a.put("title", "A");
        a.put("type", "integer");

        Map<String, String> b = new HashMap<>();
        b.put("description", "The second operand");
        b.put("title", "B");
        b.put("type", "integer");

        Map<String, Object> properties = new HashMap<>();
        properties.put("a", a);
        properties.put("b", b);

        Map<String, Object> argsSchema = new HashMap<>();
        argsSchema.put("properties", properties);
        argsSchema.put("required", List.of("a", "b"));
        argsSchema.put("title", "add");
        argsSchema.put("type", "object");

        Map<String, Object> metadata = new HashMap<>();
        metadata.put("name", "add");
        metadata.put("description", "Calculate the sum of a and b.\n");
        metadata.put("args_schema", argsSchema);

        serialized.put("metadata", metadata);

        Map<String, String> func = new HashMap<>();
        func.put("func_type", "PythonFunction");
        func.put(
                "module",
                "flink_agents.plan.tests.compatibility.python_agent_plan_compatibility_test_agent");
        func.put("qualname", "PythonAgentPlanCompatibilityTestAgent.add");
        serialized.put("func", func);
        PythonSerializableResourceProvider serializableResourceProvider =
                new PythonSerializableResourceProvider(
                        "add",
                        ResourceType.TOOL,
                        "flink_agents.plan.tools.function_tool",
                        "FunctionTool",
                        serialized);

        Map<ResourceType, Map<String, ResourceProvider>> resourceProviders = new HashMap<>();
        Map<String, ResourceProvider> chatModels = new HashMap<>();
        Map<String, ResourceProvider> tools = new HashMap<>();
        chatModels.put("chat_model", resourceProvider);
        tools.put("add", serializableResourceProvider);
        resourceProviders.put(ResourceType.CHAT_MODEL, chatModels);
        resourceProviders.put(ResourceType.TOOL, tools);
        assertEquals(resourceProviders, agentPlan.getResourceProviders());
    }
}
