from dotenv import load_dotenv
import gradio as gr
import os
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
import firebase_admin
from firebase_admin import db, credentials
import datetime
import uuid
import random

# Load environment variables
load_dotenv()

# Initialize Firebase with provided credentials and URL
cred = credentials.Certificate("redfernstech-fd8fe-firebase-adminsdk-g9vcn-0537b4efd6.json")
firebase_admin.initialize_app(cred, {"databaseURL": "https://redfernstech-fd8fe-default-rtdb.firebaseio.com/"})

# Configure 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 directories
PERSIST_DIR = "db"
PDF_DIRECTORY = 'data'

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

# Dictionary to store chat histories for different sessions
session_chat_histories = {}

def select_random_name():
    names = ['Clara', 'Lily']
    return random.choice(names)

def data_ingestion_from_directory():
    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(session_id, query):
    chat_text_qa_msgs = [
        (
            "user",
            """
            As Clara, your goal is to provide code to the user.
            Your task is to give code to the model and offer guidance on creating a website using Django, HTML, CSS, and Bootstrap.
            {context_str}
            Question:
            {query_str}
            """
        )
    ]
    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)

    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    index = load_index_from_storage(storage_context)

    context_str = ""
    if session_id in session_chat_histories:
        for past_query, response in reversed(session_chat_histories[session_id]):
            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."

    if session_id not in session_chat_histories:
        session_chat_histories[session_id] = []

    session_chat_histories[session_id].append((query, response))
    message_data = {
        "query": query,
        "response": response,
        "timestamp": datetime.datetime.now().isoformat()
    }

    save_chat_message(session_id, message_data)
    return response

def save_chat_message(session_id, message_data):
    ref = db.reference(f'/chat_history/{session_id}')
    ref.push().set(message_data)

def chat_interface(message, history):
    # Retrieve or create a new session ID based on history
    session_id = history[0][1] if history and history[0][1] else str(uuid.uuid4())
    history.append((message, session_id))  # Append the session ID to history
    response = handle_query(session_id, message)
    return response, history

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

# Load data and start Gradio interface
print("Processing PDF ingestion from directory:", PDF_DIRECTORY)
data_ingestion_from_directory()

gr.ChatInterface(fn=chat_interface, 
                 css=css, 
                 description="Clara", 
                 clear_btn=None, undo_btn=None, retry_btn=None).launch()