/**
 * Copyright (c) 2025 SoftCorporation LLC. All rights reserved.
 *
 * The Software License, Version 1.0
 *
 * SoftCorporation LLC. grants you ("Licensee") a non-exclusive, royalty free,
 * license to use, modify and redistribute this software in source and binary
 * code form, provided that the following conditions are met:
 *
 * 1. The end-user documentation included with the redistribution,
 *    if any, must include the following acknowledgment:
 *       "This product includes software developed by the
 *        SoftCorporation LLC. (http://www.softcorporation.com)."
 *    Alternately, this acknowledgment may appear in the software itself,
 *    if and wherever such third-party acknowledgments normally appear.
 *
 * 2. This software and SoftCorporation name must not be used to
 *    promote products derived from the software without prior
 *    written permission. For written permission, please contact
 *    info@softcorporation.com.
 *
 * This software is provided "AS IS," without a warranty of any kind.
 * ALL EXPRESS OR IMPLIED CONDITIONS, REPRESENTATIONS AND WARRANTIES,
 * INCLUDING ANY IMPLIED WARRANTY OF MERCHANTABILITY, FITNESS FOR A
 * PARTICULAR PURPOSE OR NON-INFRINGEMENT, ARE HEREBY EXCLUDED.
 * IN NO EVENT SHALL THE SOFTCORPORATION BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
 * USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION).
 *
 */
package com.softcorporation.sandi.base.service;

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

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import com.softcorporation.sandi.base.model.Client;
import com.softcorporation.sandi.base.model.TextChunk;
import com.softcorporation.sandi.base.util.Constants;
import com.softcorporation.sandi.base.util.Utils;
import com.softcorporation.util.VecUtil;
import com.softcorporation.sandi.base.util.EmbeddingException;

@Service
public class TextEmbeddingService {

	private static final Logger LOGGER = LoggerFactory.getLogger(TextEmbeddingService.class);

	/**
	 * Max array of texts (depends on NVIDIA card memory size)
	 */
    @Value("${sandi.embedding.process.max:50}")
    private int embeddingProcessMax;
	
	/**
	 * 1500 max chunk length for EN 512 tokens
	 */
    @Value("${sandi.embedding.chunk.length.max:1500}")
    private int embeddingChunkLengthMax;
	
	/**
	 * overlap chunk length
	 */
    @Value("${sandi.embedding.chunk.length.overlap:200}")
    private int embeddingChunkLengthOverlap;
	
	/**
	 * min chunk length
	 */
    @Value("${sandi.embedding.chunk.length.min:100}")
    private int embeddingChunkLengthMin;
	
	/**
	 * Embedding chunk weight dump
	 */
    @Value("${sandi.embedding.chunk.weight.dump:0.3}")
    private double embeddingChunkWeightDump;
    
	/**
	 * Embedding chunk weight high
	 */
    @Value("${sandi.embedding.chunk.weight.high:0.2}")
    private double embeddingChunkWeightHigh;
    
	/**
	 * Embedding chunk weight low
	 */
    @Value("${sandi.embedding.chunk.weight.low:0.1}")
    private double embeddingChunkWeightLow;
    
    @Autowired
    private LanguageModelService languageModelService;
    
    /**
	 * Embedding byte scale
	 */
    @Value("${sandi.embedding.byte.scale:500}")
    private int embByteScale;
    
	/**
 	 * Get chunk embeddings
 	 * 
	 * @param chunksHigh
	 * @param chunksLow
	 * @param chunksContent
	 * @return map
	 * @throws IndexerException
	 */
	public List<Map<String, Object>> getDocumentEmbeddings(final Client client, final List<TextChunk> chunksHigh,
			final List<TextChunk> chunksLow, final List<TextChunk> chunksContent) throws EmbeddingException {
	    List<Map<String, Object>> embeddings = new ArrayList<>();

	    List<TextChunk> chunkList = new ArrayList<>();
		try {
		    for (TextChunk textChunk : chunksHigh) {
		    	chunkList.add(textChunk);
			    if (chunkList.size() >= embeddingProcessMax) {
			    	addEmbedding(client, chunkList, embeddings);
			    	chunkList = new ArrayList<>();
			    }
		    }
		    for (TextChunk textChunk : chunksLow) {
		    	chunkList.add(textChunk);
			    if (chunkList.size() >= embeddingProcessMax) {
			    	addEmbedding(client, chunkList, embeddings);
			    	chunkList = new ArrayList<>();
			    }
		    }
		    for (TextChunk textChunk : chunksContent) {
		    	chunkList.add(textChunk);
			    if (chunkList.size() >= embeddingProcessMax) {
			    	addEmbedding(client, chunkList, embeddings);
			    	chunkList = new ArrayList<>();
			    }
		    }
		    if (chunkList.size() > 0) {
		    	addEmbedding(client, chunkList, embeddings);
		    }
		    
            adjustEmbeddings(embeddings);
            
		} catch (IOException e) {
            LOGGER.error("Error processing embedding: {}", e.getMessage());
            throw new EmbeddingException(e.getMessage());
		}
	    return embeddings;
	}

	/**
	 * Request texts embedding and add to list
	 * @throws IOException 
	 */
	private void addEmbedding(final Client client, final List<TextChunk> chunks, final List<Map<String, Object>> embeddings)
			throws IOException {
		List<String> texts = new ArrayList<>();
		for (TextChunk chunk : chunks) {
			texts.add(chunk.getText());
		}
		List<double[]> embList = languageModelService.getTextsEmbedding(client, texts);
		for (int i = 0; i < embList.size(); i++) {
			double[] emb = embList.get(i);
	        TextChunk chunk = chunks.get(i);
	        Map<String, Object> embDoc = new HashMap<>();
//	        embDoc.put(Constants.ID, id); // no need
	        embDoc.put(Constants.CHUNK_FIELD, chunk.getField());
	        embDoc.put(Constants.CHUNK_TYPE, chunk.getType());
	        embDoc.put(Constants.EMB, emb);
	        embDoc.put(Constants.CHUNK_BEGIN, chunk.getBegin());
	        embDoc.put(Constants.CHUNK_END, chunk.getEnd());
	        embDoc.put(Constants.CHUNK_TEXT, chunk.getText()); // added for slop
	        embeddings.add(embDoc);
		}
	}

	/**
	 * Break text in chunks depending on type
	 * 
	 * @param text String
	 * @return chunks
	 */
	public List<TextChunk> breakTextInChunks(String text, String type) {
		List<TextChunk> chunks = new ArrayList<>();
		String textT = text.trim();
		if (textT.length() > 0) {
			if (textT.length() < embeddingChunkLengthMax) {
				String textL = textT.toLowerCase();
				if (Constants.TYPE_HIGH.equals(type)) {
					textT = textT + "\n" + textL + "\n" + textT + " " + textL;
					textT = Utils.getSubText(textT, 0, embeddingChunkLengthMax);
					TextChunk chunk = new TextChunk(textT, 0, text.length());
					chunks.add(chunk);
				} else if (Constants.TYPE_LOW.equals(type)) {
					textT = textT + "\n" + textL;
					textT = Utils.getSubText(textT, 0, embeddingChunkLengthMax);
					TextChunk chunk = new TextChunk(textT, 0, text.length());
					chunks.add(chunk);
				} else if (Constants.TYPE_CONTENT.equals(type)) {
					textT = textT + "\n" + textL;
					textT = Utils.getSubText(textT, 0, embeddingChunkLengthMax);
					TextChunk chunk = new TextChunk(textT, 0, text.length());
					chunks.add(chunk);
				}
			} else {
			    chunks = breakTextInChunks(text);
			}
		}
		return chunks;
	}
	
	/**
	 * Break text in chunks
	 * 
	 * @param text String
	 * @return chunks
	 */
	public List<TextChunk> breakTextInChunks(String text) {
		List<TextChunk> chunks = new ArrayList<>();
		text = text.trim();
		if (text.length() == 0) {
			return chunks;
		}
		int chunkBegin = 0;
		int paraBegin = 0;
		int lineBegin = 0;
		int sentenceBegin = 0;
		int phraseBegin = 0;
		int wordBegin = 0;
		int chunkParaBegin = 0;
		int chunkLineBegin = 0;
		int chunkSentenceBegin = 0;
		int chunkPhraseBegin = 0;
		int chunkWordBegin = 0;
		boolean paraBreak = false;
		boolean lineBreak = false;
		boolean sentenceBreak = false;
		boolean phraseBreak = false;
		boolean spaceBreak = false;
		char prevChar = 0;
		StringBuilder sb = new StringBuilder();
		for (int i = 0; i < text.length(); i++) {
			char c = text.charAt(i);
			switch (c) {
			case ' ':
				if (!spaceBreak) {
					spaceBreak = true;
					if (!lineBreak) {
						sb.append(' ');
					}
				}
				if (prevChar == '.') {
					sentenceBreak = true;
					phraseBreak = true;
				} else if (prevChar == ',' || prevChar == ';') {
					phraseBreak = true;
				}
				break;
			case '\t':
				if (!spaceBreak) {
					spaceBreak = true;
					if (!lineBreak) {
						sb.append('\t');
					}
				}
				if (prevChar == '.') {
					sentenceBreak = true;
					phraseBreak = true;
				} else if (prevChar == ',' || prevChar == ';') {
					phraseBreak = true;
				}
				break;
			case '\n':
				if (!lineBreak) {
					lineBreak = true;
					sb.append('\n');
				} else {
					if (prevChar != '\r') { // windows
						paraBreak = true;
					}
				}
				if (prevChar == '.') {
					sentenceBreak = true;
					phraseBreak = true;
				} else if (prevChar == ',' || prevChar == ';') {
					phraseBreak = true;
				}
				break;
			case '\r':
				if (!lineBreak) {
					lineBreak = true;
					sb.append('\n');
				} else {
					paraBreak = true;
				}
				if (prevChar == '.') {
					sentenceBreak = true;
					phraseBreak = true;
				} else if (prevChar == ',' || prevChar == ';') {
					phraseBreak = true;
				}
				break;
			default:
				sb.append(c);
				int chunkLength = sb.length();
				int pos = chunkLength - 1;
				if (spaceBreak) {
					spaceBreak = false;
					wordBegin = i;
					chunkWordBegin = pos;
				}
				if (phraseBreak) {
					phraseBreak = false;
					phraseBegin = i;
					chunkPhraseBegin = pos;
					wordBegin = i;
					chunkWordBegin = pos;
				}
				if (sentenceBreak) {
					sentenceBreak = false;
					sentenceBegin = i;
					chunkSentenceBegin = pos;
					phraseBegin = i;
					chunkPhraseBegin = pos;
					wordBegin = i;
					chunkWordBegin = pos;
				}
				if (lineBreak) {
					lineBreak = false;
					lineBegin = i;
					chunkLineBegin = pos;
					sentenceBegin = i;
					chunkSentenceBegin = pos;
					phraseBegin = i;
					chunkPhraseBegin = pos;
					wordBegin = i;
					chunkWordBegin = pos;
				}
				if (paraBreak) {
					paraBreak = false;
					paraBegin = i;
					chunkParaBegin = pos;
					lineBegin = i;
					chunkLineBegin = pos;
					sentenceBegin = i;
					chunkSentenceBegin = pos;
					phraseBegin = i;
					chunkPhraseBegin = pos;
					wordBegin = i;
					chunkWordBegin = pos;
				}
				if (chunkLength >= embeddingChunkLengthMax) {
					int nextChunkBegin;
					String nextChunk;
					if (chunkLength - chunkParaBegin < embeddingChunkLengthOverlap) {
						nextChunk = sb.substring(chunkParaBegin);
						nextChunkBegin = paraBegin;
					} else if (chunkLength - chunkLineBegin < embeddingChunkLengthOverlap) {
						nextChunk = sb.substring(chunkLineBegin);
						nextChunkBegin = lineBegin;
					} else if (chunkLength - chunkSentenceBegin < embeddingChunkLengthOverlap) {
						nextChunk = sb.substring(chunkSentenceBegin);
						nextChunkBegin = sentenceBegin;
					} else if (chunkLength - chunkPhraseBegin < embeddingChunkLengthOverlap) {
						nextChunk = sb.substring(chunkPhraseBegin);
						nextChunkBegin = phraseBegin;
					} else if (chunkLength - chunkWordBegin < embeddingChunkLengthOverlap) {
						nextChunk = sb.substring(chunkWordBegin);
						nextChunkBegin = wordBegin;
					} else {
						nextChunk = "";
						nextChunkBegin = i;
					}
					
					// overlap
					int chunkCut, chunkEnd;
					if (chunkLength - chunkWordBegin < embeddingChunkLengthOverlap) {
						chunkEnd = wordBegin - 1;
						chunkCut = chunkWordBegin - 1;
					} else if (chunkLength - chunkPhraseBegin < embeddingChunkLengthOverlap) {
						chunkEnd = phraseBegin - 1;
						chunkCut = chunkPhraseBegin - 1;
					} else if (chunkLength - chunkSentenceBegin < embeddingChunkLengthOverlap) {
						chunkEnd = sentenceBegin - 1;
						chunkCut = chunkSentenceBegin - 1;
					} else if (chunkLength - chunkLineBegin < embeddingChunkLengthOverlap) {
						chunkEnd = lineBegin - 1;
						chunkCut = chunkLineBegin - 1;
					} else if (chunkLength - chunkParaBegin < embeddingChunkLengthOverlap) {
						chunkEnd = paraBegin - 1;
						chunkCut = chunkParaBegin - 1;
					} else {
						chunkEnd = i;
						chunkCut = sb.length();
					}
					
					String chunk = sb.subSequence(0, chunkCut).toString();
					TextChunk textChunk = new TextChunk(chunk, chunkBegin, chunkEnd);
					chunks.add(textChunk);

					sb.setLength(0);
					sb.append(nextChunk);

					chunkBegin = nextChunkBegin;
				}
			}
			prevChar = c;
		}
		if (sb.length() >= embeddingChunkLengthMin) {
			String chunk = sb.toString();
			TextChunk textChunk = new TextChunk(chunk, chunkBegin, text.length());
			chunks.add(textChunk);
		} else if (sb.length() > 0) {
			String chunk = Utils.getSubText(text, text.length() - embeddingChunkLengthMin, text.length());
			TextChunk textChunk = new TextChunk(chunk, text.length() - chunk.length(), text.length());
			chunks.add(textChunk);
		}
		if (chunks.size() == 0) {
			TextChunk textChunk = new TextChunk(text, 0, text.length());
			chunks.add(textChunk);
		}
		return chunks;
	}
	
    /**
     * Adjust embedding
     * 
     * @param embeddings
     * @throws IllegalArgumentException
     */
    private void adjustEmbeddings(final List<Map<String, Object>> embeddings)
    {
    	List<Map<String, Object>> embeddingsHigh = new ArrayList<>();
    	List<Map<String, Object>> embeddingsLow = new ArrayList<>();
    	List<Map<String, Object>> embeddingsContent = new ArrayList<>();
        for (Map<String, Object> embedding : embeddings)
        {
        	if (Constants.TYPE_HIGH.equals(embedding.get(Constants.CHUNK_TYPE))) {
        		embeddingsHigh.add(embedding);
        	} else if (Constants.TYPE_LOW.equals(embedding.get(Constants.CHUNK_TYPE))) {
        		embeddingsLow.add(embedding);
        	} else if (Constants.TYPE_CONTENT.equals(embedding.get(Constants.CHUNK_TYPE))) {
        		embeddingsContent.add(embedding);
        	}
        }

        double[] embHigh = null;
    	for (int i = embeddingsHigh.size() - 1; i >= 0; i--) {
        	Map<String, Object> embeddingHigh = embeddingsHigh.get(i);
        	@SuppressWarnings("unchecked")
			double[] emb = (double[]) embeddingHigh.get(Constants.EMB);
    		if (embHigh == null) {
            	embHigh = emb;
    		} else {
            	embHigh = VecUtil.sumArrays(emb, embHigh, embeddingChunkWeightDump);
    		}
    		// TODO review normalization: breaks sum, but keeps direction?
    		VecUtil.normalizeArray(embHigh);
    		List<Byte> embByteHigh = VecUtil.toByteList(embHigh, embByteScale);
    		embeddingHigh.put(Constants.EMB, embByteHigh);
    	}

    	double[] embLow = null;
    	for (int i = embeddingsLow.size() - 1; i >= 0; i--) {
        	Map<String, Object> embeddingLow = embeddingsLow.get(i);
        	@SuppressWarnings("unchecked")
        	double[] emb = (double[]) embeddingLow.get(Constants.EMB);
    		if (embLow == null) {
            	embLow = emb;
    		} else {
            	embLow = VecUtil.sumArrays(emb, embLow, embeddingChunkWeightDump);
    		}
    		if (embHigh != null) {
    			embLow = VecUtil.sumArrays(emb, embHigh, embeddingChunkWeightHigh);
        	}
    		VecUtil.normalizeArray(embLow);
    		List<Byte> embByteLow = VecUtil.toByteList(embLow, embByteScale);
    		embeddingLow.put(Constants.EMB, embByteLow);
    	}
        
    	for (int i = 0; i < embeddingsContent.size(); i++) {
        	Map<String, Object> embeddingContent = embeddingsContent.get(i);
        	@SuppressWarnings("unchecked")
        	double[] embContent = (double[]) embeddingContent.get(Constants.EMB);
        	VecUtil.normalizeArray(embContent);
        	if (embLow != null) {
    			embContent = VecUtil.sumArrays(embContent, embLow, embeddingChunkWeightLow);
        	}
    		if (embHigh != null) {
    			embContent = VecUtil.sumArrays(embContent, embHigh, embeddingChunkWeightHigh);
        	}
        	VecUtil.normalizeArray(embContent);
    		List<Byte> embByteContent = VecUtil.toByteList(embContent, embByteScale);
    		embeddingContent.put(Constants.EMB, embByteContent);
    	}
    }

}