import os

import gradio as gr

from dotenv import load_dotenv

from qdrant_client import models, QdrantClient
from sentence_transformers import SentenceTransformer
from openai import OpenAI


################## Set-up ##################

load_dotenv()  # load environment variables

embedding_model = SentenceTransformer("antoinelouis/french-multilingual-e5-small")

db_client = QdrantClient(":memory:")

db_client.create_collection(
    collection_name="IncubAI",
    vectors_config=models.VectorParams(
        size=embedding_model.get_sentence_embedding_dimension(),  # Vector size is defined by used model
        distance=models.Distance.COSINE,
    ),
)


################# Filling the DB ################

documents = [
    {
        "page_content": "La Tour Eiffel a été construite en 1889 comme pièce maîtresse de l'Exposition universelle à Paris.",
        "source": "encyclopedia"
    },
    {
        "page_content": "Pour préparer un bon café filtre, il est crucial de choisir une mouture moyenne et d’utiliser de l’eau à 92–96°C.",
        "source": "blog"
    },
    {
        "page_content": "Les bienfaits de la méditation incluent la réduction du stress, l'amélioration de la concentration et un meilleur sommeil.",
        "source": "health_guide"
    },
    {
        "page_content": "Le framework LangChain permet d’intégrer facilement la recherche documentaire dans les applications alimentées par des modèles de langage.",
        "source": "developer_docs"
    },
    {
        "page_content": "Selon une étude de 2023, les énergies renouvelables ont représenté 30 % de la production mondiale d'électricité.",
        "source": "news"
    },
    {
        "page_content": "Un bon pitch d’ascenseur doit être clair, concis et présenter en moins de 30 secondes ce qui rend votre idée unique.",
        "source": "business_tips"
    },
    {
        "page_content": "Le Louvre abrite des œuvres célèbres comme la Joconde, la Vénus de Milo et le Sacre de Napoléon.",
        "source": "museum_guide"
    },
    {
        "page_content": "Si vous voulez apprendre Python, commencez par comprendre les variables, les boucles et les fonctions.",
        "source": "tutorial"
    },
    {
        "page_content": "Les voitures électriques produisent moins d’émissions locales, mais leur impact écologique dépend de la production des batteries.",
        "source": "science_article"
    },
    {
        "page_content": "« Interstellar » explore les thèmes de la relativité temporelle, de l’amour et de la survie humaine à travers une aventure spatiale.",
        "source": "movie_review"
    }
]


# Upload to DB
for id, doc in enumerate(documents):

    # 1. embed the content
    doc_embedding = embedding_model.encode(doc["page_content"]).tolist()
    
    # 2. upload the document
    db_client.upload_points(
        collection_name="IncubAI",
        points=[
            models.PointStruct(
                id=id, vector=doc_embedding, payload=doc
            )
        ],
    )


################# Preparing Generation ################

template = """Répondez à la question en vous basant uniquement sur le contexte suivant:
{context}

Question : {question}
"""

llm_client = OpenAI(
    # This is the default and can be omitted
    api_key=os.environ.get("OPENAI_API_KEY"),
)


################# Combining Retrieval and Generation for RAG ################

def execute_RAG(user_query):
    query_embedding = embedding_model.encode(user_query).tolist()

    results = db_client.query_points(
        collection_name="IncubAI",
        query = query_embedding,
        limit=2
    )
    
    context = "\n\n".join(result.payload["page_content"] for result in results.points)
    prompt = template.format(context=context, question=user_query)
    response = llm_client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3
        )

    return response.choices[0].message.content



################# Prepare Interface ################

def respond(message, chat_history):
    
    bot_response = execute_RAG(message)
    
    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": bot_response})
    return "", chat_history


examples = [
    "Quand la Tour Eiffel a-t-elle été construite ?",
    "Quels sont les effets de la méditation ?",
    "Comment fonctionne LangChain ?",
    "Quel est l’impact écologique des voitures électriques ?",
]

msg = gr.Textbox()  # define outside to define location


with gr.Blocks() as demo:
    gr.HTML("<h1 style='text-align:center'>Example de un Chatbot minimal 🤖</h1>")
    chatbot = gr.Chatbot(type="messages")
    gr.Examples(examples, msg)
    msg.render()
    clear = gr.ClearButton([msg, chatbot])
    
    msg.submit(respond, [msg, chatbot], [msg, chatbot])


if __name__ == "__main__":
    demo.launch()