import gradio as gr 
import os 
from tqdm import tqdm

from unstructured.partition.pdf import partition_pdf
from langchain.schema.document import Document
import google.generativeai as genai
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
import shutil

from wasabi import msg

from PIL import Image

genai.configure(api_key = 'AIzaSyB342Fh-nkRaO38BshbyI4-s0T9orVpsMw')
model = genai.GenerativeModel('gemini-1.5-flash')

files = os.listdir('uploads')
files = [f'uploads/{file}' for file in files]

documents = []

for file in tqdm(files , total = len(files) , leave = False) : 

    elements = partition_pdf(
        filename = file , 
        extract_images_in_pdf = True , 
        infer_table_structure = True , 
        chunking_strategy = 'by_title' , 
        max_characters = 4000 , 
        new_after_n_chars = 3800 , 
        combine_text_under_n_chars = 2000 , 
        extract_image_block_output_dir = 'outputs'
    )

    for element in elements : 

        element = element.to_dict()

        metadata = element['metadata']

        if 'text_as_html' in metadata : documents.append(
            Document(
                page_content = metadata['text_as_html'] , 
                metadata = {
                    'type' : 'text' , 
                    'metadata' : element
                }
            )
        )         

        else : documents.append(
            Document(
                page_content = element['text'] , 
                metadata = {
                    'type' : 'text' , 
                    'metadata' : element
                }
            )
        )

    images = os.listdir('outputs')
    images = [f'outputs/{image}' for image in images]

    for image in tqdm(images , total = len(images) , leave = False) : 

        image = Image.open(image)

        try : 

            response = model.generate_content([
                image , 
                'Explain the Image'
            ])

            response = response.text

        except Exception as e : msg.fail(f'----| FAIL : COULDNT CALL THE IMAGE DESCRIPTION API : {e}') ; response = 'COuldnt Call Model for this'

        documents.append(
            Document(
                page_content = response , 
                metadata = {
                    'type' : 'image' , 
                    'metadata' : {
                        'image' : image
                    }
                }
            )
        )

    shutil.rmtree('uploads')

vc = FAISS.from_documents(
    documents = documents , 
    embedding = HuggingFaceEmbeddings(model_name = 'all-MiniLM-L6-v2')
)

def run_rag(query) : 

    similar_docs = vc.similarity_search(query , k = 4)
    context = [doc.page_content for doc in similar_docs]

    prompt = f'''
    You are a Helpfull Chatbot that helps users with their queries

    - You will be provided with a query
    - You will be provided with a context as well 

    Your task is to generate a response to the query based on the context provided

    Context : {context}

    Query : {query}

    '''

    response = model.generate_content(prompt)

    return response.text

demo = gr.Interface(
    fn = run_rag , 
    inputs = 'text' , 
    outputs = 'text'
)

demo.launch()