import sys  
import os  
current_directory = os.path.dirname(os.path.abspath(__file__))    
sys.path.append(current_directory)  
import agentscope
from agentscope.message import Msg
import os

from agent.flex_service_toolkit import *
from AlgoriAgent.projects.algoriAgent.agent.algori_agent import *

from agentscope.service import (
    ServiceToolkit,
    ServiceResponse,
    ServiceExecStatus,
)
import configparser

config = configparser.ConfigParser()
config.read('config.ini')
openai_api_key = config['Global']['api_key']



OPENAI_CFG_DICT = {
    "config_name": "openai_cfg",    # 此配置的名称，必须保证唯一
    "model_type": "openai_chat",         # 模型类型
    "model_name": config['Global']['model'],  # 模型名称
    #"model_name": "gpt-4",  # 模型名称
    #"model_name": "llama3",
    "api_key": openai_api_key,      # OpenAI API key. 如果没有设置，将使用环境变量中的 OPENAI_API_KEY

    "client_args": {
        "base_url": config['Global']['base_chat_url']
    },

}

import uuid
from AlgoriAgent.projects.algoriAgent.tools.judge_tools import judge
EMPTY_CHAPTER_CHAIN = [ Chapter(1, CHAPTER_FOCUS, "本章是未打开某个具体章节时的默认章节。", "处于本章节时不会有任何章节跳转。请作为一名经验丰富的算法教师，回答用户的问题。") ]

def tool_name_to_tool(tool_name_list):
    tools = []
    for tool_name in tool_name_list:
        if tool_name == "":
            continue
        if tool_name == "judge":
            tools.append(judge)
    return tools

def tool_name_to_tool_with_args(tool_name_list, tool_args_list) -> list[tuple[Callable, dict]]:
    tools = []
    for tool_name, tool_args in zip(tool_name_list, tool_args_list):
        if tool_name == "":
            continue
        if tool_name == "judge":
            tools.append((judge, {"name": tool_args[0]}))
    return tools
    
class AgentManager:
    def __init__(self,max_iter = 5, app=None, socketio=None):

        agentscope.init(model_configs=[OPENAI_CFG_DICT])#, studio_url="http://0.0.0.0:5000")
        self.agents = {}
        self.max_iter = max_iter
        self.agent_to_id = {}
        self.app = app
        self.socketio = socketio
    def new_agent(self, course_id, lesson_id, markdown:str, markdown_prompt:str, score_prompt:str, id = "Defult Assistant", root_path="."):
        '''
        markdown: 教案的markdown文件内容
        markdown_prompt: 教案的prompt的markdown文件内容
        score_prompt: 评分的prompt的markdown文件内容
        根据3个教案的prompt，生成一个agent，返回agent的id与agent
        '''
        markdown_list = markdown.split("\n")
        markdown_prompt_list = markdown_prompt.split("\n")
        score_prompt_list = score_prompt.split("\n")
        # 获取 H1 标题
        title = ""
        for line in markdown_list:
            if line.startswith("# "):
                title = line[2:]
                break
        
        # 对于 H3 标题 构造每一个 Chapter
        # 首先找到所有的 H3 标题 
        chapter_dict = {}
        chapter_sequence = []
        for line in markdown_list:
            if line.startswith("### "):
                chapter_name = line[4:]
                chapter_dict[chapter_name] = {}
                chapter_sequence.append(chapter_name)
        # 将 H3 标题 和 其对应的内容 构造成一个 Chapter

        h3_name = ""
        content = ""
        for i in range(len(markdown_list)):
            if markdown_list[i].startswith("### "): 
                if(h3_name != ""):
                    chapter_dict[h3_name]["markdown"] = content
                h3_name = markdown_list[i][4:]
                content = ""
                continue
            if h3_name != "":
                content += markdown_list[i]+"\n"
        if(h3_name != ""):
            chapter_dict[h3_name]["markdown"] = content
        

        h3_name = ""
        content = ""
        require_tools = []
        require_tools_args = []
        for i in range(len(markdown_prompt_list)):
            if markdown_prompt_list[i].startswith("### "): 
                if(h3_name != ""):
                    chapter_dict[h3_name]["markdown_prompt"] = content
                    chapter_dict[h3_name]["require_tools"] = tool_name_to_tool_with_args(require_tools, require_tools_args)
                h3_name = markdown_prompt_list[i][4:]
                content = ""
                require_tools=[]
                require_tools_args = []
                continue

            if h3_name != "":
                if markdown_prompt_list[i].startswith("_require_tools"):
                    require_tools.append(markdown_prompt_list[i].split("=")[1].strip().split(",")[0])
                    require_tools_args.append((markdown_prompt_list[i].split("=")[1].strip().split(",")[1:]))
                    continue
                content += markdown_prompt_list[i]+"\n"
        if(h3_name != ""):
            chapter_dict[h3_name]["markdown_prompt"] = content
            chapter_dict[h3_name]["require_tools"] = tool_name_to_tool_with_args(require_tools, require_tools_args)
            
        
        h3_name = ""
        content = ""
        for i in range(len(score_prompt_list)):
            if score_prompt_list[i].startswith("### "): 
                if(h3_name != ""):
                    chapter_dict[h3_name]["score_prompt"] = content
                h3_name = score_prompt_list[i][4:]
                content = ""
                continue

            if h3_name != "":
                content += score_prompt_list[i]+"\n"
        if(h3_name != ""):
            chapter_dict[h3_name]["score_prompt"] = content


        chapter_chain = []
        No = 1
        print (chapter_dict)
        for chapter_name in chapter_sequence:
            chapter_chain.append(Chapter(course_id, lesson_id, No, CHAPTER_LATTER, chapter_name, chapter_dict[chapter_name]["markdown"], chapter_dict[chapter_name]["markdown_prompt"], chapter_dict[chapter_name]["score_prompt"],chapter_dict[chapter_name]["require_tools"]))
            No+=1
        chapter_chain[0].Focus()


        # Prepare the tools for the agent
        service_toolkit = FlexServiceToolkit()
        # for tool_function in unity_function_list:
        #     service_toolkit.add(tool_function)
        agent = ChapterChainAgent(
                    name=id,
                    model_config_name="openai_cfg",
                    verbose=True,
                    service_toolkit=service_toolkit,
                    max_iters=5,
                    chapter_chain=chapter_chain,
                    root_path = root_path,
                )
        self.agents[id] = agent
        self.agent_to_id[agent] = id
        agent.agent_manager = self
        return id, agent
        
    def message_pass(self, agent, messagetype, message):
        id = self.agent_to_id[agent]
        print("message_pass", id, messagetype, message)
        with self.app.app_context(): 
            try:
                self.socketio.emit(messagetype, message, room=id, namespace='/agent')
            except Exception as e:
                print("message_pass",e)

    def system_message(self, message):
        with self.app.app_context(): 
            try:
                self.socketio.emit('system_message', message, room=id, namespace='/agent')
            except Exception as e:
                print("ERROR: system_message",e)


    def save_chapter_memory(self, agent,course_id, lesson_id, subchapter_title, mem_list, score, is_rebuttal):
        id = self.agent_to_id[agent]
        self.app.my_function.save_chapter_memory(id, course_id, lesson_id, subchapter_title, mem_list, score, is_rebuttal)


    def new_agent_with_chain(self, chapter_chain = None, id = None):
        if chapter_chain is None:
            chapter_chain = EMPTY_CHAPTER_CHAIN
        
        # Prepare the tools for the agent
        service_toolkit = FlexServiceToolkit()
        # for tool_function in unity_function_list:
        #     service_toolkit.add(tool_function)
        agent = ChapterChainAgent(
                    name="assistant",
                    model_config_name="openai_cfg",
                    verbose=True,
                    service_toolkit=service_toolkit,
                    max_iters=5,
                    chapter_chain=EMPTY_CHAPTER_CHAIN
                )
        if id is None: id = uuid.uuid4()
        self.agents[id] = agent
        return id, agent


    def get_agent(self, id=None):
        '''
            如果在get agent之前没有new agent，并用id进行访问，就返回一个章节链为空的agent
        '''
        if id:
            if  id in self.agents:
                return self.agents[id]
            self.agents[id] = self.new_agent(id)[1]
        else:
            id = str(uuid.uuid4())
            self.agents[id] = self.new_agent()[1]
        return id, self.agents[id]
        

    def invoke(self,id, query, user_backboard, reset=False):
        msg = Msg("user", query, role="user")
        agent = self.agents[id]
        res  = agent(msg, user_backboard=user_backboard)
        return res
        # 暂时取消生成器的写法
        # 假设 agent 是一个生成器
        # if reset or not hasattr(agent, '_generator'):
        #     agent._generator = agent(msg, user_backboard=user_backboard)
        
        # try:
        #     response = next(agent._generator)
        #     yield response
        # except StopIteration:
        #     return
    
    def function_call(self, id, arg_function):
        agent = self.agents[id]
        return agent.function_call(arg_function)

    def change_language(self, id, language):
        agent = self.agents[id]
        agent.change_language(language)

    def sample_judge(self, id, bb):
        agent = self.agents[id]
        
        with self.app.app_context():
            try:
                    print(bb.get_active_file_reletive_path())
                    assert bb is not None and bb.active_file_path != "" , "no path specified"
                    if bb.get_active_file_reletive_path().endswith(".py"):
                        self.socketio.emit('terminal', {'data': f"python {bb.get_active_file_reletive_path()}"}, room=id, namespace='/vscode')
                        self.socketio.emit('system_message', "代码试运行尚在建设，推荐直接通过命令行执行！", room=id, namespace='/agent')
                    else:
                        self.socketio.emit('system_message', "代码试运行尚仅支持python文件", room=id, namespace='/agent')
            except Exception as e:
                print(e)
                if (str(e) == "no path specified"):
                    self.socketio.emit("system_message", "未加载IDE文件同步，请尝试重新打开IDE文件或刷新页面，并确保插件已连接", room=id, namespace='/agent')

    def judge(self, id, bb):
        agent = self.agents[id]
        with self.app.app_context():
            try:
                if (agent.chapter_chain[agent.chapter_chain_now].append_tools is not None and len(agent.chapter_chain[agent.chapter_chain_now].append_tools)):
                    for tool, tool_kwargs in agent.chapter_chain[agent.chapter_chain_now].append_tools:
                        if tool.__name__ == "judge":
                                print(bb.get_active_file_reletive_path())
                                assert bb is not None and bb.active_file_path != "" , "no path specified"
                                if bb.get_active_file_reletive_path().endswith(".py"):
                                    # self.socketio.emit('terminal', {'data': f"python {bb.get_active_file_reletive_path()}"}, room=id, namespace='/vscode')
                                    self.socketio.emit('system_message', "代码评测中，请稍后", room=id, namespace='/agent')
                                    agent.function_call({'name':'judge','arguments':{'filepath': bb.get_active_file_reletive_path()}})
                                else:
                                    self.socketio.emit('system_message', "代码评测尚仅支持python文件", room=id, namespace='/agent')

                        else:
                            with self.app.app_context():
                                    self.socketio.emit("system_message", "本章学习无需评测", room=id, namespace='/agent')
                else:
                    with self.app.app_context():
                            self.socketio.emit("system_message", "本章学习无需评测", room=id, namespace='/agent')
            except Exception as e:
                if (str(e) == "no path specified"):
                    self.socketio.emit("system_message", "未加载IDE文件同步，请尝试重新打开IDE文件或刷新页面，并确保插件已连接", room=id, namespace='/agent')


import time
if __name__ == '__main__':
    manager = AgentManager()
    id, agent = manager.new_agent()

    # Main thread logic
    try:
        # Your code that might be interrupted
        while True:
            input_data = input("用户：")
            print(manager.invoke(id, input_data))
            
    except KeyboardInterrupt:
        print("Process interrupted by user.")
        # You can add any cleanup code here if needed
    finally:
        print("Exiting program.")
        # Code to run before the program exits