package net.byteboost.duck.utils;

import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.Path;
import java.nio.file.Paths;

import dev.langchain4j.chain.ConversationalRetrievalChain;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.ParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.huggingface.HuggingFaceChatModel;
import dev.langchain4j.model.huggingface.HuggingFaceEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.retriever.EmbeddingStoreRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import net.byteboost.duck.ApiKeys;
import net.byteboost.duck.App;

import static java.time.Duration.ofSeconds;


/**
 * Contains all the AI utilities used in Duck like questioning with the AI consultant.
 */
public class  AIutils {
    
    //Transforms DOCX,PDF,TXT,PPT to Document
    public static Document toDoc(String file){
    Path filePath = toPath(file);
    return FileSystemDocumentLoader.loadDocument(filePath);
    }

    //Loads the file and asks a question to the AI that returns the answer
    public static String loadIntoOpenAI(Document file, String question){

        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.withApiKey(ApiKeys.OPENAI_API_KEY);

        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .splitter(new ParagraphSplitter())
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        ingestor.ingest(file);

        ConversationalRetrievalChain chain = ConversationalRetrievalChain.builder()
                .chatLanguageModel(OpenAiChatModel.withApiKey(ApiKeys.OPENAI_API_KEY))
                .retriever(EmbeddingStoreRetriever.from(embeddingStore, embeddingModel))
                // .chatMemory() // you can override default chat memory
                // .promptTemplate() // you can override default prompt template
                .build();

        return chain.execute(question);
    }
    public static String loadIntoHugging(Document file, String question){

        EmbeddingModel embeddingModel = HuggingFaceEmbeddingModel.builder()
                .accessToken(ApiKeys.HF_API_KEY)
                .modelId("sentence-transformers/all-MiniLM-L6-v2")
                .waitForModel(true)
                .timeout(ofSeconds(60))
                .build();

        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();


        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .splitter(new ParagraphSplitter())
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        ingestor.ingest(file);


        ConversationalRetrievalChain chain = ConversationalRetrievalChain.builder()
                .chatLanguageModel(HuggingFaceChatModel.withAccessToken(ApiKeys.HF_API_KEY))
                .retriever(EmbeddingStoreRetriever.from(embeddingStore, embeddingModel))
                // .chatMemory() // you can override default chat memory
                // .promptTemplate() // you can override default prompt template
                .build();

        return chain.execute(question);
    }

    private static Path toPath(String fileName) {
        try {
            URL fileUrl = App.class.getResource(fileName);
            return Paths.get(fileUrl.toURI());
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }
}



[((1783, 1987), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((1783, 1962), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((1783, 1914), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((1783, 1866), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((2066, 2459), 'dev.langchain4j.chain.ConversationalRetrievalChain.builder'), ((2066, 2280), 'dev.langchain4j.chain.ConversationalRetrievalChain.builder'), ((2066, 2191), 'dev.langchain4j.chain.ConversationalRetrievalChain.builder'), ((2623, 2875), 'dev.langchain4j.model.huggingface.HuggingFaceEmbeddingModel.builder'), ((2623, 2850), 'dev.langchain4j.model.huggingface.HuggingFaceEmbeddingModel.builder'), ((2623, 2810), 'dev.langchain4j.model.huggingface.HuggingFaceEmbeddingModel.builder'), ((2623, 2774), 'dev.langchain4j.model.huggingface.HuggingFaceEmbeddingModel.builder'), ((2623, 2707), 'dev.langchain4j.model.huggingface.HuggingFaceEmbeddingModel.builder'), ((3007, 3211), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((3007, 3186), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((3007, 3138), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((3007, 3090), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((3291, 3690), 'dev.langchain4j.chain.ConversationalRetrievalChain.builder'), ((3291, 3511), 'dev.langchain4j.chain.ConversationalRetrievalChain.builder'), ((3291, 3422), 'dev.langchain4j.chain.ConversationalRetrievalChain.builder'), ((3830, 3861), 'net.byteboost.duck.App.class.getResource')]