import gradio as gr
import os
from http.cookies import SimpleCookie
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
import datetime
from gradio_client import Client
import requests

# Load environment variables
load_dotenv()

# Configure the Llama index settings with updated API
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)

# Function to save chat history to cookies
def save_chat_history_to_cookies(chat_id, query, response, cookies):
    if cookies is None:
        cookies = {}
    history = cookies.get('chat_history', '[]')
    history_list = eval(history)
    history_list.append({
        "chat_id": chat_id,
        "query": query,
        "response": response,
        "timestamp": str(datetime.datetime.now())
    })
    cookies['chat_history'] = str(history_list)

def handle_query(query, cookies=None):
    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 = ""
    if cookies:
        history = cookies.get('chat_history', '[]')
        history_list = eval(history)
        for entry in reversed(history_list):
            if entry["query"].strip():
                context_str += f"User asked: '{entry['query']}'\nBot answered: '{entry['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())
    save_chat_history_to_cookies(chat_id, query, response, cookies)

    return response

# Define the button click function
def retrieve_history_and_redirect(cookies):
    # Initialize the Gradio client
    client = Client("vilarin/Llama-3.1-8B-Instruct")

    # Retrieve and format chat history
    history = cookies.get('chat_history', '[]')
    history_list = eval(history)
    history_str = "\n".join(
        [f"User: {entry['query']}\nBot: {entry['response']}" for entry in history_list]
    )

    # Prepare the message for summarization
    message = f"""
    Chat history:
    {history_str}
    """

    # Call the Gradio API for summarization
    result = client.predict(
        message=message,
        system_prompt="Summarize the text and provide client interest in 30-40 words in bullet points.",
        temperature=0.8,
        max_new_tokens=1024,
        top_p=1,
        top_k=20,
        penalty=1.2,
        api_name="/chat"
    )

    # Print the result for debugging
    print(f"Summary: {result}")
    
    # Prepare the data to be sent via POST request
    data = {
        "summary": result,
        "timestamp": str(datetime.datetime.now())
    }

    # Send the result to the URL
    response = requests.post("https://redfernstech.com/api/receive_result", json=data)

    # Debugging response status
    print(f"POST request response: {response.status_code}, {response.text}")

    return response.status_code, response.text

# Define your Gradio chat interface function
def chat_interface(message, history):
    cookies = {}  # You might need to get cookies from the request in a real implementation
    try:
        # Process the user message and generate a response
        response = handle_query(message, cookies)

        # 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.svelte-1rjryqp{display: none;}
div.progress-text.svelte-z7cif2.meta-text {display: none;}
'''

# Use Gradio Blocks to wrap components
with gr.Blocks(css=css) as demo:
    chatbot = gr.Chatbot(placeholder="<strong>Your Personal Yes-Man</strong><br>Ask Me Anything")
    with gr.Row():  # Align horizontally in the same line
        with gr.Column(scale=8):  # 80% for the input and submit button
            submit_button = gr.ChatInterface(fn=chat_interface, chatbot=chatbot, clear_btn=None, undo_btn=None, retry_btn=None)
        with gr.Column(scale=2):  # 20% for custom button
            custom_button = gr.Button("Close Chat")

    # Connect the button with the function, and handle the redirection
    custom_button.click(fn=retrieve_history_and_redirect, inputs=[gr.State()])
    
    # Add a JavaScript function to handle redirection after the Gradio event is processed
    custom_button.click(fn=None, js="() => { window.open('https://redfernstech.com/chat-bot-test', '_blank'); }")

# Launch the Gradio interface
demo.launch()