import os  
from dotenv import load_dotenv  
import gradio as gr  
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings  
from llama_index.llms.huggingface import HuggingFaceInferenceAPI  
from llama_index.embeddings.huggingface import HuggingFaceEmbedding  

load_dotenv()  

# Configure the Llama index settings  
Settings.llm = HuggingFaceInferenceAPI(  
    model_name="meta-llama/Meta-Llama-3-8B-Instruct",  
    tokenizer_name="meta-llama/Meta-Llama-3-8B-Instruct",  
    context_window=3000,  
    token=os.getenv("HF_TOKEN"),  
    max_new_tokens=512,  
    generate_kwargs={"temperature": 0.1},  
)  

Settings.embed_model = HuggingFaceEmbedding(  
    model_name="BAAI/bge-small-en-v1.5"  
)  

# Define the directory for persistent storage and data  
PERSIST_DIR = "db"  
PDF_DIRECTORY = 'data'  # Directory containing PDFs  

# Ensure directories exist  
os.makedirs(PDF_DIRECTORY, exist_ok=True)  
os.makedirs(PERSIST_DIR, exist_ok=True)  

# Variable to store current chat conversation  
current_chat_history = []  

def data_ingestion_from_directory():  
    try:  
        # Use SimpleDirectoryReader on the directory containing the PDF files  
        documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data()  
        storage_context = StorageContext.from_defaults()  
        index = VectorStoreIndex.from_documents(documents)  
        index.storage_context.persist(persist_dir=PERSIST_DIR)  
    except Exception as e:  
        print(f"Error during data ingestion: {e}")  

def handle_query(query):  
    chat_text_qa_msgs = [  
        (  
            "user",  
            """  
            You are the JackNJill Solutions chatbot. Your goal is to provide accurate, professional, and helpful answers to user queries based on the company's data. Always ensure your responses are clear and concise.   

            Context:  
            {context_str}  

            Question:  
            {query_str}  
            """  
        )  
    ]  

    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)  

    # Load index from storage  
    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)  
    index = load_index_from_storage(storage_context)  

    # Use chat history to enhance response  
    context_str = "\n".join([f"User asked: '{past_query}'\nBot answered: '{response}'"  
                             for past_query, response in reversed(current_chat_history) if past_query.strip()])  

    query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str)  
    answer = query_engine.query(query)  

    response = answer.response if hasattr(answer, 'response') else answer.get('response', "I'm sorry, I can't answer that.")  

    # Remove sensitive information and unwanted sections from the response  
    sensitive_keywords = [PERSIST_DIR, PDF_DIRECTORY, "/", "\\", ".pdf", ".doc", ".txt"]  
    for keyword in sensitive_keywords:  
        response = response.replace(keyword, "")  

    # Remove sections starting with specific keywords  
    unwanted_sections = ["Page Label", "Page Label:", "page_label", "page_label:", "file_path:", "file_path"]  
    for section in unwanted_sections:  
        if section in response:  
            response = response.split(section)[0]  

    # Additional cleanup for any remaining artifacts from replacements  
    response = ' '.join(response.split())  

    # Update current chat history  
    current_chat_history.append((query, response))  

    return response  

# Example usage: Process PDF ingestion from directory  
print("Processing PDF ingestion from directory:", PDF_DIRECTORY)  
data_ingestion_from_directory()  

# Define the input and output components for the Gradio interface  
input_component = gr.Textbox(  
    show_label=False,  
    placeholder="Ask me anything about JackNJill Solutions..."  
)  

output_component = gr.Textbox()  

# Function to handle queries  
def chatbot_handler(query):  
    response = handle_query(query)  
    return response  

# Create the Gradio interface  
interface = gr.Interface(  
    fn=chatbot_handler,  
    inputs=input_component,  
    outputs=output_component,  
    title="Welcome to JackNJill Solutions",  
    description="I am here to assist you with any questions you have about JackNJill Solutions. How can I help you today?"  
)  

# Launch the Gradio interface  
interface.launch()