import gradio as gr
import os
from PIL import Image
import pytesseract
from pdf2image import convert_from_path
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain_groq import ChatGroq
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter

# Initialize the Groq API Key and the model
os.environ["GROQ_API_KEY"] = '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'
llm = ChatGroq(
    model='llama3-70b-8192',
    temperature=0.5,
    max_tokens=None,
    timeout=None,
    max_retries=2
)

# OCR functions
def ocr_image(image_path, language='eng+guj'):
    img = Image.open(image_path)
    return pytesseract.image_to_string(img, lang=language)

def ocr_pdf(pdf_path, language='eng+guj'):
    images = convert_from_path(pdf_path)
    all_text = "\n".join(pytesseract.image_to_string(img, lang=language) for img in images)
    return all_text

def ocr_file(file_path):
    ext = os.path.splitext(file_path)[1].lower()
    if ext == ".pdf":
        return ocr_pdf(file_path)
    elif ext in [".jpg", ".jpeg", ".png", ".bmp"]:
        return ocr_image(file_path)
    else:
        return "Unsupported file format."

def get_text_chunks(text):
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
    return splitter.split_text(text)

def get_vector_store(chunks):
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vector_store = FAISS.from_texts(chunks, embedding=embeddings)
    os.makedirs("faiss_index", exist_ok=True)
    vector_store.save_local("faiss_index")
    return vector_store

# Conversational chain
def get_conversational_chain():
    template = """<Insert your prompt here>"""
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-MiniLM-L6-v2")
    vector_store = FAISS.load_local("faiss_index", embeddings)
    qa_chain = RetrievalQA.from_chain_type(
        llm,
        retriever=vector_store.as_retriever(),
        chain_type='stuff',
        verbose=True,
        chain_type_kwargs={
            "prompt": PromptTemplate(input_variables=["history", "context", "question"], template=template),
            "memory": ConversationBufferMemory(memory_key="history", input_key="question"),
        }
    )
    return qa_chain

# File and question handling
def process_files(files, question):
    text = ""
    for file in files:
        file_path = os.path.join("temp", file.name)
        with open(file_path, "wb") as f:
            f.write(file.read())
        text += ocr_file(file_path) + "\n"

    chunks = get_text_chunks(text)
    vector_store = get_vector_store(chunks)

    qa_chain = get_conversational_chain()
    response = qa_chain({"query": question})
    return response.get("result", "No result found.")

# Gradio Interface
def app(files, question):
    return process_files(files, question)

iface = gr.Interface(
    fn=app,
    inputs=[gr.File(file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp"], label="Upload Files"), gr.Textbox(label="Ask a Question")],
    outputs="text",
    title="OCR and Document Query System",
    description="Upload PDF or image files and ask questions based on their content."
)

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