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
from .utils import load_agent_configs
import os
from .retrieval_group import init_retrieval_group
from .constant import config_list, STREAM, SILENT, WORK_DIR, CACHE



agent_configs = load_agent_configs(os.path.join(os.path.dirname(os.path.abspath(__file__)), "config/converter_group.yaml"))
venv_context = create_virtual_env(WORK_DIR)
llm_config = {"config_list": config_list, "stream": STREAM, "cache_seed": CACHE}


def init_converter_group():

    outer_converter_agent = autogen.ConversableAgent(
        name="Outer_Converter_Admin",
        human_input_mode="NEVER",
        # human_input_mode="TERMINATE",
        code_execution_config={
            "work_dir": WORK_DIR,
            "use_docker": False,
        },
        is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        description="Outer_Converter_Admin",
        default_auto_reply="continue",
        max_consecutive_auto_reply=1
    )

    scheme_converter_name = "scheme_converter"
    scheme_converter = autogen.AssistantAgent(
        name=scheme_converter_name,
        system_message=agent_configs[scheme_converter_name]["system_message"],
        llm_config=llm_config,    
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        human_input_mode="TERMINATE",
        description="scheme_converter",
        max_consecutive_auto_reply=3
    )


    converter_critic_name = "converter_critic"
    converter_critic = autogen.AssistantAgent(
        name=converter_critic_name,
        system_message=agent_configs[converter_critic_name]['system_message'],
        llm_config=llm_config,
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        human_input_mode="TERMINATE",
        description="converter_critic",
        max_consecutive_auto_reply=3
    )


    mergrid_ploter_name = "mergrid_ploter"
    mergrid_ploter = autogen.AssistantAgent(
        name=mergrid_ploter_name,
        system_message=agent_configs[mergrid_ploter_name]['system_message'],
        llm_config=llm_config,
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        human_input_mode="TERMINATE",
        description="mergrid_ploter",
        max_consecutive_auto_reply=1
    )

    scheme_code_writer_name = "scheme_code_writer"
    scheme_code_writer = autogen.AssistantAgent(
        name=scheme_code_writer_name,
        system_message=agent_configs[scheme_code_writer_name]['system_message'],
        llm_config=llm_config,
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        human_input_mode="TERMINATE",
        description="scheme_code_writer",
        max_consecutive_auto_reply=3
    )

    scheme_code_critic_name = "scheme_code_critic"
    scheme_code_critic = autogen.AssistantAgent(
        name=scheme_code_critic_name,
        system_message=agent_configs[scheme_code_critic_name]['system_message'],
        llm_config=llm_config,
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        human_input_mode="TERMINATE",
        description="scheme_code_critic",
        max_consecutive_auto_reply=3
    )

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

        if last_speaker is outer_converter_agent:
            return scheme_converter
        elif last_speaker is scheme_converter:
            return converter_critic
        elif last_speaker is converter_critic:
            if "OPTIMIZE" in messages[-1]["content"]:
                return scheme_converter
            elif "TERMINATE" in messages[-1]["content"]:
                return mergrid_ploter
        elif last_speaker is mergrid_ploter:
            return scheme_code_writer
        elif last_speaker is scheme_code_writer:
            return scheme_code_critic
        elif last_speaker is scheme_code_critic:
            if "OPTIMIZE" in messages[-1]["content"]:
                return scheme_code_writer
            elif "TERMINATE" in messages[-1]["content"]:
                return outer_converter_agent

    converter_group = autogen.GroupChat(
        agents=[outer_converter_agent, scheme_converter, converter_critic, mergrid_ploter, scheme_code_writer, scheme_code_critic],
        messages=[],
        speaker_selection_method=state_transition, # custom speaker selection method
        max_round=12,
    )

    inner_converter_admin = autogen.GroupChatManager(
        name="Converter_Group_Admin",
        description="Converter_Group_Admin",
        groupchat=converter_group,
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        llm_config=llm_config,
        system_message=agent_configs["admin"]['system_message']
    )

    return inner_converter_admin, outer_converter_agent