import gradio as gr
import spaces
import os
from pathlib import Path
from llama_index.core.tools import QueryEngineTool
from llama_index.core import VectorStoreIndex
from llama_index.core import Settings
from llama_index.core import SimpleDirectoryReader
from llama_index.llms.groq import Groq
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from typing import Tuple
from llama_index.core import StorageContext, load_index_from_storage
from llama_index.core.objects import ObjectIndex
from llama_index.core.agent import ReActAgent
import torch
import sys
import io

@spaces.GPU
def create_doc_tools(document_fp: str, doc_name: str, verbose: bool = True) -> Tuple[QueryEngineTool,]:
    documents = SimpleDirectoryReader(input_files=[document_fp]).load_data()

    llm = Groq(model="mixtral-8x7b-32768")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-mpnet-base-v2", device=device)
    Settings.llm = llm
    Settings.embed_model = embed_model

    load_dir_path = f"/home/user/app/agentic_index_st/{doc_name}"
    storage_context = StorageContext.from_defaults(persist_dir=load_dir_path)
    vector_index = load_index_from_storage(storage_context)
    vector_query_engine = vector_index.as_query_engine()

    vector_tool = QueryEngineTool.from_defaults(
        name=f"{doc_name}_vector_query_engine_tool",
        query_engine=vector_query_engine,
        description=f"Useful for retrieving specific context from the {doc_name}.",
    )

    return vector_tool

def find_tex_files(directory: str):
    tex_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(('.tex', '.txt')):
                file_path = os.path.abspath(os.path.join(root, file))
                tex_files.append(file_path)
    tex_files.sort()
    return tex_files

def initialize_agent(apikey):
    os.environ["GROQ_API_KEY"] = apikey

    llm = Groq(model="mixtral-8x7b-32768")

    try:
        directory = '/home/user/app/rag_docs_final_review_tex_merged'
        tex_files = find_tex_files(directory)
        
        paper_to_tools_dict = {}
        for paper in tex_files:
            path = Path(paper)
            vector_tool = create_doc_tools(doc_name=path.stem, document_fp=path)
            paper_to_tools_dict[path.stem] = [vector_tool]
    
        initial_tools = [t for paper in tex_files for t in paper_to_tools_dict[Path(paper).stem]]
    
        obj_index = ObjectIndex.from_objects(
            initial_tools,
            index_cls=VectorStoreIndex,
        )
    
        obj_retriever = obj_index.as_retriever(similarity_top_k=6)
    
        context = """You are an agent designed to answer scientific queries over a set of given documents.
        Please always use the tools provided to answer a question. Do not rely on prior knowledge.
        """
    
        agent = ReActAgent.from_tools(
            tool_retriever=obj_retriever,
            llm=llm,
            verbose=True,
            context=context
        )
        return agent
    except Exception as e:
        return f"Error: {str(e)}"

def chat_with_agent(prompt, agent, verbose_toggle):
    try:
        # Check if agent initialization was successful
        if isinstance(agent, str) and agent.startswith("Error:"):
            return agent  # Return the error message if initialization failed

        # Redirect stdout
        original_stdout = sys.stdout
        sys.stdout = io.StringIO()

        # query the agent
        response = agent.query(prompt)

        # Get the captured output and restore stdout
        output = sys.stdout.getvalue()
        sys.stdout = original_stdout

        # format the received verbose output
        verbose = ''
        for output_string in output.split('==='):
            verbose += output_string
            verbose += '\n'

        # assistant response
        msg = f'{verbose}' if verbose_toggle else f'{response.response[10:]}'
        
        return msg
    except Exception as e:
        return str(e)

def main():
    agent = None

    def set_apikey(apikey):
        nonlocal agent
        agent = initialize_agent(apikey)
        if isinstance(agent, str) and agent.startswith("Error:"):
            return agent  # Return the error message if initialization failed
        return "API Key Set. You may start asking questions now."

    def reset_chat():
        nonlocal agent
        agent = None
        return "Chat reset. You may start asking questions now."

    def chat_function(prompt, apikey, verbose_toggle):
        nonlocal agent
        if not agent:
            api_response = set_apikey(apikey)
            if isinstance(agent, str) and agent.startswith("Error:"):
                return api_response  # Return the error message if initialization failed
        return chat_with_agent(prompt, agent, verbose_toggle)

    with gr.Blocks() as demo:
        gr.Markdown("# AMGPT, Powered by LlamaIndex")

        with gr.Row():
            apikey = gr.Textbox(label="Enter your Groq API Key", type="password")
            set_apikey_button = gr.Button("Set API Key")
        
        set_apikey_button.click(set_apikey, inputs=apikey, outputs=None)

        with gr.Row():
            verbose_toggle = gr.Checkbox(label="Verbose", value=True)
            reset = gr.Button("Reset Chat")
        
        reset.click(reset_chat, outputs=None)

        prompt = gr.Textbox(label="Ask a question")
        output = gr.Textbox(label="Response")

        prompt.submit(chat_function, inputs=[prompt, apikey, verbose_toggle], outputs=output)

    demo.launch()

if __name__ == "__main__":
    main()