/*
 *  Copyright 2025 Anyware Services
 *
 *  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 org.ametys.plugins.ai.provider.impl;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.List;

import org.apache.avalon.framework.activity.Disposable;
import org.apache.avalon.framework.activity.Initializable;
import org.apache.avalon.framework.component.Component;
import org.apache.avalon.framework.configuration.Configurable;
import org.apache.avalon.framework.configuration.Configuration;
import org.apache.avalon.framework.configuration.ConfigurationException;
import org.apache.avalon.framework.service.ServiceException;
import org.apache.avalon.framework.service.ServiceManager;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.hc.client5.http.classic.methods.HttpGet;
import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;

import org.ametys.core.util.HttpUtils;
import org.ametys.plugins.ai.AIHelper;
import org.ametys.plugins.ai.provider.AIProvider;
import org.ametys.runtime.config.Config;
import org.ametys.runtime.i18n.I18nizableText;
import org.ametys.runtime.plugin.component.AbstractLogEnabled;
import org.ametys.runtime.plugin.component.DeferredServiceable;
import org.ametys.runtime.plugin.component.PluginAware;
import org.ametys.runtime.util.AmetysHomeHelper;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.exception.AuthenticationException;
import dev.langchain4j.exception.RateLimitException;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.output.Response;

/**
 * This is an abstract class for AI providers.
 */
public abstract class AbstractAIProvider extends AbstractLogEnabled implements AIProvider, Configurable, Initializable, PluginAware, DeferredServiceable, Component, Disposable
{
    /** The configuration parameter for the API key  */
    public static final String CONFIG_API_KEY = "ai.apikey";
    
    /** The configuration parameter for URL */
    public static final String CONFIG_SERVER_URL = "ai.serverURL";
    
    /** The configuration parameter for text model */
    public static final String CONFIG_TEXT_MODEL = "ai.model.text";
    
    /** The configuration parameter for image model */
    public static final String CONFIG_IMAGE_MODEL = "ai.model.image";
    
    /** The configuration parameter for image generation timeout */
    public static final String CONFIG_IMAGE_TIMEOUT = "ai.model.image.timeout";

    /** The configuration parameter for embedding model */
    public static final String CONFIG_EMBEDDING_MODEL = "ai.model.embedding";
    
    /** The provider id */
    protected String _id;
    /** The provider label */
    protected I18nizableText _label;
    /** The plugin name */
    protected String _pluginName;
    
    /** The provider models existing for text generation */
    protected List<String> _textModelNames;
    /** The provider models existing for image generation */
    protected List<String> _imageModelNames;
    /** The provider models existing for embedding */
    protected List<EmbeddingModelInfo> _embeddingModelInfos;
    
    /** Current text model */
    protected ChatModel _textModel;
    /** The context window for text generation */
    protected int _textModelContextWindow;
    /** Current image model */
    protected ImageModel _imageModel;
    /** The context window for image generation */
    protected int _imageModelContextWindow;
    /** Current embedding model */
    protected EmbeddingModel _embeddingModel;
    
    /** The AI helper */
    protected AIHelper _aiHelper;
    
    /** The HTTP client used to fetch images */
    protected CloseableHttpClient _httpClient;

    /**
     * Info for an embedding model
     * @param name The name of the model
     * @param vectorDimension The vector dimension
     */
    public record EmbeddingModelInfo(String name, int vectorDimension) { }
    

    public void setPluginInfo(String pluginName, String featureName, String id)
    {
        _pluginName = pluginName;
        _id = id;
    }

    public void configure(Configuration configuration) throws ConfigurationException
    {
        _label = I18nizableText.parseI18nizableText(configuration.getChild("label"), "plugin." + _pluginName, "");
        
        Configuration modelConfiguration = configuration.getChild("model");
        
        // Text
        Configuration textConfiguration = modelConfiguration.getChild("text");
        
        _textModelNames = Arrays.stream(textConfiguration.getChildren("name"))
            .map(c -> c.getValue(null))
            .filter(StringUtils::isNotBlank)
            .toList();
        
        int contextWindow = textConfiguration.getChild("contextWindow").getValueAsInteger(-1);
        if (contextWindow == -1)
        {
            throw new ConfigurationException("Missing <contextWindow> for model <text>", modelConfiguration);
        }
        _textModelContextWindow = contextWindow;

        // Image
        Configuration imageConfiguration = modelConfiguration.getChild("image", false);
        if (imageConfiguration != null)
        {
            _imageModelNames = Arrays.stream(imageConfiguration.getChildren("name"))
                    .map(c -> c.getValue(null))
                    .filter(StringUtils::isNotBlank)
                    .toList();
                
            contextWindow = imageConfiguration.getChild("contextWindow").getValueAsInteger(-1);
            if (contextWindow == -1)
            {
                throw new ConfigurationException("Missing <contextWindow> for model <image>", modelConfiguration);
            }
            
            _imageModelContextWindow = contextWindow;
        }
        else
        {
            _imageModelNames = List.of();
        }


        // Embedding
        Configuration embeddingConfiguration = modelConfiguration.getChild("embedding", false);
        if (embeddingConfiguration != null)
        {
            _embeddingModelInfos = Arrays.stream(embeddingConfiguration.getChildren("name"))
                    .map(c -> new EmbeddingModelInfo(c.getValue(null), Integer.parseInt(c.getAttribute("vectorDimension", ""))))
                    .filter(v -> StringUtils.isNotBlank(v.name()))
                    .toList();
        }
        else
        {
            _embeddingModelInfos = List.of();
        }
    }
    
    private synchronized void _initModels()
    {
        if (_textModel == null)
        {
            _textModel = createChatModel(Config.getInstance().getValue(CONFIG_TEXT_MODEL));
            
            if (_imageModelContextWindow != -1)
            {
                String imageModelName = Config.getInstance().getValue(CONFIG_IMAGE_MODEL);
                if (StringUtils.isNotBlank(imageModelName))
                {
                    _imageModel = createImageModel(imageModelName);
                }
            }

            if (_embeddingModelInfos.size() > 0)
            {
                String embeddingModelName = Config.getInstance().getValue(CONFIG_EMBEDDING_MODEL);
                if (StringUtils.isNotBlank(embeddingModelName))
                {
                    _embeddingModel = createEmbeddingModel(embeddingModelName);
                }
            }
        }
    }
    
    /**
     * Creates the text model implementation
     * @param modelName The model name
     * @return The chat model
     */
    protected abstract ChatModel createChatModel(String modelName);
    
    /**
     * Creates the image model implementation
     * @param modelName The model name
     * @return The image model
     */
    protected ImageModel createImageModel(String modelName)
    {
        throw new UnsupportedOperationException(this.getClass().getName() + " does not support image generation");
    }
    
    /**
     * Creates the embedding model implementation
     * @param modelName The model name
     * @return The embedding model
     */
    protected abstract EmbeddingModel createEmbeddingModel(String modelName);

    public void initialize() throws Exception
    {
        _httpClient = HttpUtils.createHttpClient(0, 2, false);
    }
    
    public void dispose()
    {
        try
        {
            _httpClient.close();
        }
        catch (IOException e)
        {
            throw new RuntimeException(e);
        }
    }
    
    public void deferredService(ServiceManager manager) throws ServiceException
    {
        _aiHelper = (AIHelper) manager.lookup(AIHelper.ROLE);
    }
    
    public String getId()
    {
        return _id;
    }

    public I18nizableText getLabel()
    {
        return _label;
    }
    
    /**
     * Get the list of known text models for this provider
     * @return The list
     */
    public List<String> getKnownTextModels()
    {
        return _textModelNames;
    }

    /**
     * Get the list of known image models for this provider
     * @return The non null list
     */
    public List<String> getKnownImageModels()
    {
        return _imageModelNames;
    }
    
    /**
     * Get the list of known embedding models for this provider
     * @return the non null list
     */
    public List<EmbeddingModelInfo> getKnownEmbeddingModels()
    {
        return _embeddingModelInfos;
    }
    
    /**
     * Creates a token count estimator
     * @return The non null token count estimator
     */
    protected TokenCountEstimator _getTokenCountEstimator()
    {
        return new DefaultTokenCountEstimator();
    }

    public String textToSummary(String text, int maxLength) throws Exception
    {
        _initModels();
        
        TokenCountEstimator estimator = _getTokenCountEstimator();
        
        // Get the token number of the prompt without the input text, to not exceed context window.
        int tokenForPrompt = estimator.estimateTokenCountInText(_aiHelper.getTextSummaryPrompt(maxLength));
        int contextWindow = _textModelContextWindow;
        int tokenMax = contextWindow - tokenForPrompt;
        
        List<String> summaries = new ArrayList<>();
        DocumentSplitter splitter = DocumentSplitters.recursive(tokenMax , 5, estimator);
        List<TextSegment> segments = splitter.split(Document.from(text));
        for (TextSegment segment : segments)
        {
            boolean success = false;
            int retries = 0;
            while (!success && retries < 3)
            {
                try
                {
                    ChatResponse response = _textModel.chat(new SystemMessage(_aiHelper.getTextSummaryPrompt(maxLength)), new UserMessage(segment.text()));
                    summaries.add(response.aiMessage().text());
                    success = true;
                }
                catch (RateLimitException e)
                {
                    retries++;
                    Thread.sleep(40000); // Waiting a little to low the token per minute
                }
            }
            
            if (!success)
            {
                getLogger().warn("An error occurred during summary generation of the segment. So we ignore it.");
            }
        }
        
        // Only if one summary exists, return the summary
        if (summaries.size() == 1)
        {
            return summaries.get(0);
        }
        // Else, return a final summary of all summaries
        String summary = String.join(" ", summaries);
        String prompt = _aiHelper.getTextSummaryPrompt(maxLength);
        
        ChatResponse response = _textModel.chat(SystemMessage.from(prompt), UserMessage.from(summary));
        return response.aiMessage().text();
    }
    
    public boolean isImageGenerationSupported()
    {
        _initModels();
        
        return _imageModel != null;
    }
    
    public Path textToImage(String text) throws Exception
    {
        _initModels();
        
        TokenCountEstimator estimator = new DefaultTokenCountEstimator();
        
        // Get the token number of the prompt without the input text, to not exceed context window.
        int tokenForPrompt = estimator.estimateTokenCountInText(_aiHelper.getImageGenerationPrompt());
        int contextWindow = _imageModelContextWindow;
        int tokenMax = contextWindow - tokenForPrompt;
        
        DocumentSplitter splitter = DocumentSplitters.recursive(tokenMax, 0, estimator);
        List<TextSegment> segments = splitter.split(Document.from(text));

        if (segments.isEmpty())
        {
            throw new IllegalArgumentException("An error occurred generating the image. Can't build a segment for the given text " + text);
        }
        String res = segments.get(0).text();

        String prompt = _aiHelper.getImageGenerationPrompt();
        
        Response<Image> response = _imageModel.generate(prompt + " Here is the user text: " + res);
        Image image = response.content();
        
        if (image.base64Data() != null)
        {
            String tmpFolderName = Long.toString(Math.round(Math.random() * 1000000.0));
            String fileName = "image.png";
            
            Path targetFile = AmetysHomeHelper.getAmetysHomeTmp().toPath().resolve(tmpFolderName).resolve(fileName);
            Files.createDirectories(targetFile.getParent());

            byte[] imageBytes = Base64.getDecoder().decode(image.base64Data());
            Files.write(targetFile, imageBytes);
            return targetFile;
        }
        else
        {
            return _download(image.url());
        }
    }
    
    /**
     * Helper to download the given url to a temporary file
     * @param uri The uri to download
     * @return The created file
     * @throws IOException If an error occurred
     */
    protected Path _download(URI uri) throws IOException
    {
        HttpGet httpGet = new HttpGet(uri);
        
        String tmpFolderName = Long.toString(Math.round(Math.random() * 1000000.0));
        String fileName = StringUtils.substringAfterLast(uri.getPath(), "/");
        
        Path targetFile = AmetysHomeHelper.getAmetysHomeTmp().toPath().resolve(tmpFolderName).resolve(fileName);
        Files.createDirectories(targetFile.getParent());
        
        return _httpClient.execute(httpGet, httpResponse ->
        {
            if (httpResponse.getCode() != 200)
            {
                throw new IOException("Error " + httpResponse.getCode() + " while accessing to the URI " + uri);
            }
            
            try (InputStream is = httpResponse.getEntity().getContent();
                    OutputStream os = Files.newOutputStream(targetFile))
            {
                IOUtils.copy(is, os);
            }
            
            return targetFile;
        });
    }

    
    /**
     * Get the API key from the configuration
     * @return the API key
     */
    protected String getAPIKey()
    {
        return Config.getInstance().getValue(CONFIG_API_KEY);
    }

    /**
     * Get the server URL from the configuration
     * @return the server URL
     */
    protected String getServerURL()
    {
        return Config.getInstance().getValue(CONFIG_SERVER_URL);
    }
    
    public boolean isEmbeddingSupported()
    {
        _initModels();
        
        return _embeddingModel != null;
    }
    
    public String askChatbot(ChatMemory memory)
    {
        _initModels();
        
        ChatRequest chatRequest = ChatRequest.builder()
                .messages(memory.messages())
                .build();
        
        ChatResponse chatResponse = _textModel.chat(chatRequest);
        return chatResponse.aiMessage().text();
    }
    
    public List<Float> computeEmbedding(String text)
    {
        try
        {
            _initModels();
            
            Embedding embedding = _embeddingModel.embed(text).content();
            List<Float> embeddingList = embedding.vectorAsList();
            
            if (isEmbeddingNormalizationRequired())
            {
                return normalize(embeddingList);
            }
            else
            {
                return embeddingList;
            }
        }
        catch (AuthenticationException e)
        {
            getLogger().warn("Authentication error while computing embedding. Please check the API key.", e);
            return null;
        }
    }
    
    /**
     * Is normalization required for the embeddings of this provider
     * @return true if normalization is required
     */
    protected boolean isEmbeddingNormalizationRequired()
    {
        return false;
    }

    /**
     * Necessary for open models who often are not normalized.
     * @param vector the vector to normalize
     * @return the normalized vector
     */
    public List<Float> normalize (List<Float> vector)
    {
        double sumSquares = 0.0;
        for (Float v : vector)
        {
            sumSquares += v * v;
        }
        double norm = Math.sqrt(sumSquares);
        if (norm == 0)
        {
            return vector;
        }
        List<Float> normalized = new ArrayList<>(vector.size());
        for (Float v : vector)
        {
            normalized.add((float) (v / norm));
        }
        return normalized;
    }
    
    public int getEmbeddingDimension()
    {
        String modelName = Config.getInstance().getValue(CONFIG_EMBEDDING_MODEL);
        if (StringUtils.isNotBlank(modelName))
        {
            for (EmbeddingModelInfo info : _embeddingModelInfos)
            {
                if (info.name().equals(modelName))
                {
                    return info.vectorDimension();
                }
            }
        }
        return -1;
    }
    
    public String getEmbeddingKey()
    {
        if (isEmbeddingSupported())
        {
            return this.getId() + "_" + Config.getInstance().getValue(CONFIG_EMBEDDING_MODEL) + ":"  + getEmbeddingDimension();
        }
        return null;
    }
}