import gradio as gr
import os
from dotenv import load_dotenv
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 simple_salesforce import Salesforce, SalesforceLogin
import random
import datetime

# Load environment variables
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'

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

# Variable to store current chat conversation in a dictionary
current_chat_history = {}
kkk = random.choice(['Clara', 'Lily'])

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 Lily Redfernstech 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. give response within 10-15 words only       
           {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.values()):
        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, I couldn't find an answer."

    # Update current chat history dictionary (use unique ID as key)
    chat_id = str(datetime.datetime.now().timestamp())
    current_chat_history[chat_id] = (query, response)

    return response

def save_chat_history_to_salesforce():
    username =os.getenv("username")
    password =os.getenv("password")
    security_token =os.getenv("security_token")
    domain = 'test' 

    # Log in to Salesforce
    session_id, sf_instance = SalesforceLogin(username=username, password=password, security_token=security_token, domain=domain)

    # Create Salesforce object
    sf = Salesforce(instance=sf_instance, session_id=session_id)

    # Iterate over chat history dictionary and push to Salesforce
    for chat_id, (user_message, bot_response) in current_chat_history.items():
        data = {
            'Name': 'Chat with user',
            'Bot_Message__c': bot_response,
            'User_Message__c': user_message,
            'Date__c': str(datetime.datetime.now().date())
        }
        # Insert into the custom object (replace 'Chat_History__c' with your custom object's API name)
        sf.Chat_History__c.create(data)

# Define the function to handle predictions
def predict(message, history):
    logo_html = '''
    <div class="circle-logo">
      <img src="https://rb.gy/8r06eg" alt="FernAi">
    </div>
    '''
    response = handle_query(message)
    response_with_logo = f'<div class="response-with-logo">{logo_html}<div class="response-text">{response}</div></div>'
    
    # Return the response with the logo
    return response_with_logo

# Define your Gradio chat interface function
def chat_interface(message, history):
    try:
        # Process the user message and generate a response
        response = handle_query(message)

        # Return the bot response
        return response
    except Exception as e:
        return str(e)

# Custom CSS for styling
css = '''
  .circle-logo {
  display: inline-block;
  width: 40px;
  height: 40px;
  border-radius: 50%;
  overflow: hidden;
  margin-right: 10px;
  vertical-align: middle;
}
.circle-logo img {
  width: 100%;
  height: 100%;
  object-fit: cover;
}
.response-with-logo {
  display: flex;
  align-items: center;
  margin-bottom: 10px;
}
footer {
    display: none !important;
    background-color: #F8D7DA;
  }
label.svelte-1b6s6s {display: none}
div.svelte-rk35yg {display: none;}
div.progress-text.svelte-z7cif2.meta-text {display: none;}
'''

# Use Gradio Blocks to wrap components
with gr.Blocks() as demo:
    chat = gr.ChatInterface(chat_interface, css=css, description="Lily", clear_btn=None, undo_btn=None, retry_btn=None)
    
    # Add a button to save chat history
    save_button = gr.Button("Save History")
    save_button.click(fn=save_chat_history_to_salesforce)

# Launch the Gradio interface
demo.launch()