import os
import pymongo
import ModelDef
from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader, set_global_service_context
from llama_index.llms.azure_openai import AzureOpenAI
from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
from llama_index.core.indices.vector_store.base import VectorStoreIndex
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
from llama_index.core import Settings
from llama_index.core.indices.vector_store import VectorStoreIndex
import os
import pymongo
import json
import ModelDef
from flask import Flask, request

# ---------------------------
# Constantes de l'application
# ---------------------------

# Définition des models de déploiement AI

_modelGPT4o = ModelDef.ModelGPT("gpt-4o", "gpt-4o", "2024-10-21")
_modelAda2 = ModelDef.ModelGPT("ada2", "text-embedding-ada-002", "2024-10-21")

_server4o = ModelDef.ServerGPT("gpt4o", "https://.azure.com", "26f3ea90247b1a9286057d53c2539", "c59f1006a64015a7b083ed29", "eastus2", _modelGPT4o, _modelAda2)

_models = [ _server4o]
_model = _models[0]

# Constants

_index = "LeJourSeLeveIFC"
_directory = "C:\\Projets\\ifcdoctest"
_mongoURI = os.environ["MONGO_URI"] = "url of mongo data base"

# ----------------------
# Démarrage du programme
# ----------------------

# Initialisation OpenAI

print("Initialisation OpenAI...")

llm = AzureOpenAI(
    #model=_model.ChatModel.Model,
    #deployment_name=_model.ChatModel.Name,
    api_key=_model.Key1,
    azure_endpoint=_model.Server,
    api_version=_model.ChatModel.ApiVersion,
)

embed_model = AzureOpenAIEmbedding(
    model=_model.LearningModel.Model,
    deployment_name=_model.LearningModel.Name,
    api_key=_model.Key1,
    azure_endpoint=_model.Server,
    api_version=_model.LearningModel.ApiVersion,
)
#Settings.llm = AzureOpenAI(model=llm)
#Settings.embed_model = AzureOpenAIEmbedding(model=_modelAda2)
#service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)
#set_global_service_context(service_context)
# Initialisation des paramètres pour les requètes sur MongoDB Atlas
print("Initialisation MongoDB...")

mongodb_client = pymongo.MongoClient(_mongoURI)
store = MongoDBAtlasVectorSearch(mongodb_client, db_name=_index)
storage_context = StorageContext.from_defaults(vector_store=store)

# On parcours chaque fichier
print("Démarrage de l'importation...")

reader = SimpleDirectoryReader(_directory, recursive=True, encoding="latin-1", required_exts=[".pdf", ".docx", ".pptx", ".csv", ".txt"])

for docs in reader.iter_data():
        print("F > " + docs[0].metadata['file_name'])
        VectorStoreIndex.from_documents(docs, storage_context=storage_context)
# Fin du programme

print("Terminée.")

        # Récupération info model

        model = GetModel(requestDTO.Model)

        # Initialisation OpenAI

        llm = AzureOpenAI(
            model=model.ChatModel.Model,
            deployment_name=model.ChatModel.Name,
            temperature=requestDTO.Temperature,
            api_key=model.Key1,
            azure_endpoint=model.Server,
            api_version=model.ChatModel.ApiVersion,
        )

        embed_model = AzureOpenAIEmbedding(
            model=model.LearningModel.Model,
            deployment_name=model.LearningModel.Name,
            api_key=model.Key1,
            azure_endpoint=model.Server,
            api_version=model.LearningModel.ApiVersion,
        )

        Settings.llm = llm
        Settings.embed_model = embed_model
        Settings.context_window = _contextWindow
        Settings.num_output = _numOutput

         # Génération de l'historique

        messages = []

        if (requestDTO.History != None):
            for item in requestDTO.History:
                messages.append(ChatMessage(role=item.Type.lower(), content=item.Prompt))
      
        # Initialisation des paramètres pour les requètes sur MongoDB Atlas

        mongodb_client = pymongo.MongoClient(_mongoURI)
        store = MongoDBAtlasVectorSearch(mongodb_client, db_name=requestDTO.Index)

        # Initialisation de l'index via les index sur MongoDB Atlas
        # Et inversement, commenter/décommenter si on veut juste query l'index existant
        # build index
        
        vector_index = VectorStoreIndex.from_vector_store(store)

        # configure retriever
        
        retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=requestDTO.LinkNumber)
        
        # configure response synthesizer, ici on peut chosir le mode réponse pour la query(ici on peut influer sur la maniere dont le contexte et le prompt sont itéré sur la query)
        
        response_synthesizer = get_response_synthesizer(response_mode=GetResponseMode(requestDTO.Mode), text_qa_template=qa_template)
        
        # assemble query engine, o recupere l'index grace au retriver et aussi le mode de reponse du "ResponseMode.TREE_SUMMARIZE" via response_synthesizer 
        
        query_retreiver = RetrieverQueryEngine(retriever=retriever, response_synthesizer=response_synthesizer)

        # Partie permettant de créer la réponse

        #query_engine = vector_index.as_query_engine(text_qa_template=qa_template)
        #gpt_result = query_engine.query(requestDTO.Prompt)

        gpt_result = query_retreiver.query(requestDTO.Prompt)

        resultDTO = ServiceQueryResponse.ServiceQueryResponseResultDTO(gpt_result.response, [])
        
        for item in gpt_result.source_nodes:
            node = ServiceQueryResponse.ServiceQueryResponseNodeDTO(item.node.extra_info.get("file_name"), item.node.extra_info.get("page_label"), item.node.text, item.score)
            resultDTO.Nodes.append(node)

        # Construction de la réponse

        responseDTO = ServiceQueryResponse.ServiceQueryResponseDTO(False, None, resultDTO)

        # Terminée, on envoi la réponse définitive

        return GenerateQueryResponse(requestDTO, responseDTO), 200
    
    except Exception as error:

        return str(error), 400

        mongodb_client = pymongo.MongoClient(_mongoURI)
        store = MongoDBAtlasVectorSearch(mongodb_client, db_name=indexName)

        # Initialisation de l'index via les index sur MongoDB Atlas
        # Et inversement, commenter/décommenter si on veut juste query l'index existant
        # build index
        index = VectorStoreIndex.from_vector_store(store)


        # configure retriever
        retriever = VectorIndexRetriever(
            index=index,
        )
        # configure response synthesizer
        response_synthesizer = get_response_synthesizer(
            response_mode= ResponseMode.TREE_SUMMARIZE,
        )
        # Template obligatoire du system prompt définissant le comprtement du LLM)

        qa_template = Prompt(requestDTO.get_system_template())
        
        if not qa_template or qa_template == "" or qa_template == None:
            responseDTO = ServiceQueryResponse.ServiceQueryResponseDTO(True, "Un historique de type 'System' est obligatoire.", None)        
            return GenerateQueryResponse(requestDTO, responseDTO), 400
        
        # assemble query engine
        query_engine = RetrieverQueryEngine(
            retriever=retriever,
            response_synthesizer=response_synthesizer,
        )

        #query_engine = query_engine.query(similarity_top_k=requestDTO.LinkNumber,text_qa_template=qa_template)


        # Partie permettant de créer la réponse

        query_engine = index.as_query_engine(similarity_top_k=requestDTO.LinkNumber,text_qa_template=qa_template)
        query_text = requestDTO.Prompt
        gpt_result = query_engine.query(query_text)



        resultDTO = ServiceQueryResponse.ServiceQueryResponseResultDTO(gpt_result.response, answer.message.content, [])
        
        for item in gpt_result.source_nodes:
            node = ServiceQueryResponse.ServiceQueryResponseNodeDTO(item.node.extra_info.get("file_name"), item.node.extra_info.get("page_label"), item.node.text, item.score)
            resultDTO.Nodes.append(node)



        # Construction de la réponse
        responseDTO = ServiceQueryResponse.ServiceQueryResponseDTO(False, None, resultDTO)
        # Terminée, on envoi la réponse définitive

        return GenerateQueryResponse(requestDTO, responseDTO), 200