import sys
import os
import argparse
from gradio_base import WebUI, UIHelper, PORT, HOST, Client
from gradio_config import GradioConfig as gc
from typing import List, Tuple, Any
import gradio as gr
import time
from Agent import Agent
from design_states import get_desgin_states,get_cot_result
from gen_utils import *
from utils import get_embedding,cos_sim
import torch
import json
import openai

def get_embedding(sentence,api_key):
    openai.api_key = api_key
    embedding_model = openai.Embedding
    embed = embedding_model.create(
        model="text-embedding-ada-002",
        input=sentence
    )
    embed = embed["data"][0]["embedding"]
    embed = torch.tensor(embed,dtype=torch.float32)
    if len(embed.shape)==1:
        embed = embed.unsqueeze(0)
    return embed

class GeneralUI(WebUI):
    def render_and_register_ui(self):
        # bind the agent with avatar
        self.agent_name:list = [self.cache["agents_name"]] if isinstance(self.cache["agents_name"], str) else self.cache['agents_name']
        gc.add_agent(self.agent_name)
    
    def handle_message(self, history, state, agent_name, token, node_name):
        if state % 10 == 0:
            self.data_history.append({agent_name: token})
        elif state % 10 == 1:
            # Same state. Need to add new bubble in same bubble.
            self.data_history[-1][agent_name] += token
        elif state % 10 == 2:
            # New state. Need to add new bubble.
            history.append([None, ""])
            self.data_history.clear()
            self.data_history.append({agent_name: token})
        else:
            assert False, "Invalid state."
        render_data = self.render_bubble(history, self.data_history, node_name, render_node_name= True)
        return render_data
    
    def __init__(
        self,
        client_cmd: list,
        socket_host: str = HOST,
        socket_port: int = PORT,
        bufsize: int = 1024,
        ui_name: str = "GeneralUI"
    ):
        super(GeneralUI, self).__init__(client_cmd, socket_host, socket_port, bufsize, ui_name)
        self.first_recieve_from_client()
        self.current_node_name = ""
        self.data_history = None
        for _ in ['agents_name', 'api_key']:
            assert _ in self.cache
    
    def generate_sop(self,api_key,proxy,target):
        os.environ["API_KEY"] = api_key
        # os.environ["PROXY"] = proxy
        self.design_assistant = "An assistant that can help users create content such as articles, blogs, advertising copy, etc"
        self.tutor = "A tutor who provides personalized learning resources for students to help them understand complex concepts and problems"
        self.online_medical_consultant = "An online medical consultant who offers preliminary medical advice to patients and answers common questions about diseases, symptoms, and treatments."
        self.online_legal_consultant = "An online legal advisor who can respond to inquiries related to legal matters, providing basic legal information and advice."
        self.online_financial_advisor = "An online financial advisor who can analyze financial markets and data, offering investment advice and market forecasts to users."
        self.virtual_tour_guide = "A virtual tour guide providing destination information, travel recommendations, and virtual travel experiences for travelers."
        self.design_assistant = get_embedding(self.design_assistant,api_key)
        self.tutor = get_embedding(self.tutor,api_key)
        self.online_medical_consultant = get_embedding(self.online_medical_consultant,api_key)
        self.online_legal_consultant = get_embedding(self.online_legal_consultant,api_key)
        self.online_financial_advisor = get_embedding(self.online_financial_advisor,api_key)
        self.virtual_tour_guide = get_embedding(self.virtual_tour_guide,api_key)
        self.embeddings = torch.cat([self.design_assistant,self.tutor,self.online_medical_consultant,self.online_legal_consultant,self.online_financial_advisor,self.virtual_tour_guide],dim = 0)  
        self.SOP["config"]["API_KEY"] = api_key
        # self.SOP["config"]["PROXY"] = proxy
        target_tensor = get_embedding(target,api_key)
        sim_scores = cos_sim(target_tensor, self.embeddings)[0]
        top_k_score, top_k_idx = torch.topk(sim_scores,k = 1)
        if top_k_score > 0.7:
            index = top_k_idx
        else:
            index = 0
        target = get_cot_result(target)
        design_states = get_desgin_states(target,index)
        root = design_states[0]["state_name"]
        agents = get_agents(design_states)
        relations = get_relations(design_states)
        states = gen_states(design_states)
        for state_name,state_dict in states.items():
            state_dict["begin_role"] = list(agents.keys())[0]
            state_dict["begin_query"] = "Now that we are in the **{}**, I'm glad to offer you assistance.".format(state_name)
        self.SOP["root"] = root
        self.SOP["relations"] = relations
        self.SOP["agents"] = agents
        self.SOP["states"] = states
        # 将字典写入JSON文件
        print(self.SOP)
        file_name = 'generated_sop.json'
        with open(file_name, "w",encoding="utf-8") as json_file:
            json.dump(self.SOP, json_file ,indent=4,ensure_ascii=False)
        return file_name

    def load_sop_fn(self,sop):
        return sop.name
    
    def construct_ui(self):
        with gr.Blocks(css=gc.CSS) as demo:
            with gr.Tab(label="SOP generation") as tab1:
                self.SOP = {
                    "config": {
                        "API_KEY": "sk-********",
                        "MAX_CHAT_HISTORY": "5",
                        "User_Names": '["User"]',
                    },
                    "root": "state1",
                    "relations": {
                        "state1": {"0": "state1", "1": "state2"},
                        "state2": {"0": "state2", "1": "end_state"},
                    },
                    "agents": None,
                    "states": None,
                }
                gr.Markdown("""# Generate Agent""")
                with gr.Row():
                    self.api_key_sop_generation = gr.Textbox(label="api_key")
                    self.proxy_sop_generation = gr.Textbox(label="proxy",visible=False)
                with gr.Row():
                    self.requirement_sop_generation = gr.Textbox(value ="a shopping assistant help customer to buy the commodity",label="requirement")
                with gr.Row():
                    self.generated_sop = gr.File(label="generated_file")
                self.generate_button = gr.Button(label="Generate")
                self.generate_button.click(fn = self.generate_sop,inputs=[self.api_key_sop_generation,self.proxy_sop_generation,self.requirement_sop_generation],outputs=[self.generated_sop])
            with gr.Tab(label="Chat") as tab2:
                uploaded_sop = gr.State()
                with gr.Row():
                    sop = gr.File(label="upload your custmized SOP")
                    load_sop_btn = gr.Button(value="Load SOP")
                load_sop_btn.click(self.load_sop_fn, sop,uploaded_sop)
                with gr.Column():
                    self.radio_mode = gr.Radio(
                        [Client.SINGLE_MODE],
                        label = Client.MODE_LABEL,
                        info = Client.MODE_INFO,
                        value= Client.SINGLE_MODE,
                        interactive=True
                        # label="Select the execution mode",
                        # info="Single mode refers to when the current agent output ends, it will stop running until you click to continue. Auto mode refers to when you complete the input, all agents will continue to output until the task ends."
                    )
                    self.text_api = gr.Textbox(
                        value = self.cache["api_key"],
                        placeholder="openai key",
                        label="Please input valid openai key for gpt-3.5-turbo-16k."
                    )
                    self.btn_start = gr.Button(
                        value="Start😁(Click here to start!)",
                    )
                    self.chatbot = gr.Chatbot(
                        elem_id="chatbot1",
                        label="Dialog",
                        visible=False,
                        height=700
                    )
                    self.btn_next = gr.Button(
                        value="Next Agent Start",
                        visible=False
                    )
                    with gr.Row():
                        self.text_input = gr.Textbox(
                            placeholder="Please enter your content.",
                            label="Input",
                            scale=9,
                            visible=False
                        )
                        self.btn_send = gr.Button(
                            value="Send",
                            visible=False
                        )
                    self.btn_reset = gr.Button(
                        value="Restart",
                        visible=False
                    )

                all_components = [self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
                
                self.btn_start.click(
                    fn = self.btn_start_when_click,
                    inputs=[self.radio_mode, self.text_api,uploaded_sop],
                    outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next, self.radio_mode, self.text_api]
                ).then(
                    fn = self.btn_start_after_click,
                    inputs=[self.chatbot],
                    outputs=all_components
                )
                
                self.btn_send.click(
                    fn=self.btn_send_when_click,
                    inputs=[self.text_input, self.chatbot],
                    outputs=all_components
                ).then(
                    fn=self.btn_send_after_click,
                    inputs=[self.text_input, self.chatbot],
                    outputs=all_components
                )
                
                self.text_input.submit(
                    fn=self.btn_send_when_click,
                    inputs=[self.text_input, self.chatbot],
                    outputs=all_components
                ).then(
                    fn=self.btn_send_after_click,
                    inputs=[self.text_input, self.chatbot],
                    outputs=all_components
                )
                
                self.btn_reset.click(
                    fn=self.btn_reset_when_click,
                    inputs=[],
                    outputs=all_components
                ).then(
                    fn=self.btn_reset_after_click,
                    inputs=[],
                    outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next, self.radio_mode, self.text_api]
                )
                
                self.btn_next.click(
                    fn=self.btn_next_when_click,
                    inputs=[self.chatbot],
                    outputs=all_components
                ).then(
                    fn=self.btn_next_after_click,
                    inputs=[self.chatbot],
                    outputs=all_components
                )
            
            self.demo = demo
    
    def btn_start_when_click(self, mode, api,sop):
        """
        inputs=[mode, api]
        outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next, self.radio_mode]
        """
        print("server: send ", mode, api)
        self.send_start_cmd({"mode": mode, "api_key":api,"uploaded_sop": sop})
        agents,roles_to_names,names_to_roles = Agent.from_config(str(sop))
        agents_name = []
        for i in names_to_roles :
            for j in names_to_roles[i]:
                agents_name.append(j+"("+names_to_roles[i][j]+")")
        self.new_render_and_register_ui(agents_name)
        return gr.Button.update(visible=False), \
            gr.Button.update(visible=False),\
            gr.Button.update(visible=False),\
            gr.Chatbot.update(visible=True),\
            gr.Textbox.update(visible=False),\
            gr.Button.update(visible=False),\
            gr.Radio.update(visible=False),\
            gr.Textbox.update(visible=False)
    
    def new_render_and_register_ui(self,agent_names):
        gc.add_agent(agent_names, 0)

    def btn_start_after_click(self, history):
        """
        inputs=[self.chatbot]
        outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
        """
        if self.data_history is None:
            self.data_history = list()
        receive_server = self.receive_server
        while True:
            data_list: List = receive_server.send(None)
            for item in data_list:
                data = eval(item)
                assert isinstance(data, list)
                state, agent_name, token, node_name = data
                self.current_node_name = node_name
                assert isinstance(state, int)
                assert state in [10, 11, 12, 30, 99, 98]
                if state == 99:
                    # finish
                    yield gr.Button.update(visible=False),\
                        gr.Button.update(visible=True, interactive=False),\
                        gr.Button.update(visible=True, interactive=True),\
                        history,\
                        gr.Textbox.update(visible=True, interactive=False),\
                        gr.Button.update(visible=False)
                    return
                elif state == 98:
                    # single mode
                    yield gr.Button.update(visible=False), \
                            gr.Button.update(visible=False),\
                            gr.Button.update(visible=True),\
                            history,\
                            gr.Textbox.update(visible=False),\
                            gr.Button.update(visible=True, value=f"Next Agent: 🤖{agent_name} | Next Node: ⭕{node_name}")
                    return 
                elif state == 30:
                    # user input
                    yield gr.Button.update(visible=False), \
                            gr.Button.update(visible=True),\
                            gr.Button.update(visible=True),\
                            history,\
                            gr.Textbox.update(visible=True, value=""),\
                            gr.Button.update(visible=False)
                    return
                history = self.handle_message(history, state, agent_name, token, node_name)
                yield gr.Button.update(visible=False), \
                    gr.Button.update(visible=False),\
                    gr.Button.update(visible=False),\
                    history,\
                    gr.Textbox.update(visible=False),\
                    gr.Button.update(visible=False)
    
    def btn_send_when_click(self, text_input, history):
        '''
        inputs=[self.text_input, self.chatbot]
        outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
        '''
        history = self.handle_message(history, 10, 'User', text_input, self.current_node_name)
        self.send_message("<USER>"+text_input+self.SIGN["SPLIT"])
        yield gr.Button.update(visible=False), \
            gr.Button.update(visible=False),\
            gr.Button.update(visible=False),\
            history,\
            gr.Textbox.update(visible=False),\
            gr.Button.update(visible=False)
        return
    
    def btn_send_after_click(self, text_input, history):
        '''
        inputs=[self.text_input, self.chatbot]
        outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
        '''
        yield from self.btn_start_after_click(history=history)
        return 
    
    def btn_reset_when_click(self):
        """
        outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
        """
        return gr.Button.update(interactive=False), gr.Button.update(interactive=False), gr.Button.update(interactive=False, value="Restarting....."), gr.Chatbot.update(label="Dialog"), \
            gr.Textbox.update(interactive=False), gr.Button.update(visible=False)
    
    def btn_reset_after_click(self):
        """
        outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next, self.radio_mode]
        """
        self.reset()
        self.first_recieve_from_client(reset_mode=True)
        self.current_node_name = ""
        self.data_history = None
        return gr.Button.update(interactive=True, visible=True), \
            gr.Button.update(interactive=True, visible=False), \
            gr.Button.update(interactive=True, value="Restart", visible=False), \
            gr.Chatbot.update(label="Dialog", visible=False, value=None), \
            gr.Textbox.update(interactive=True, visible=False),\
            gr.Button.update(visible=False),\
            gr.Radio.update(visible=True), \
            gr.Textbox.update(visible=True)
    
    def btn_next_when_click(self, history):
        """
        outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
        """
        yield gr.Button.update(visible=False), \
                gr.Button.update(visible=False),\
                gr.Button.update(visible=False),\
                history,\
                gr.Textbox.update(visible=False),\
                gr.Button.update(visible=False)
        self.send_message("nothing")
        return
    
    def btn_next_after_click(self, history):
        time.sleep(1)
        yield from self.btn_start_after_click(history=history)
        return
        
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='A demo of chatbot')
    parser.add_argument('--agent', type=str, help='path to SOP json')
    args = parser.parse_args()
    
    ui = GeneralUI(client_cmd=["python3","gradio_backend.py"])
    ui.construct_ui()
    ui.run(share=True)