from autogen.coding import LocalCommandLineCodeExecutor
import autogen
from .utils import send_instruction_to_robot, get_uv_latest_file, get_pl_latest_file
from .constant import config_list, STREAM, CACHE
llm_config = {"config_list": config_list, "stream": STREAM, "cache_seed": CACHE}


def init_executor_group(work_dir, venv_context):
    outer_executor_agent = autogen.ConversableAgent(
        name="Outer_Executor_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_Executor_Admin",
        default_auto_reply="continue",
        max_consecutive_auto_reply=1
    )

    executor_env = LocalCommandLineCodeExecutor(
        work_dir=work_dir, 
        virtual_env_context=venv_context, 
        functions=[send_instruction_to_robot, get_uv_latest_file, get_pl_latest_file]
    )

    expriment_executor = autogen.UserProxyAgent(
        name="experiment_executor",
        human_input_mode="NEVER",
        llm_config=False,
        code_execution_config={
            "executor": executor_env,
        },
        description="experiment_executor",
        max_consecutive_auto_reply=1
    )

    expriment_coder_system_message = """
    You are a expriment_coder. 
    You are given a task and some code snippets. Your task is to write code to call the function of send_instruction_to_robot, that will send the instruction to the robot.
    You need to write the rest of the code to answer the user's request. 
    DONOT ANSWER THE USER'S REQUEST DIRECTLY.
    """
    expriment_coder_system_message += executor_env.format_functions_for_prompt()
    expriment_code_writer = autogen.AssistantAgent(
        name="expriment_code_writer",
        human_input_mode="NEVER",
        system_message=expriment_coder_system_message,
        # llm_config={"config_list": config_list},
        llm_config=llm_config,
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        # human_input_mode="TERMINATE",
        description="expriment_code_writer",
        max_consecutive_auto_reply=1
    )

    data_collector = autogen.UserProxyAgent(
        name="data_collector",
        human_input_mode="NEVER",
        llm_config=False,
        code_execution_config={
            "executor": executor_env,
        },
        description="data_collector",
        max_consecutive_auto_reply=1
    )


    collector_coder_system_message = """
    You are a collector_coder. 
    You are given a task and some code snippets. Your task is to write code to call the functions of get_uv_latest_file and get_pl_latest_file at the same time.
    You need to write the rest of the code to answer the user's request. 
    DONOT ANSWER THE USER'S REQUEST DIRECTLY.
    """
    collector_coder_system_message += executor_env.format_functions_for_prompt()
    collector_code_writer = autogen.AssistantAgent(
        name="collector_code_writer",
        human_input_mode="NEVER",
        system_message=collector_coder_system_message,
        # llm_config={"config_list": config_list},
        llm_config=llm_config,
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        # human_input_mode="TERMINATE",
        description="collector_code_writer",
        max_consecutive_auto_reply=1
    )

    executor_group = autogen.GroupChat(
        agents=[outer_executor_agent, expriment_code_writer, expriment_executor, collector_code_writer, data_collector],
        messages=[],
        speaker_selection_method='round_robin',  # With two agents, this is equivalent to a 1:1 conversation.
        allow_repeat_speaker=False,
        max_round=6,
    )

    inner_executor_admin = autogen.GroupChatManager(
        name="Inner_Executor_Admin",
        groupchat=executor_group,
        # is_termination_msg=lambda x: x.get("content", "").find("TERMINATE") >= 0,
        # llm_config={"config_list": config_list},
        llm_config=llm_config,
        human_input_mode="NEVER",
        description="Inner_Executor_Admin",
    )

    return inner_executor_admin, outer_executor_agent