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
import uuid
import json

# 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
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):
        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
    current_chat_history.append((query, response))

    return response

def save_chat_history(history):
    # Save the chat history to a local file or Firebase
    session_id = str(uuid.uuid4())
    chat_history_path = f"chat_history_{session_id}.json"
    
    with open(chat_history_path, 'w') as f:
        json.dump(history, f)
    print(f"Chat history saved as {chat_history_path}")

    # Save to Salesforce
    save_to_salesforce(current_chat_history)

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

    session_id, sf_instance = SalesforceLogin(username=username, password=password, security_token=security_token, domain=domain)
    sf = Salesforce(instance=sf_instance, session_id=session_id)
    for past_query, response in history:
        data = {
            'Name': 'Chat with user',
            'Bot_Message__c': response,
            'User_Message__c': past_query,
            'Date__c': str(datetime.datetime.now().date())
        }
        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>'
    
    # Save the updated history
    save_chat_history(current_chat_history)
    
    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)

        # Update the history and save it
        save_chat_history(current_chat_history)

        # 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;}
'''

demo = gr.ChatInterface(chat_interface, 
                        css=css,
                        description="Lily",
                        clear_btn=None, undo_btn=None, retry_btn=None,
)

# Add a button to save chat history
gr.Button("Close Chat").click(fn=save_chat_history)

# Launch the interface
demo.launch()