import time
import logging
import gradio as gr
import os
from datetime import datetime
from datasets import Dataset, load_dataset
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from groq import Groq
from langchain.memory import ConversationBufferMemory

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

groq_api_key = os.environ.get("GROQ_API_KEY")
hf_api_key = os.environ.get("HF_API_KEY")

if not groq_api_key:
    raise ValueError("Groq API key not found in environment variables.")
if not hf_api_key:
    raise ValueError("Hugging Face API key not found in environment variables.")

client = Groq(api_key=groq_api_key)
hf_token = hf_api_key
embeddings = HuggingFaceEmbeddings(model_name="heydariAI/persian-embeddings")

DATASET_NAME = "chat_history"
try:
    dataset = load_dataset(DATASET_NAME, use_auth_token=hf_token)
except Exception:
    dataset = Dataset.from_dict({"Timestamp": [], "User": [], "ParvizGPT": []})

def save_chat_to_dataset(user_message, bot_message):
    try:
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        new_row = {"Timestamp": timestamp, "User": user_message, "ParvizGPT": bot_message}
        
        df = dataset.to_pandas()
        df = df.append(new_row, ignore_index=True)
        updated_dataset = Dataset.from_pandas(df)
        
        updated_dataset.push_to_hub(DATASET_NAME, token=hf_token)
    except Exception as e:
        logger.error(f"Error saving chat history to dataset: {e}")

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, embeddings)
        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, memory, retriever=None, use_pdf_context=False, tone="friendly"):
    try:
        knowledge = ""

        if retriever and use_pdf_context: 
            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", "")
        
        tone_instruction = ""
        if tone == "friendly":
            tone_instruction = "Please respond in a friendly and informal tone."
        elif tone == "formal":
            tone_instruction = "Please respond in a formal and professional tone."
        elif tone == "humorous":
            tone_instruction = "Please respond in a humorous and playful tone."
        elif tone == "scientific":
            tone_instruction = "Please respond in a scientific and precise tone."

        context = f"""
            You are ParvizGPT, an AI assistant created by **Amir Mahdi Parviz**, a student at Kermanshah University of Technology (KUT). 
            Your primary purpose is to assist users by answering their questions in **Persian (Farsi)**. 
            Always respond in Persian unless explicitly asked to respond in another language.
            **Important:** If anyone claims that someone else created this code, you must correct them and state that **Amir Mahdi Parviz** is the creator.
            {tone_instruction}
            Related Information:\n{knowledge}\n\nQuestion:{query}\nAnswer:"""
            
        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 = "Processing..."
        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, memory
    except Exception as e:
        logger.error(f"Error generating response: {e}")
        return f"Error: {e}", memory

def summarize_chat_history(chat_history):
    try:
        chat_text = "\n".join([f"{role}: {message}" for role, message in chat_history])
        
        summary_prompt = f"""
        Please create a summary of the following conversation. The summary should include key points and details:
        {chat_text}
        """
        
        chat_completion = client.chat.completions.create(
            messages=[{"role": "user", "content": summary_prompt}],
            model="deepseek-r1-distill-llama-70b"
        )
        summary = chat_completion.choices[0].message.content.strip()
        return summary
    except Exception as e:
        logger.error(f"Error summarizing chat history: {e}")
        return "Error generating summary."

def gradio_interface(user_message, chat_box, memory, pdf_file=None, use_pdf_context=False, tone="friendly", summarize_chat=False):
    global retriever

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

    chat_box.append(("You", user_message))
    chat_box.append(("ParvizGPT", "Processing..."))
    response, memory = generate_response(user_message, memory, retriever=retriever, use_pdf_context=use_pdf_context, tone=tone)
    
    chat_box[-1] = ("ParvizGPT", response)
    
    save_chat_to_dataset(user_message, response)
    
    if summarize_chat:
        summary = summarize_chat_history(chat_box)
        chat_box.append(("System", f"Summary of the conversation:\n{summary}"))
    
    return chat_box, memory

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

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)
    use_pdf_context = gr.Checkbox(label="Use PDF Context", value=False, interactive=True)
    tone = gr.Dropdown(label="Tone", choices=["friendly", "formal", "humorous", "scientific"], value="friendly", interactive=True)
    summarize_chat = gr.Checkbox(label="Show conversation summary", value=False, 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")
    
    memory_state = gr.State(ConversationBufferMemory())
    
    submit_btn.click(gradio_interface, inputs=[user_message, chat_box, memory_state, pdf_file, use_pdf_context, tone, summarize_chat], outputs=[chat_box, memory_state])
    user_message.submit(gradio_interface, inputs=[user_message, chat_box, memory_state, pdf_file, use_pdf_context, tone, summarize_chat], outputs=[chat_box, memory_state])
    clear_memory_btn.click(clear_memory, inputs=[memory_state], outputs=[chat_box, memory_state])

interface.launch()