import os
from langchain_groq import ChatGroq
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import create_retrieval_chain
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from dotenv import load_dotenv
import gradio as gr

# Load environment variables
load_dotenv()
groq_api_key = os.getenv('GROQ_API_KEY')

if not groq_api_key:
    raise ValueError("GROQ_API_KEY not found in environment variables. Please add it to your .env file.")

# Global variables
conversation_history = []

# Initialize chatbot components
llm = ChatGroq(groq_api_key=groq_api_key, model_name="mixtral-8x7b-32768")
prompt = ChatPromptTemplate.from_template(
    """Answer the questions based on the provided context only. Please provide the most accurate response based on the question <context> {context} </context> Questions:{input}"""
)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
loader = PyPDFLoader("TumorPDf.pdf")
docs = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
final_documents = text_splitter.split_documents(docs)
doc_texts = [doc.page_content for doc in final_documents]
embeddings_result = embeddings.embed_documents(doc_texts)
if embeddings_result:
    vectors = FAISS.from_documents(final_documents, embeddings)
else:
    raise ValueError("Failed to generate embeddings. Please check your input documents or try a different embedding model.")

document_chain = create_stuff_documents_chain(llm, prompt)
retriever = vectors.as_retriever()
retrieval_chain = create_retrieval_chain(retriever, document_chain)

def add_message(history, message):
    history.append((message, None))
    return history, gr.Textbox(value="", interactive=True)

def bot(history):
    global conversation_history
    
    if not history:
        initial_message = "Hello! How can I assist you today?"
        history.append((None, initial_message))
        conversation_history.append(f"AI: {initial_message}")
        return history

    message = history[-1][0]
    
    conversation_history.append(f"Human: {message}")
    
    full_context = "\n".join(conversation_history)
    
    try:
        response = retrieval_chain.invoke({
            'input': f"{full_context}\nHuman: {message}",
        })['answer']
    except Exception as e:
        response = f"I'm sorry, but I encountered an error: {str(e)}"
    
    conversation_history.append(f"AI: {response}")
    
    history[-1] = (history[-1][0], response)
    return history

def clear_chat():
    global conversation_history
    conversation_history = []
    return []

with gr.Blocks() as demo:
    with gr.Column():
        chatbot = gr.Chatbot([], elem_id="chatbot")
        msg = gr.Textbox(label="Message")
        clear = gr.Button("Clear")

        chat_msg = msg.submit(add_message, [chatbot, msg], [chatbot, msg], queue=False).then(
            bot, chatbot, chatbot
        )

        clear.click(clear_chat, None, chatbot, queue=False)

    demo.queue()

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