/*  
 *  Imixs-Workflow 
 *  
 *  Copyright (C) 2001-2020 Imixs Software Solutions GmbH,  
 *  http://www.imixs.com
 *  
 *  This program is free software; you can redistribute it and/or 
 *  modify it under the terms of the GNU General Public License 
 *  as published by the Free Software Foundation; either version 2 
 *  of the License, or (at your option) any later version.
 *  
 *  This program is distributed in the hope that it will be useful, 
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of 
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 
 *  General Public License for more details.
 *  
 *  You can receive a copy of the GNU General Public
 *  License at http://www.gnu.org/licenses/gpl.html
 *  
 *  Project: 
 *      https://www.imixs.org
 *      https://github.com/imixs/imixs-workflow
 *  
 *  Contributors:  
 *      Imixs Software Solutions GmbH - Project Management
 *      Ralph Soika - Software Developer
 */

package ai.dynamixs.rag;

import java.util.List;
import java.util.Optional;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;

import org.eclipse.microprofile.config.inject.ConfigProperty;
import org.imixs.ai.workflow.OpenAIAPIService;
import org.imixs.workflow.ItemCollection;
import org.imixs.workflow.SignalAdapter;
import org.imixs.workflow.engine.WorkflowService;
import org.imixs.workflow.exceptions.AdapterException;
import org.imixs.workflow.exceptions.PluginException;

import ai.dynamixs.rag.cluster.ClusterException;
import ai.dynamixs.rag.cluster.ClusterService;
import ai.dynamixs.rag.cluster.RetrievalResult;
import jakarta.inject.Inject;

/**
 * The RAGAdapter is used for Retrieval-Augmented Generation workflows.
 * <p>
 * The adapter automatically index and embed text data.
 * <p>
 * The imixs-ai prompt definition defines the API endpoint and processing events
 * to adapt the indexing and the prompt template.
 * <p>
 * The Adapter can be configured by a event entity.
 * <p>
 * INDEX configuration - to add embeddings:
 * 
 * <pre>
 * {@code
        <imixs-ai name="INDEX">
            <endpoint><propertyvalue>llm.service.endpoint</propertyvalue></endpoint>
            <debug>true</debug>
        </imixs-ai>
 * }
 * </pre>
 * 
 * The Endpoint defines the Rest API endpoint of the llama-cpp http server or
 * any compatible OpenAI / Open API rest service endpoint.
 * <p>
 * Example Prompt Template:
 * 
 * 
 * <pre>
 * {@code
    <?xml version="1.0" encoding="UTF-8"?>
    <IndexDefinition>
    <index_options>{"n_predict": 4096, "temperature": 0}</index_options>
    <![CDATA[<text><itemvalue>$workflowgroup</itemvalue>: <itemvalue>$workflowstatus</itemvalue>

    # Summary
    <itemvalue>$workflowsummary</itemvalue>

    <itemvalue>offer.summary</itemvalue>

    ]]></text>
    </IndexDefinition> 
 * }
 * </pre>
 * 
 * @author Ralph Soika
 * @version 1.0
 *
 */
public class RAGAdapter implements SignalAdapter {

    public static final String ML_ENTITY = "entity";
    public static final String API_ERROR = "API_ERROR";

    public static final String RAG_INDEX = "INDEX";
    public static final String RAG_DELETE = "DELETE";
    public static final String RAG_RETRIEVAL = "RETRIEVAL";

    private static Logger logger = Logger.getLogger(RAGAdapter.class.getName());

    @Inject
    @ConfigProperty(name = OpenAIAPIService.ENV_LLM_SERVICE_ENDPOINT)
    Optional<String> serviceEndpoint;

    @Inject
    private WorkflowService workflowService;

    @Inject
    private OpenAIAPIService llmService;

    @Inject
    private ClusterService clusterService;

    /**
     * Default Constructor
     */
    public RAGAdapter() {
        super();
    }

    /**
     * CDI Constructor to inject WorkflowService
     * 
     * @param workflowService
     */
    @Inject
    public RAGAdapter(WorkflowService workflowService) {
        super();
        this.workflowService = workflowService;
    }

    public void setWorkflowService(WorkflowService workflowService) {
        this.workflowService = workflowService;
    }

    /**
     * This method parses the LLM Event definitions.
     * 
     * For each PROMPT the method posts a context data (e.g text from an attachment)
     * to the Imixs-AI Analyse service endpoint
     * 
     * @throws PluginException
     */
    public ItemCollection execute(ItemCollection workitem, ItemCollection event)
            throws AdapterException, PluginException {

        logger.finest("...running api adapter...");

        // read optional configuration form the model or imixs.properties....
        List<ItemCollection> llmIndexDefinitions = workflowService.evalWorkflowResultXML(event, "imixs-ai",
                RAG_INDEX, workitem, false);
        List<ItemCollection> llmRetrievalDefinitions = workflowService.evalWorkflowResultXML(event, "imixs-ai",
                RAG_RETRIEVAL, workitem, false);

        List<ItemCollection> llmDeleteDefinitions = workflowService.evalWorkflowResultXML(event, "imixs-ai",
                RAG_DELETE, workitem, false);

        // verify if we have an DELETE configuration
        if (llmDeleteDefinitions != null && llmDeleteDefinitions.size() > 0) {
            try {
                clusterService.removeAllEmbeddings(workitem.getUniqueID());
            } catch (ClusterException e) {
                throw new PluginException(RAGAdapter.class.getSimpleName(), API_ERROR, e.getMessage(), e);
            }
        }

        // verify if we have an INDEX configuration
        if (llmIndexDefinitions != null && llmIndexDefinitions.size() > 0) {
            createIndex(llmIndexDefinitions, workitem, event);
        }

        // verify if we have an RETRIEVAL configuration
        if (llmRetrievalDefinitions != null && llmRetrievalDefinitions.size() > 0) {
            createRetrieval(llmRetrievalDefinitions, workitem, event);
        }

        return workitem;
    }

    /**
     * Creates a RAG Index for a Worktiem based on IndexDefinitions
     * 
     * @param llmIndexDefinitions
     * @param workitem
     * @param event
     * @throws PluginException
     * @throws AdapterException
     */
    private void createIndex(List<ItemCollection> llmIndexDefinitions, ItemCollection workitem, ItemCollection event)
            throws PluginException, AdapterException {
        long processingTime = System.currentTimeMillis();
        String llmAPIEndpoint = null;
        boolean llmAPIDebug = false;
        try {
            for (ItemCollection indexDefinition : llmIndexDefinitions) {
                llmAPIEndpoint = parseLLMEndpointByBPMN(indexDefinition);
                // do we have a valid endpoint?
                if (llmAPIEndpoint == null || llmAPIEndpoint.isEmpty()) {
                    throw new PluginException(RAGAdapter.class.getSimpleName(), API_ERROR,
                            "imixs-ai configuration error: no llm service endpoint defined!");
                }
                if ("true".equalsIgnoreCase(indexDefinition.getItemValueString("debug"))) {
                    llmAPIDebug = true;
                }

                logger.info("├── post RAG index request: " + llmAPIEndpoint);
                String promptTemplate = llmService.loadPromptTemplate(event);
                logger.info("│   ├── PromptTemplate: ");
                logger.info(promptTemplate);
                String llmPrompt = llmService.buildPrompt(promptTemplate, workitem);
                if (llmPrompt.isBlank()) {
                    throw new PluginException(
                            RAGAdapter.class.getSimpleName(), API_ERROR,
                            "Unable to parse prompt for RAG indexing");
                }

                if (llmAPIDebug) {
                    logger.info("│   ├── Total Prompt Length = " + llmPrompt.length());
                    logger.info("│   ├── Prompt: ");
                    logger.info(llmPrompt);
                }

                // remove old embeddings
                clusterService.removeAllEmbeddings(workitem.getUniqueID());
                // chunk text....
                List<String> chunk_list = RAGUtil.chunkMarkupDocument(llmPrompt);
                for (String chunk : chunk_list) {
                    if (llmAPIDebug) {
                        logger.info("│   ├── 🔸 chunk: ");
                        logger.info(chunk);
                    }
                    List<Float> indexResult = llmService.postEmbedding(chunk, llmAPIEndpoint, llmAPIDebug);
                    // write to cassandra
                    clusterService.insertVector(workitem.getUniqueID(),
                            workitem.getModelVersion(),
                            workitem.getTaskID(),
                            chunk,
                            indexResult);
                    if (llmAPIDebug) {
                        logger.info("│   ├── ⇨ " + indexResult.size() + " floats stored in RAG db.");
                    }
                }

                if (llmAPIDebug) {
                    logger.info(
                            "├── ✅ Total processing time: " + (System.currentTimeMillis() - processingTime) + "ms");
                }
            }
        } catch (PluginException e) {
            logger.severe("├── ⚠️ Failed to post embeddings: " + e.getMessage());
            throw e;
        } catch (ClusterException e) {
            throw new PluginException(RAGAdapter.class.getSimpleName(), API_ERROR, e.getMessage(), e);
        }
    }

    /**
     * Creates a RAG retrieval for a Worktiem based on RetrievalDefinitions
     * 
     * @param llmRetrievalfDefinitions
     * @param workitem
     * @param event
     * @throws PluginException
     * @throws AdapterException
     */
    private void createRetrieval(List<ItemCollection> llmRetrievalfDefinitions, ItemCollection workitem,
            ItemCollection event) throws PluginException, AdapterException {
        long processingTime = System.currentTimeMillis();
        String llmAPIEndpoint = null;
        boolean llmAPIDebug = false;
        try {
            for (ItemCollection indexDefinition : llmRetrievalfDefinitions) {
                llmAPIEndpoint = parseLLMEndpointByBPMN(indexDefinition);
                // do we have a valid endpoint?
                if (llmAPIEndpoint == null || llmAPIEndpoint.isEmpty()) {
                    throw new PluginException(RAGAdapter.class.getSimpleName(), API_ERROR,
                            "imixs-ai configuration error: no llm service endpoint defined!");
                }
                if ("true".equalsIgnoreCase(indexDefinition.getItemValueString("debug"))) {
                    llmAPIDebug = true;
                }

                logger.info("├── post RAG Retrievalf request: " + llmAPIEndpoint);
                String promptTemplate = llmService.loadPromptTemplate(event);
                logger.info("│   ├── PromptTemplate: ");
                logger.info(promptTemplate);
                String llmPrompt = llmService.buildPrompt(promptTemplate, workitem);
                if (llmPrompt.isBlank()) {
                    throw new PluginException(
                            RAGAdapter.class.getSimpleName(), API_ERROR,
                            "Unable to parse prompt for RAG indexing");
                }
                if (llmAPIDebug) {
                    logger.info("│   ├── Total Prompt Length = " + llmPrompt.length());
                    logger.info("│   ├── Prompt: ");
                    logger.info(llmPrompt);
                }

                // retrieve prompt....
                List<Float> embeddings = llmService.postEmbedding(llmPrompt, llmAPIEndpoint, llmAPIDebug);
                if (llmAPIDebug) {
                    logger.info("├── ⇨ " + embeddings.size() + " floats stored in RAG db.");
                }
                // search cassandra
                List<RetrievalResult> retrievalResultList = clusterService.searchEmbeddings(embeddings, 5);
                List<String> listOfIds = retrievalResultList.stream()
                        .map(RetrievalResult::getUniqueId)
                        .collect(Collectors.toList());
                workitem.setItemValue("product.ref", listOfIds);
                if (llmAPIDebug) {
                    logger.info(
                            "├── found " + retrievalResultList.size() + " matches");
                    logger.info(
                            "├── ✅ Total processing time: " + (System.currentTimeMillis() - processingTime) + "ms");
                }
            }
        } catch (PluginException e) {
            logger.severe("├── ⚠️ Failed to post embeddings: " + e.getMessage());
            throw e;
        } catch (ClusterException e) {
            throw new PluginException(RAGAdapter.class.getSimpleName(), API_ERROR, e.getMessage(), e);
        }
    }

    /**
     * This helper method parses the ml api endpoint either provided by a model
     * definition or a imixs.property or an environment variable.
     * <p>
     * If not api endpoint is defined by the model the adapter uses the default api
     * endpoint.
     * 
     * @param llmPrompt
     * @return
     * @throws PluginException
     */
    private String parseLLMEndpointByBPMN(ItemCollection llmPrompt) throws PluginException {
        boolean debug = logger.isLoggable(Level.FINE);
        String llmAPIEndpoint = null;

        // Test if the model provides a API Endpoint.
        llmAPIEndpoint = null;
        if (llmPrompt != null) {
            llmAPIEndpoint = llmPrompt.getItemValueString("endpoint");
        }

        // switch to default api endpoint?
        if (llmAPIEndpoint == null || llmAPIEndpoint.isEmpty()) {
            // set defautl api endpoint if defined
            if (serviceEndpoint.isPresent() && !serviceEndpoint.get().isEmpty()) {
                llmAPIEndpoint = serviceEndpoint.get();
            }
        }
        if (debug) {
            logger.info("......llm api endpoint " + llmAPIEndpoint);
        }

        // adapt text...
        llmAPIEndpoint = workflowService.adaptText(llmAPIEndpoint, null);
        if (!llmAPIEndpoint.endsWith("/")) {
            llmAPIEndpoint = llmAPIEndpoint + "/";
        }
        return llmAPIEndpoint;
    }

}