/*
 * 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.runtime.python.utils;

import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.prompt.Prompt;
import org.apache.flink.agents.api.resource.Resource;
import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
import org.apache.flink.agents.api.tools.Tool;
import org.apache.flink.agents.api.vectorstores.Document;
import org.apache.flink.agents.api.vectorstores.VectorStoreQuery;
import org.apache.flink.agents.api.vectorstores.VectorStoreQueryResult;
import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class PythonResourceAdapterImpl implements PythonResourceAdapter {

    static final String PYTHON_IMPORTS = "from flink_agents.runtime import python_java_utils";

    static final String JAVA_RESOURCE = "j_resource";

    static final String JAVA_RESOURCE_ADAPTER = "j_resource_adapter";

    static final String RESOURCE_CONTEXT_KEY = "resource_context";

    static final String PYTHON_MODULE_PREFIX = "python_java_utils.";

    static final String GET_RESOURCE_CONTEXT = PYTHON_MODULE_PREFIX + "get_resource_context";

    static final String CALL_METHOD = PYTHON_MODULE_PREFIX + "call_method";

    static final String CREATE_RESOURCE = PYTHON_MODULE_PREFIX + "create_resource";

    static final String FROM_JAVA_RESOURCE = PYTHON_MODULE_PREFIX + "from_java_resource";

    static final String FROM_JAVA_TOOL = PYTHON_MODULE_PREFIX + "from_java_tool";

    static final String FROM_JAVA_PROMPT = PYTHON_MODULE_PREFIX + "from_java_prompt";

    static final String FROM_JAVA_CHAT_MESSAGE = PYTHON_MODULE_PREFIX + "from_java_chat_message";

    static final String TO_JAVA_CHAT_MESSAGE = PYTHON_MODULE_PREFIX + "update_java_chat_message";

    static final String FROM_JAVA_DOCUMENT = PYTHON_MODULE_PREFIX + "from_java_document";

    static final String UPDATE_JAVA_DOCUMENT = PYTHON_MODULE_PREFIX + "update_java_document";

    static final String FROM_JAVA_VECTOR_STORE_QUERY =
            PYTHON_MODULE_PREFIX + "from_java_vector_store_query";

    static final String GET_PYTHON_TOOL_METADATA =
            PYTHON_MODULE_PREFIX + "get_python_tool_metadata";

    static final String INVOKE_PYTHON_TOOL = PYTHON_MODULE_PREFIX + "invoke_python_tool";

    private final ResourceContext resourceContext;
    private final PythonInterpreter interpreter;
    private final JavaResourceAdapter javaResourceAdapter;
    private PyObject pythonResourceContext;

    public PythonResourceAdapterImpl(
            ResourceContext resourceContext,
            PythonInterpreter interpreter,
            JavaResourceAdapter javaResourceAdapter) {
        this.resourceContext = resourceContext;
        this.interpreter = interpreter;
        this.javaResourceAdapter = javaResourceAdapter;
    }

    public void open() {
        interpreter.exec(PYTHON_IMPORTS);
        pythonResourceContext = (PyObject) interpreter.invoke(GET_RESOURCE_CONTEXT, this);
    }

    public Object getResource(String resourceName, String resourceType) {
        Resource resource;
        try {
            resource =
                    this.resourceContext.getResource(
                            resourceName, ResourceType.fromValue(resourceType));
        } catch (RuntimeException e) {
            throw e;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        if (resource instanceof PythonResourceWrapper) {
            PythonResourceWrapper pythonResource = (PythonResourceWrapper) resource;
            return pythonResource.getPythonResource();
        }
        if (resource instanceof Tool) {
            return convertToPythonTool((Tool) resource);
        }
        if (resource instanceof Prompt) {
            return convertToPythonPrompt((Prompt) resource);
        }
        return toPythonResource(resourceType, resource);
    }

    private Object toPythonResource(String resourceType, Resource resource) {
        Map<String, Object> kwargs = new HashMap<>();
        kwargs.put(JAVA_RESOURCE, resource);
        kwargs.put(JAVA_RESOURCE_ADAPTER, javaResourceAdapter);
        kwargs.put(RESOURCE_CONTEXT_KEY, pythonResourceContext);
        return interpreter.invoke(FROM_JAVA_RESOURCE, resourceType, kwargs);
    }

    @Override
    public PyObject initPythonResource(String module, String clazz, Map<String, Object> kwargs) {
        kwargs.put(RESOURCE_CONTEXT_KEY, pythonResourceContext);
        return (PyObject) interpreter.invoke(CREATE_RESOURCE, module, clazz, kwargs);
    }

    @Override
    public Object toPythonChatMessage(ChatMessage message) {
        return interpreter.invoke(FROM_JAVA_CHAT_MESSAGE, message);
    }

    @Override
    public ChatMessage fromPythonChatMessage(Object pythonChatMessage) {
        // TODO: Update this method after the pemja findClass method is fixed.
        ChatMessage chatMessage = new ChatMessage();

        String roleValue =
                (String) interpreter.invoke(TO_JAVA_CHAT_MESSAGE, pythonChatMessage, chatMessage);
        chatMessage.setRole(MessageRole.fromValue(roleValue));
        return chatMessage;
    }

    @Override
    public Object toPythonDocuments(List<Document> documents) {
        List<Object> pythonDocuments = new ArrayList<>();
        for (Document document : documents) {
            pythonDocuments.add(interpreter.invoke(FROM_JAVA_DOCUMENT, document));
        }
        return pythonDocuments;
    }

    @Override
    public List<Document> fromPythonDocuments(List<PyObject> pythonDocuments) {
        List<Document> documents = new ArrayList<>();
        for (PyObject pythonDocument : pythonDocuments) {
            Document document =
                    new Document(
                            pythonDocument.getAttr("content").toString(),
                            (Map<String, Object>) pythonDocument.getAttr("metadata", Map.class),
                            pythonDocument.getAttr("id").toString());
            documents.add(document);
        }
        return documents;
    }

    @Override
    public Object toPythonVectorStoreQuery(VectorStoreQuery query) {
        return interpreter.invoke(FROM_JAVA_VECTOR_STORE_QUERY, query);
    }

    @Override
    public VectorStoreQueryResult fromPythonVectorStoreQueryResult(
            PyObject pythonVectorStoreQueryResult) {
        List<PyObject> pythonDocuments =
                (List<PyObject>) pythonVectorStoreQueryResult.getAttr("documents", List.class);
        return new VectorStoreQueryResult(fromPythonDocuments(pythonDocuments));
    }

    @Override
    public Object convertToPythonTool(Tool tool) {
        return interpreter.invoke(FROM_JAVA_TOOL, tool);
    }

    private Object convertToPythonPrompt(Prompt prompt) {
        return interpreter.invoke(FROM_JAVA_PROMPT, prompt);
    }

    @Override
    public Object callMethod(Object obj, String methodName, Map<String, Object> kwargs) {
        return interpreter.invoke(CALL_METHOD, obj, methodName, kwargs);
    }

    @Override
    public Object invoke(String name, Object... args) {
        return interpreter.invoke(name, args);
    }

    @Override
    public Map<String, String> getPythonToolMetadata(String module, String qualName) {
        @SuppressWarnings("unchecked")
        Map<String, String> result =
                (Map<String, String>)
                        interpreter.invoke(GET_PYTHON_TOOL_METADATA, module, qualName);
        if (result == null) {
            throw new IllegalStateException(
                    "Python get_python_tool_metadata returned null for " + module + ":" + qualName);
        }
        return result;
    }

    @Override
    public Object invokePythonTool(String module, String qualName, Map<String, Object> kwargs) {
        return interpreter.invoke(INVOKE_PYTHON_TOOL, module, qualName, kwargs);
    }
}
