import gradio as gr
import random
import time
import pymongo
import certifi
import os
from dotenv import load_dotenv


import argparse
from dataclasses import dataclass
from langchain.vectorstores.chroma import Chroma
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_openai.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

from deep_translator import GoogleTranslator

uri = "mongodb+srv://clementrof:t5fXqwpDQYFpvuCk@cluster0.rl5qhcj.mongodb.net/?retryWrites=true&w=majority"

# Create a new client and connect to the server
client = pymongo.MongoClient(uri, tlsCAFile=certifi.where())

# Send a ping to confirm a successful connection
try:
    client.admin.command('ping')
    print("Pinged your deployment. You successfully connected to MongoDB!")
except Exception as e:
    print(e)

# Access your database
db = client.get_database('camila')
records = db.info

# Load environment variables from .env
load_dotenv()

# Access the private key
private_key = os.getenv("OPENAI_API_KEY")
os.environ["OPENAI"] = "OPENAI_API_KEY"

CHROMA_PATH = "ch_chatbot"


####### F R ################
PROMPT_TEMPLATE = """
Réponds à la question en te basant sur le contexte suivant :

{context}

---

Voici l'historique de cette conversation, utilise l'historique comme une mémoire:

{memory}

---

Réponds à la question en se basant sur le contexte ci-dessus et parle de la même manière que le contexte. Ne dis pas que tu utilises le contexte pour répondre : {question}

"""





def message(question,memory):

    # Prepare the DB.
    embedding_function = OpenAIEmbeddings()
    db = Chroma(persist_directory=CHROMA_PATH,
                embedding_function=embedding_function)

    # Search the DB.
    results = db.similarity_search_with_relevance_scores(question, k=3)
    if len(results) == 0 or results[0][1] < 0.7:
        print("Unable to find matching results.")
        return

    context_text = "\n\n---\n\n".join(
        [doc.page_content for doc, _score in results])
    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=context_text, memory=memory, question=question)
    print(prompt)

    model = ChatOpenAI()
    response_text = model.invoke(prompt)
    content = response_text.content
    return content



def Chat_call(question):
    existing_user_doc = records.find_one({'ID': '1'})

    message_log = []
    messages = existing_user_doc['message']
    if len(messages)>1:
        messages = messages[-1:]

    message_log.extend(messages)
    # Convert each dictionary into a string representation
    message_strings = [f"{message['role']}: {message['content']}" for message in message_log]
    # Join the strings with newline characters
    memory = '\n'.join(message_strings)
            


    response = message(question,memory)


    records.update_one({'ID': '1'}, 
    {'$push':{'message':  {'role': 'user', 'content': f'{question}'}}})
    records.update_one({'ID': '1'}, 
    {'$push':{'message': {'role': 'assistant', 'content': f'{response}'}}})

   
    return response




with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, chat_history):
        bot_message = Chat_call(message)
        chat_history.append((message, bot_message))
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

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