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
from sentence_transformers import SentenceTransformer
import spaces

MARKDOWN = """
This demo utilizes <a href="https://huggingface.co/meta-llama">LLaMA 3 8B Instructor Model</a> by Meta LLaMA.

Furthermore, the feature extraction is performed using <a href="https://huggingface.co/BAAI/bge-large-zh-v1.5"> BGE Model Series </a> by BAAI.

I am looking for more specific data to refine the responses of the chatbot, so if any specialist wants to collaborate, you are welcome to do so. My details are provided below.

**Demo by [Sunder Ali Khowaja](https://sander-ali.github.io) - [X](https://x.com/SunderAKhowaja) -[Github](https://github.com/sander-ali) -[Hugging Face](https://huggingface.co/SunderAli17)**

**The Savior Bot is here to assist you with any questions you have about Heart Disease Preventions. How can the Savior Bot help you?"**
"""

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-large-zh-v1.5"
)

# Define the directory for persistent storage and data
PERSIST_DIR = "db"
PDF_DIRECTORY = 'data'  # Changed to the 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 = []

#@spaces.GPU


def data_ingestion_from_directory():
    # 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)

def handle_query(query):
    chat_text_qa_msgs = [
        (
            "user",
            """
            You are the Savior Bot. Your goal is to provide accurate, preventions, and helpful answers to user queries based on the available 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 = ""
    for past_query, response in reversed(current_chat_history):
        if past_query.strip():
            context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"

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

    if hasattr(answer, 'response'):
        response = answer.response
    elif isinstance(answer, dict) and 'response' in answer:
        response = answer['response']
    else:
          response = "Sorry, SAK solutions are trying to improve my knowledge further, however, as per my current knowledge I am unable to answer this question. Is there anything else I can help you with?"

    # 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="Savior Bot is at your service ... Let me know what you are feeling"
)

output_component = gr.Textbox()

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

theme = gr.themes.Soft(
    font=[gr.themes.GoogleFont('Bree Serif'), gr.themes.GoogleFont('Public Sans'), 'system-ui', 'sans-serif'],
)
js_func = """
function refresh() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""

# Create the Gradio interface
interface = gr.Interface(
    fn=chatbot_handler,
    inputs=input_component,
    outputs=output_component,
    title="Welcome to SAK solutions",
    description=MARKDOWN,
    theme = theme,
    js = js_func
)

# Launch the Gradio interface
interface.launch(share=True)