import glob
import gradio as gr
import openai
import os
from dotenv import load_dotenv
import phoenix as px
import llama_index
from llama_index import Prompt, ServiceContext, VectorStoreIndex, SimpleDirectoryReader
from llama_index.chat_engine.types import ChatMode
from llama_index.llms import ChatMessage, MessageRole
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.text_splitter import SentenceSplitter
from llama_index.extractors import TitleExtractor
from llama_index.ingestion import IngestionPipeline
from chat_template import CHAT_TEXT_QA_PROMPT, TEXT_QA_SYSTEM_PROMPT
from schemas import ChatbotVersion, ServiceProvider
from chatbot import Chatbot, IndexBuilder
from custom_io import MarkdownReader, UnstructuredReader, default_file_metadata_func
from qdrant import client as qdrantClient
from llama_index import set_global_service_context

from service_provider_config import get_service_provider_config

load_dotenv()
# initial service setup
px.launch_app()
llama_index.set_global_handler("arize_phoenix")
# llama_index.set_global_handler("wandb", run_args={"project": "llamaindex"})
openai.api_key = os.getenv("OPENAI_API_KEY")

IS_LOAD_FROM_VECTOR_STORE = True 
VDB_COLLECTION_NAME = "demo-v4"
MODEL_NAME = ChatbotVersion.CHATGPT_4.value


CHUNK_SIZE = 8191
LLM, EMBED_MODEL = get_service_provider_config(
    service_provider=ServiceProvider.OPENAI, model_name=MODEL_NAME)
service_context = ServiceContext.from_defaults(
    chunk_size=CHUNK_SIZE,
    llm=LLM,
    embed_model=EMBED_MODEL,
)
set_global_service_context(service_context)


class AwesumIndexBuilder(IndexBuilder):
    def _load_doucments(self):
        directory = "./awesumcare_data/awesumcare_manual_data"
        # all_files = glob.glob(os.path.join(directory, '*.md'))
        # faq_files = [f for f in all_files if 'FAQ' in os.path.basename(f)]
        # print(faq_files)
        dir_reader = SimpleDirectoryReader(directory, file_extractor={
            ".pdf": UnstructuredReader(),
            ".docx": UnstructuredReader(),
            ".pptx": UnstructuredReader(),
            ".md": MarkdownReader()
        },
            recursive=True,
            # input_files=faq_files,
            exclude=["*.png", "*.pptx", "*.docx", "*.pdf"],
            file_metadata=default_file_metadata_func)

        self.documents = dir_reader.load_data()
        print(f"Loaded {len(self.documents)} docs")

    def _setup_service_context(self):
        super()._setup_service_context()

    def _setup_vector_store(self):
        self.vector_store = QdrantVectorStore(
            client=qdrantClient, collection_name=self.vdb_collection_name)
        super()._setup_vector_store()

    def _setup_index(self):
        super()._setup_index()
        if self.is_load_from_vector_store:
            self.index = VectorStoreIndex.from_vector_store(self.vector_store)
            print("set up index from vector store")
            return
        pipeline = IngestionPipeline(
            transformations=[
                # SentenceSplitter(),
                self.embed_model,
            ],
            vector_store=self.vector_store,
        )
        pipeline.run(documents=self.documents, show_progress=True)
        self.index = VectorStoreIndex.from_vector_store(self.vector_store)


class AwesumCareToolChatbot(Chatbot):
    DENIED_ANSWER_PROMPT = ""
    SYSTEM_PROMPT = ""
    CHAT_EXAMPLES = [
        "什麼是安心三寶？",
        "點樣立平安紙？",
        "甚麼是⾒證？",
        "訂立每份⽂件需要多少錢以及付款⽅法？",
        "通過安⼼三寶製作的⽂件有法律效⼒嗎？",

    ]

    def _setup_observer(self):
        pass

    def _setup_index(self):
        super()._setup_index()

    # def _setup_index(self):
    #     self.index = VectorStoreIndex.from_documents(
    #         self.documents,
    #         service_context=self.service_context
    #     )
    #     super()._setup_index()

    def _setup_query_engine(self):
        super()._setup_query_engine()
        self.query_engine = self.index.as_query_engine(
            text_qa_template=CHAT_TEXT_QA_PROMPT)

    def _setup_tools(self):
        from llama_index.tools.query_engine import QueryEngineTool
        self.tools = QueryEngineTool.from_defaults(
            query_engine=self.query_engine)
        return super()._setup_tools()

    def _setup_chat_engine(self):
        # testing #
        from llama_index.agent import OpenAIAgent
        self.chat_engine = OpenAIAgent.from_tools(
            tools=[self.tools],
            llm=LLM,
            similarity_top_k=1,
            verbose=True
        )
        print("set up agent as chat engine")
        # testing #
        # self.chat_engine = self.index.as_chat_engine(
        #     chat_mode=ChatMode.BEST,
        #     similarity_top_k=5,
        #     text_qa_template=CHAT_TEXT_QA_PROMPT)
        super()._setup_chat_engine()

class AweSumCareContextChatbot(AwesumCareToolChatbot):
    def _setup_query_engine(self):
        pass
    def _setup_tools(self):
        pass
    def _setup_chat_engine(self):
        self.chat_engine = self.index.as_chat_engine(
            chat_mode=ChatMode.CONTEXT,
            similarity_top_k=5,
            system_prompt=TEXT_QA_SYSTEM_PROMPT.content,
            text_qa_template=CHAT_TEXT_QA_PROMPT)

class AweSumCareSimpleChatbot(AwesumCareToolChatbot):
    def _setup_query_engine(self):
        pass
    def _setup_tools(self):
        pass
    def _setup_chat_engine(self):
        self.chat_engine = self.index.as_chat_engine(
            chat_mode=ChatMode.SIMPLE)

model_name = MODEL_NAME
index_builder = AwesumIndexBuilder(vdb_collection_name=VDB_COLLECTION_NAME,
                                   embed_model=EMBED_MODEL,
                                   is_load_from_vector_store=IS_LOAD_FROM_VECTOR_STORE)

# gpt-3.5-turbo-1106, gpt-4-1106-preview
awesum_chatbot = AwesumCareToolChatbot(model_name=model_name, index_builder=index_builder)
awesum_chatbot_context = AweSumCareContextChatbot(model_name=model_name, index_builder=index_builder)
awesum_chatbot_simple = AweSumCareSimpleChatbot(model_name=model_name, index_builder=index_builder)


def service_setup(model_name):
    CHUNK_SIZE = 1024
    LLM, EMBED_MODEL = get_service_provider_config(
        service_provider=ServiceProvider.OPENAI, model_name=model_name)
    service_context = ServiceContext.from_defaults(
        chunk_size=CHUNK_SIZE,
        llm=LLM,
        embed_model=EMBED_MODEL,
    )
    set_global_service_context(service_context)
    return LLM, EMBED_MODEL


def vote(data: gr.LikeData):
    if data.liked:
        gr.Info("You up-voted this response: " + data.value)
    else:
        gr.Info("You down-voted this response: " + data.value)


chatbot = gr.Chatbot()

with gr.Blocks() as demo:

    gr.Markdown("# Awesum Care demo")



    # with gr.Row():
    #     model_selector = gr.Radio(
    #         value=ChatbotVersion.CHATGPT_35.value,
    #         choices=[ChatbotVersion.CHATGPT_35.value, ChatbotVersion.CHATGPT_4.value],
    #         label="Select Chatbot Model (To be implemented)"
    #     )

    with gr.Tab("With relevant context sent to system prompt"):
        context_interface = gr.ChatInterface(
            awesum_chatbot_context.stream_chat,
            examples=awesum_chatbot.CHAT_EXAMPLES,
        )
        chatbot.like(vote, None, None)

    # with gr.Tab("With function calling as tool to retrieve"):
    #     function_call_interface = gr.ChatInterface(
    #         awesum_chatbot.stream_chat,
    #         examples=awesum_chatbot.CHAT_EXAMPLES,
    #     )
    #     chatbot.like(vote, None, None)


    # with gr.Tab("Vanilla ChatGPT without modification"):
    #     vanilla_interface = gr.ChatInterface(
    #         awesum_chatbot_simple.stream_chat,
    #         examples=awesum_chatbot.CHAT_EXAMPLES)

    gr.Markdown("instructions:\n"
                "\nUsing model gpt-4-preview-1106, the most advanced model now in the market.\n"
                "\n(Note that it can be much slower than gpt-3.5, openai's api can be unstable sometimes.)\n"
                # "\nThree Tabs:\n"
                # "1. Relevant context: retreiving relevant documents and send to ChatGPT.\n"
                # "2. Give tools to chatgpt to retrieve context: the most advanced, slowest (>30s to use the tools, before answering).\n"
                # "3. Vanilla ChatGPT: self-explanatory.\n"
                )
    # @model_selector.change(inputs=[model_selector, chatbot], outputs=[context_interface, function_call_interface, vanilla_interface])
    # def switch_model(model_name, my_chatbot):
    #     print(model_name)
    #     print(my_chatbot.config())
    #     LLM, EMBED_MODEL = service_setup(model_name)
    #     # global awesum_chatbot, awesum_chatbot_context, awesum_chatbot_simple
    #     # Logic to switch models - create new instances of the chatbots with the selected model
    #     index_builder = AwesumIndexBuilder(vdb_collection_name=VDB_COLLECTION_NAME,
    #                                     embed_model=EMBED_MODEL,
    #                                     is_load_from_vector_store=IS_LOAD_FROM_VECTOR_STORE)
    #     awesum_chatbot = AwesumCareToolChatbot(model_name=model_name, index_builder=index_builder, llm=LLM)
    #     awesum_chatbot_context = AweSumCareContextChatbot(model_name=model_name, index_builder=index_builder)
    #     awesum_chatbot_simple = AweSumCareSimpleChatbot(model_name=model_name, index_builder=index_builder)
    #     # return awesum_chatbot.stream_chat, awesum_chatbot_context.stream_chat, awesum_chatbot_simple.stream_chat

    #     new_context_interface = gr.ChatInterface(
    #         awesum_chatbot_context.stream_chat,
    #     )
    #     new_function_call_interface = gr.ChatInterface(
    #         awesum_chatbot.stream_chat,
    #     )
    #     new_vanilla_interface = gr.ChatInterface(
    #         awesum_chatbot_simple.stream_chat,
    #     )
    #     return new_context_interface, new_function_call_interface, new_vanilla_interface


demo.queue()
demo.launch(share=False, auth=("demo", os.getenv("PASSWORD")))