import gradio as gr
import os
import subprocess
import fitz
from dotenv import load_dotenv
from langchain_community.document_loaders import UnstructuredPDFLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import CharacterTextSplitter
from langchain_groq import ChatGroq
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from gtts import gTTS
import sys
import pytesseract
from pdf2image import convert_from_path
from huggingface_hub import Repository, login
from huggingface_hub import hf_hub_download




# Load environment variables
load_dotenv()
secret_key = os.getenv("GROQ_API_KEY")
hf_key = os.getenv("HF_TOKEN")

os.environ["GROQ_API_KEY"] = secret_key
login(token=hf_key,add_to_git_credential=True)

embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2")

# Ensure the necessary folders exist
UPLOAD_FOLDER = 'uploads/'
AUDIO_FOLDER = 'audio/'
for folder in [UPLOAD_FOLDER, AUDIO_FOLDER]:
    if not os.path.exists(folder):
        os.makedirs(folder)


def load_pdf(file_path):
    """Load and preprocess Arabic text from a PDF file."""
    pages = convert_from_path(file_path, 500)
    documents = []
    for pageNum, imgBlob in enumerate(pages):
        text = pytesseract.image_to_string(imgBlob, lang="ara")
        documents.append(text)
    return documents


import os
from langchain.vectorstores import FAISS
from huggingface_hub import Repository

def save_faiss_index_to_hub(vectorstore, repo_id="MohammedNasser/faiss-index"):
    index_dir = "faiss_index"
    
    # Ensure the index directory exists
    if not os.path.exists(index_dir):
        os.makedirs(index_dir)
    
    # Save FAISS index locally
    vectorstore.save_local(index_dir)
    
    # Initialize Hugging Face repository
    repo = Repository(local_dir=index_dir, clone_from=repo_id, repo_type="dataset")
    
    # Push the FAISS index files to the Hugging Face Hub
    repo.push_to_hub(commit_message="Pushing FAISS index")

    print(f"FAISS index saved to Hugging Face Hub: {repo_id}")

def prepare_vectorstore(data):
    index_dir = "faiss_index"
    if not os.path.exists(index_dir):
        os.makedirs(index_dir)
    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=20, separator="\n")
    texts = data
    vectorstore = FAISS.from_texts(texts, embeddings)
    save_faiss_index_to_hub(vectorstore)
    return vectorstore
    

def load_vectorstore(repo_id="MohammedNasser/faiss-index"):
    index_dir = "faiss_index"
    
    # Ensure the index directory exists
    if not os.path.exists(index_dir):
        os.makedirs(index_dir)

    
    
    # Download the FAISS index files from Hugging Face Hub
    hf_hub_download(repo_id=repo_id, filename="index.faiss", local_dir=index_dir, repo_type="dataset")
    hf_hub_download(repo_id=repo_id, filename="index.json", local_dir=index_dir, repo_type="dataset")
    vectorstore = FAISS.load_local(index_dir, embeddings, allow_dangerous_deserialization=True)
    return vectorstore

def create_chain(vectorstore):
    llm = ChatGroq(model="gemma2-9b-it", temperature=0)
    retriever = vectorstore.as_retriever()
    memory = ConversationBufferMemory(llm=llm, output_key="answer", memory_key="chat_history", return_messages=True)
    chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory,
        verbose=False,
        chain_type="map_reduce"
    )
    return chain

def process_pdf(pdf_file):
    file_path = os.path.join(UPLOAD_FOLDER, pdf_file.name)
    with open(file_path, "wb") as f:
        f.write(pdf_file.read())
    data = load_pdf(file_path)
    vectorstore = prepare_vectorstore(data)
    return "PDF processed successfully. You can now start chatting!"

def chat(user_input, history):
    vectorstore = load_vectorstore()
    chain = create_chain(vectorstore)

    prompt = f"""
    You are an expert Arabic-language assistant specialized in analyzing and responding to queries about Arabic PDF documents. Your responses should be precise, informative, and reflect the professional tone and structure expected in formal Arabic communication. Focus on extracting and presenting relevant information from the document clearly and systematically, while avoiding colloquial or informal language.

    When responding, ensure the following:
       - Your answer directly reflects the content of the document.
       - If the requested information is not available in the document, clearly state that.
       - Keep your response concise yet comprehensive, addressing the question fully.
       - Always respond in formal Arabic, without using English.

    Question: {user_input}
    Helpful Answer:"""

    response = chain({"question": prompt})
    assistant_response = response["answer"]

    # Generate audio file
    tts = gTTS(text=assistant_response, lang='ar')
    audio_file = f"response_{len(history)}.mp3"
    tts.save(os.path.join(AUDIO_FOLDER, audio_file))

    return assistant_response, audio_file

custom_css = """
body {
    font-family: 'Noto Kufi Arabic', sans-serif;
    background: linear-gradient(135deg, #799351 0%, #A67B5B 100%);
    background-size: cover;
    background-position: center;
    background-attachment: fixed;
}

.gradio-container {
    max-width: 800px !important;
    margin: auto !important;
    background: rgba(255, 255, 255, 0.9);
    border-radius: 20px;
    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37);
    backdrop-filter: blur(4px);
    border: 1px solid rgba(255, 255, 255, 0.18);
    padding: 20px;
}

h1, h2, h3 {
    color: #1A4D2E;
    font-weight: bold;
    text-align: center;
}

p {
    color: #A89F91;
}

.gradio-button {
    background-color: #5F6F65 !important;
    color: #FFFFFF !important;
}

.gradio-button:hover {
    background-color: #FFFFFF !important;
    color: #5F6F65 !important;
}

.chat-message {
    border-radius: 10px;
    padding: 10px;
    margin-bottom: 10px;
}

.chat-message.user {
    background-color: #E7F0DC;
}

.chat-message.bot {
    background-color: #F7EED3;
}

.chat-message::before {
    content: '';
    display: inline-block;
    width: 24px;
    height: 24px;
    background-size: contain;
    background-repeat: no-repeat;
    margin-right: 10px;
    vertical-align: middle;
}

.chat-message.user::before {
    content: '👤';
}

.chat-message.bot::before {
    content: '🤖';
}
"""
# Gradio interface
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("# ديمو بوت للقاء مركز حضرموت للدراسات التاريخية")
    gr.Markdown("## المنعقد السبت 14 - سبتمبر 2024")
    
    with gr.Row():
        pdf_input = gr.File(label="اختر ملف PDF للدردشة")
        process_button = gr.Button("رفع وبدء الدردشة")
    
    # Let Gradio manage the Textbox internally within ChatInterface
    chat_interface = gr.ChatInterface(
        chat,
        title="الدردشة مع البوت",
        description="اسأل أي سؤال عن محتوى الملف PDF",
        theme="soft",
        examples=["ما هو موضوع الوثيقة؟", "من هم الأشخاص المذكورون؟", "ما هي التواريخ الرئيسية المذكورة؟"],
        cache_examples=True,
        retry_btn=None,
        undo_btn="مسح آخر رسالة",
        clear_btn="مسح المحادثة",
    )

    audio_output = gr.Audio(label="الرد الصوتي")
    
    process_button.click(process_pdf, inputs=[pdf_input], outputs=[chat_interface.textbox])

    # Use the internal Textbox and Chatbot provided by the ChatInterface
    chat_interface.submit(
        fn=lambda x, y: y[-1][1],
        inputs=[chat_interface.textbox, chat_interface.chatbot],
        outputs=[audio_output]
    )

demo.launch()