import gradio as gr
import tiktoken
import os
import numpy as np
import requests
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer
from langchain_community.chat_models import ChatOllama
from basesrc.strategies import convert_pdf_to_text, split_pages_into_chunks, vectorize
from prompt_query import generate_prompt

def process_pdf(pdf_url, embedding_model, llm_model, top_k=5):
    response = requests.get(pdf_url)
    pdf_filename = pdf_url.split('/')[-1]
    
    with open(pdf_filename, 'wb') as f:
        f.write(response.content)
    
    pages = convert_pdf_to_text(pdf_filename)
    print(f"Je suis Llama3, de l'équipe DREAMS TEAM, votre assistant QA pour répondre à vos questions liés aux documents 🙂 !\nDéjà pour info, le nombre de pages de vote document est: {len(pages)}")
    
    tokenizer = tiktoken.get_encoding("cl100k_base")
    chunks = split_pages_into_chunks(pages, 128, tokenizer)
    embedding_model = SentenceTransformer(embedding_model)

    knowledge_base = vectorize(chunks, embedding_model)
    chunks, embeddings = list(zip(*knowledge_base))
    corpus_embeddings = np.vstack(embeddings)

    llm_model = ChatOllama(model=llm_model)

    def chat_interface(question):
        return generate_prompt(question, chunks, corpus_embeddings, embedding_model, llm_model, top_k)

    return chat_interface, pdf_filename

def create_interface():
    def main(pdf_url, question):
        embedding_model = "Sahajtomar/french_semantic"
        llm_model = "llama3"
        top_k = 5
        chat_interface, pdf_filename = process_pdf(pdf_url, embedding_model, llm_model, top_k)
        response = chat_interface(question)
        return response

    examples = [
        ["https://razvanbarbulescu.pages.math.cnrs.fr/inequalities.pdf", "Quel est ton rôle ?"],
        #["https://maths-olympiques.fr/wp-content/uploads/2019/11/Inegalites-Theo.pdf", "Que nous parle ce document ?"],
        # ["https://hellofuture.orange.com/app/uploads/2024/05/2024-Orange-white-paper-on-Mobile-Network-Technology-Evolutions-Beyond-2030.pdf", "Fais-moi un résumé des éléments essentiels de ce document"],
        # ["https://hellofuture.orange.com/app/uploads/2024/05/2024-Orange-white-paper-on-Mobile-Network-Technology-Evolutions-Beyond-2030.pdf", "Comment puis-je présenter ce document dans une réunion d'équipe de professionnels ?"]
    ]

    iface = gr.Interface(
        fn=main,
        inputs=[gr.Textbox(label="URL du PDF", placeholder="Entrez l'URL de votre PDF ici..."),
                gr.Textbox(lines=3, label="✨Posez votre question en français 😎", placeholder="Posez une question...")],
        outputs=gr.Textbox(label="Réponse"),
        title="Assistant de chat sur le document | by PSW",
        description="Entrez l'URL de votre PDF et posez votre question en français pour obtenir une réponse basée sur le contexte du document.",
        examples=examples
    )

    iface.launch(share=True)

if __name__ == "__main__":
    create_interface()