from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status
import requests  # 用于调用 AutoGen API
from django.http import StreamingHttpResponse
from contextlib import asynccontextmanager
from autogen.code_utils import create_virtual_env
from autogen.coding import LocalCommandLineCodeExecutor
from autogen.agentchat.contrib.capabilities.teachability import Teachability
from autogen.agentchat.contrib.capabilities.vision_capability import VisionCapability
from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent
from pathlib import Path
import autogen
import os
from .constant import config_list, STREAM, SILENT, WORK_DIR
from .utils import load_agent_configs
from .retrieval_group import init_retrieval_group
from .generate_group import init_generate_group
from .converter_group import init_converter_group
from .executor_group import init_executor_group
from .optimize_group import init_optimize_group

import io
import sys
from contextlib import redirect_stdout
from websockets.sync.client import connect as ws_connect
from autogen.io.websockets import IOWebsockets
import json
import asyncio
from channels.generic.websocket import AsyncWebsocketConsumer

venv_context = create_virtual_env(WORK_DIR)
llm_config = {"config_list": config_list, "stream": True}



class ChatConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        await self.accept()

    async def disconnect(self, close_code):
        pass

    async def receive(self, text_data):
        # 在这里处理接收到的消息
        print(f"Received message: {text_data}")

        # 异步启动 WebSocket 连接和处理过程
        await self.handle_websocket_connection(text_data)

    async def handle_websocket_connection(self, text_data):
        # 启动独立的 WebSocket 服务器，使用异步线程避免阻塞
        uri = "ws://localhost:8765"  # 示例 URI，可以替换为你的 URI
        
        # 启动 WebSocket 服务器在后台线程中运行
        await asyncio.to_thread(self.run_websocket_server, uri, text_data)

    def run_websocket_server(self, uri, text_data):
        """
        这是一个同步方法，负责启动 IOWebsockets 的服务器并处理消息
        使用 asyncio.to_thread() 将其异步化
        """
        with IOWebsockets.run_server_in_thread(on_connect=self.handle_message, port=8765) as uri:
            print(f" - test_setup() with websocket server running on {uri}.", flush=True)

            with ws_connect(uri) as websocket:
                print(f" - Connected to server on {uri}", flush=True)

                print(" - Sending message to server.", flush=True)
                websocket.send(text_data)

                import re
                current_agent = "User"
                while True:
                    message = websocket.recv()
                    message = message.decode("utf-8") if isinstance(message, bytes) else message
                    print(message, end="", flush=True)

                    cleaned_string = re.sub(r'\x1b\[[0-?]*[ -/]*[@-~]', '', message)
                    # match = re.findall(r"(\w+)\s*\(to\s+(\w+)\)", cleaned_string)
                    # print(match)
                    # if len(match)==1:
                    #     current_agent = match[0][0] if current_agent != match[0][0] else current_agent
                        # print(f"A: {match[0][0]}, B: {match[0][1]}")
                    if "Next speaker" in cleaned_string:
                        match = re.search(r"Next\s+speaker:\s+(\w+)", cleaned_string)
                        current_agent = match.group(1)
                    else:
                        if cleaned_string == "\n--------------------------------------------------------------------------------\n":
                            continue
                        if cleaned_string == "\n********************************************************************************\n":
                            continue
                        if cleaned_string == "Starting a new chat....\n":
                            continue
                        if cleaned_string == "\n>>>>>>>> USING AUTO REPLY...\n":
                            continue
                        match = re.findall(r"(\w+)\s*\(to\s+(\w+)\)", cleaned_string)
                        if len(match)==1:
                            continue

                        if current_agent in ['Outer_Retrieval_Admin', 'Outer_Generate_Admin', 'Outer_Converter_Admin']:
                            current_agent = current_agent.replace('Outer_', '')
                        if current_agent in ['vector_code_executor']:
                            continue
                        
                        if current_agent == 'User':
                            group_name = 'Planner'
                        if current_agent in ['vector_searcher','vector_code_executor', 'graphrag_searcher', 'graphrag_code_executor', 'web_searcher', 'web_summary', 'Outer_Retrieval_Admin']:
                            group_name = 'Retrieval'
                        if current_agent in ['structure_scientist','property_scientist', 'application_scientist', 'synthesis_scientist', 'scheme_critic', 'Outer_Generate_Admin']:
                            group_name = 'Generator'
                        if current_agent in ['scheme_converter','converter_critic', 'mergrid_ploter', 'scheme_code_writer', 'scheme_code_critic', 'Outer_Converter_Admin']:
                            group_name = 'Converter'
                        if current_agent in ['experiment_executor','expriment_code_writer', 'data_collector', 'collector_code_writer', 'Outer_Executor_Admin']:
                            group_name = 'Executor'
                        if current_agent in ['analysis_executor','analysis_pl_uv', 'analysis_picturer', 'Experiment_Optimizer', 'optimizer_critic', 'Outer_Analysis_Admin']:
                            group_name = 'Optimizer'

                        content = {"group_name": group_name, "agent_name": current_agent.replace("_", " ").title(), "content": cleaned_string}
                        # 通过 WebSocket 消费者发送消息
                        asyncio.run(self.send(text_data=json.dumps(content)))

                    # if "TERMINATE" in message:
                    #     print(" - Received TERMINATE message. Exiting.", flush=True)
                    #     break


    def handle_message(self, iostream: IOWebsockets):
        initial_msg = iostream.input()
        agent_configs = load_agent_configs(
            os.path.join(os.path.dirname(os.path.abspath(__file__)), "config/plan_group.yaml"))
        user = autogen.UserProxyAgent(
            name="User",
            chat_messages=None,
            human_input_mode="NEVER",
            code_execution_config={
                "work_dir": WORK_DIR,
                "use_docker": False,
            },
            is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
            description="User",
        )

        inner_retrieval_admin, outer_retrieval_agent = init_retrieval_group(WORK_DIR, venv_context)
        inner_generate_admin, outer_generate_agent = init_generate_group(outer_retrieval_agent, inner_retrieval_admin)
        inner_converter_admin, outer_converter_agent = init_converter_group()
        inner_executor_admin, outer_executor_agent = init_executor_group(WORK_DIR, venv_context)
        inner_analysis_admin, outer_analysis_agent, optimizer = init_optimize_group(WORK_DIR, venv_context)

        def state_transition(last_speaker, groupchat):
            messages = groupchat.messages

            if last_speaker is user:
                if len(messages) <= 1:
                    return outer_generate_agent
                else:
                    return "auto"
            elif last_speaker is outer_generate_agent:
                if "synthesis" in messages[-1]["content"].lower():
                    return outer_converter_agent
                else:
                    return user
            elif last_speaker is outer_converter_agent:
                return outer_executor_agent
            elif last_speaker is outer_executor_agent:
                return outer_analysis_agent
            elif last_speaker is outer_analysis_agent:
                return optimizer
            else:
                return user

        matagent_group = autogen.GroupChat(
            agents=[user, outer_generate_agent, outer_converter_agent, outer_executor_agent, outer_analysis_agent,
                    optimizer],
            messages=[],
            speaker_selection_method=state_transition,
            max_round=50,
        )

        matagent_admin_name = "Planer"
        matagent_admin = autogen.GroupChatManager(
            name=matagent_admin_name,
            groupchat=matagent_group,
            # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
            llm_config=llm_config,
            system_message=agent_configs[matagent_admin_name]['system_message'],
            description=matagent_admin_name
        )

        outer_generate_agent.register_nested_chats(
            [
                {"recipient": inner_generate_admin, "max_turn": 1, "summary_method": "last_msg", "silent": SILENT},
            ],
            trigger=matagent_admin,
            #reply_func_from_nested_chats =
        )

        outer_converter_agent.register_nested_chats(
            [
                {"recipient": inner_converter_admin, "max_turn": 1, "summary_method": "last_msg", "silent": SILENT},
            ],
            trigger=matagent_admin,
        )

        outer_executor_agent.register_nested_chats(
            [
                {"recipient": inner_executor_admin, "max_turn": 1, "summary_method": "last_msg", "silent": SILENT},
            ],
            trigger=matagent_admin,
        )

        outer_analysis_agent.register_nested_chats(
            [
                {"recipient": inner_analysis_admin, "max_turn": 1, "summary_method": "last_msg", "silent": SILENT},
            ],
            trigger=matagent_admin,
        )

        user.initiate_chat(
            matagent_admin,
            message=initial_msg,  # "如何在常温条件下制备CsPbBr3纳米立方体",
            cache=False
        )




# class ChatConsumer(AsyncWebsocketConsumer):
#     async def connect(self):
#         await self.accept()

#     async def disconnect(self, close_code):
#         pass

#     async def receive(self, text_data):
#         #text_data_json = json.loads(text_data)
#         #message = text_data_json['message']

#         # 在这里处理接收到的消息
#         print(f"Received message: {text_data}")

#         agent_configs = load_agent_configs(
#             os.path.join(os.path.dirname(os.path.abspath(__file__)), "config/plan_group.yaml"))
#         user = autogen.UserProxyAgent(
#             name="User",
#             chat_messages=None,
#             human_input_mode="NEVER",
#             code_execution_config={
#                 "work_dir": WORK_DIR,
#                 "use_docker": False,
#             },
#             is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
#             description="User",
#         )

#         inner_retrieval_admin, outer_retrieval_agent = init_retrieval_group(WORK_DIR, venv_context)
#         inner_generate_admin, outer_generate_agent = init_generate_group(outer_retrieval_agent, inner_retrieval_admin)
#         inner_converter_admin, outer_converter_agent = init_converter_group()
#         inner_executor_admin, outer_executor_agent = init_executor_group(WORK_DIR, venv_context)
#         inner_analysis_admin, outer_analysis_agent, optimizer = init_optimize_group(WORK_DIR, venv_context)

#         def state_transition(last_speaker, groupchat):
#             messages = groupchat.messages

#             if last_speaker is user:
#                 if len(messages) <= 1:
#                     return outer_generate_agent
#                 else:
#                     return "auto"
#             elif last_speaker is outer_generate_agent:
#                 if "synthesis" in messages[-1]["content"].lower():
#                     return outer_converter_agent
#                 else:
#                     return user
#             elif last_speaker is outer_converter_agent:
#                 return outer_executor_agent
#             elif last_speaker is outer_executor_agent:
#                 return outer_analysis_agent
#             elif last_speaker is outer_analysis_agent:
#                 return optimizer
#             else:
#                 return user

#         matagent_group = autogen.GroupChat(
#             agents=[user, outer_generate_agent, outer_converter_agent, outer_executor_agent, outer_analysis_agent,
#                     optimizer],
#             messages=[],
#             speaker_selection_method=state_transition,
#             max_round=50,
#         )

#         matagent_admin_name = "Planer"
#         matagent_admin = autogen.GroupChatManager(
#             name=matagent_admin_name,
#             groupchat=matagent_group,
#             # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
#             llm_config=llm_config,
#             system_message=agent_configs[matagent_admin_name]['system_message'],
#             description=matagent_admin_name
#         )

#         outer_generate_agent.register_nested_chats(
#             [
#                 {"recipient": inner_generate_admin, "max_turn": 1, "summary_method": "last_msg", "silent": SILENT},
#             ],
#             trigger=matagent_admin,
#             #reply_func_from_nested_chats =
#         )

#         outer_converter_agent.register_nested_chats(
#             [
#                 {"recipient": inner_converter_admin, "max_turn": 1, "summary_method": "last_msg", "silent": SILENT},
#             ],
#             trigger=matagent_admin,
#         )

#         outer_executor_agent.register_nested_chats(
#             [
#                 {"recipient": inner_executor_admin, "max_turn": 1, "summary_method": "last_msg", "silent": SILENT},
#             ],
#             trigger=matagent_admin,
#         )

#         outer_analysis_agent.register_nested_chats(
#             [
#                 {"recipient": inner_analysis_admin, "max_turn": 1, "summary_method": "last_msg", "silent": SILENT},
#             ],
#             trigger=matagent_admin,
#         )

#         def print_messages(recipient, messages, sender, config):
#             # if "callback" in config and config["callback"] is not None:
#             #     callback = config["callback"]
#             #     callback(sender, recipient, messages[-1])
#             print(f"~~~~~~~~~~~~~~~~: {recipient.name} | num messages: {len(messages)}")
#             self.send(text_data=json.dumps(messages))
#             return False, None  # required to ensure the agent communication flow continues

#         # user.register_reply(
#         #     [autogen.Agent, None],
#         #     reply_func=print_messages,
#         #     config={"callback": None},
#         # )

#         user.initiate_chat(
#             matagent_admin,
#             message=text_data  # "如何在常温条件下制备CsPbBr3纳米立方体",
#         )
        

#         #json.dump(user.chat_messages[matagent_admin], open(output.txt, "w"), indent=2)
#         response_text = outer_generate_agent.last_message(matagent_admin)
#         await self.send(text_data=json.dumps(response_text))



# with IOWebsockets.run_server_in_thread(on_connect=on_connect, host='192.168.42.130', port=8765) as uri:
#     print(f" - test_setup() with websocket server running on {uri}.", flush=True)
#     with ws_connect(uri) as websocket:
#         print(f" - Connected to server on {uri}", flush=True)
#
#         print(" - Sending message to server.", flush=True)
#         websocket.send("开始聊天吧")
    #
    #     while True:
    #         message = websocket.recv()
    #         message = message.decode("utf-8") if isinstance(message, bytes) else message
    #
    #         print(message, end="", flush=True)
    #         if "TERMINATE" in message:
    #             print()
    #             print(" - Received TERMINATE message. Exiting.", flush=True)
    #             break

class UserView(APIView):
    def post(self, request, *args, **kwargs):
        # 从请求中获取用户输入
        user_name = request.data.get('user_name')
        pass_word = request.data.get('pass_word')

        # 验证用户名和密码
        if user_name=='test' and pass_word=='111111':
            return Response({"token": "token_test"},status=status.HTTP_200_OK)
        else:
            return Response({"error": "用户名或密码错误"},status=status.HTTP_401_UNAUTHORIZED)

class ModelView(APIView):
    def get(self, request, *args, **kwargs):
        token = request.headers.get("token", None)
        if not token:
            return Response({"error": "请登录以后再重试"}, status=status.HTTP_203_NON_AUTHORITATIVE_INFORMATION)
        model_lists = {
            "count":"1",
            "data":
                [
                    {"model_name":"model1","model_des":"model1"},
                    {"model_name":"model1","model_des":"model2"}
                ]
        }
        return Response(model_lists,status=status.HTTP_200_OK)


