import os
import gradio as gr
from typing import List, Tuple
from gradio import ChatMessage
import base64
from llama_index.core import StorageContext, load_index_from_storage
from dotenv import load_dotenv
from retrieve import get_latest_dir, get_latest_html_file
from graph_handler import query_graph_qa, plot_subgraph
from embed_handler import query_rag_qa
from evaluate import evaluate_llm, reasoning_graph, get_coupon
import base64
import time

load_dotenv()

KG_INDEX_PATH = get_latest_dir(os.getenv("GRAPH_DIR"))
KG_PLOT_PATH = get_latest_html_file(os.getenv("GRAPH_VIS"))
RAG_INDEX_PATH = get_latest_dir(os.getenv("EMBEDDING_DIR"))

# Load Graph-RAG index
graph_rag_index = load_index_from_storage(
    StorageContext.from_defaults(persist_dir=KG_INDEX_PATH)
)

# Load RAG index
rag_index = load_index_from_storage(
    StorageContext.from_defaults(persist_dir=RAG_INDEX_PATH)
)


def query_tqa(query, search_level):
    """
    Query the Graph-RAG and RAG models for a given query.

    Args:
    query (str): The query to ask the RAGs.
    search_level (int): The max search level to use for the Graph RAG.

    Returns:
    tuple: The response, reference, and reference text for the Graph-RAG and RAG models.
    """

    if not query.strip():
        raise gr.Error("Please enter a query before asking.")

    grag_response, grag_reference, grag_reference_text = query_graph_qa(
        graph_rag_index, query, search_level
    )
    # rag_response, rag_reference, rag_reference_text = query_rag_qa(
    #     rag_index, query, search_level
    # )
    print(str(grag_response.response))
    return (
        str(grag_response.response)
    )



 
# with gr.Blocks() as demo:
#     gr.Markdown("# Comfy Virtual Assistant")
#     chatbot = gr.Chatbot(
#         label="Comfy Virtual Assistant",
#         type="messages",
#         scale=1,
#         # suggestions = [
#         #         {"text": "How much iphone cost?"},
#         #         {"text": "What phone options do i have ?"}
#         #         ],
           
#     )
#     msg = gr.Textbox(label="Input Your Query")
#     clear = gr.ClearButton([msg, chatbot])

#     def respond(message, chat_history):
#             bot_message = query_tqa(message, 2)
#             # chat_history.append((message, bot_message))
#             chat_history.append(ChatMessage(role="user", content=message))
#             chat_history.append(ChatMessage(role="assistant", content=bot_message))
#             time.sleep(1)
#             return "", chat_history
    
#     msg.submit(respond, [msg, chatbot], [msg, chatbot])

def chatbot_response(message: str, history: List[Tuple[str, str]]) -> str:
    # Use the query_tqa function to get the response
    search_level = 2  # You can adjust this or make it configurable
    response = query_tqa(message, search_level)
    return response

# Create the Gradio interface
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history):
        bot_message = chatbot_response(history[-1][0], history)
        history[-1][1] = bot_message
        return history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

    
demo.launch(auth=(os.getenv("ID"), os.getenv("PASS")), share=False)