import time
import logging
import gradio as gr
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain_core.vectorstores import InMemoryVectorStore
from groq import Groq

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

client = Groq(api_key="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")
embedding_model = HuggingFaceEmbeddings(model_name="heydariAI/persian-embeddings")

# Initialize in-memory vector store for chat history
memory = InMemoryVectorStore()

def process_pdf_with_langchain(pdf_path):
    try:
        loader = PyPDFLoader(pdf_path)
        documents = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
        split_documents = text_splitter.split_documents(documents)
        
        vectorstore = FAISS.from_documents(split_documents, embedding_model)
        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
        return retriever
    except Exception as e:
        logger.error(f"Error processing PDF: {e}")
        raise

def generate_response(query, retriever=None):
    try:
        knowledge = ""

        if retriever:
            relevant_docs = retriever.get_relevant_documents(query)
            knowledge += "\n".join([doc.page_content for doc in relevant_docs])

        chat_history = memory.load_memory_variables({}).get("chat_history", "")
        context = "This is a conversation with ParvizGPT, an AI model designed by Amir Mahdi Parviz from KUT."

        if knowledge:
            context += f"\n\nRelevant Knowledge:\n{knowledge}"
        if chat_history:
            context += f"\n\nChat History:\n{chat_history}"

        context += f"\n\nYou: {query}\nParvizGPT:"

        # ابتدا یک پیام موقت نمایش داده شود
        response = "در حال پردازش..."

        retries = 3
        for attempt in range(retries):
            try:
                chat_completion = client.chat.completions.create(
                    messages=[{"role": "user", "content": context}],
                    model="deepseek-r1-distill-llama-70b"
                )
                response = chat_completion.choices[0].message.content.strip()
                memory.save_context({"input": query}, {"output": response})
                break
            except Exception as e:
                logger.error(f"Attempt {attempt + 1} failed: {e}")
                time.sleep(2)

        return response
    except Exception as e:
        logger.error(f"Error generating response: {e}")
        return f"Error: {e}"



        
def gradio_interface(user_message, chat_box, pdf_file=None):
    global retriever
    if pdf_file is not None:
        try:
            retriever = process_pdf_with_langchain(pdf_file.name)
        except Exception as e:
            return chat_box + [("Error", f"Error processing PDF: {e}")]

    chat_box.append(("ParvizGPT", "در حال پردازش..."))
    
    response = generate_response(user_message, retriever=retriever)
    
    chat_box[-1] = ("ParvizGPT", response) 
    chat_box.append(("You", user_message))
    
    return chat_box

def clear_memory():
    memory.clear()
    return []

retriever = None

with gr.Blocks() as interface:
    gr.Markdown("## ParvizGPT")
    chat_box = gr.Chatbot(label="Chat History", value=[])
    user_message = gr.Textbox(label="Your Message", placeholder="Type your message here and press Enter...", lines=1, interactive=True)
    clear_memory_btn = gr.Button("Clear Memory", interactive=True)
    pdf_file = gr.File(label="Upload PDF for Context (Optional)", type="filepath", interactive=True, scale=1)
    submit_btn = gr.Button("Submit")
    submit_btn.click(gradio_interface, inputs=[user_message, chat_box, pdf_file], outputs=chat_box)
    user_message.submit(gradio_interface, inputs=[user_message, chat_box, pdf_file], outputs=chat_box)
    clear_memory_btn.click(clear_memory, inputs=[], outputs=chat_box)

interface.launch()