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, CACHE
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


venv_context = create_virtual_env(WORK_DIR)
llm_config = {"config_list": config_list, "stream": STREAM, "cache_seed": CACHE}


def main():
    agent_configs = load_agent_configs(os.path.join(os.path.dirname(os.path.abspath(__file__)), "config/plan_group.yaml"))
    user = autogen.UserProxyAgent(
        name="User",
        human_input_mode="ALWAYS",
        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,
    )

    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="如何在常温条件下制备CsPbBr3纳米立方体",
        message="how to synthesis of CsPbBr3 Perovskite NCs at room temperature?"
        # message="how to synthesis CsPbBr3 nanocubes at room temperature?"
        # message="什么是钙钛矿？"
        # message="Please prepare few layers graphene from graphite powder.",
        # message="Can you please prepare black phosphorusene with improved stability from black phosphorus crystals powder?",
        # message="Can you synthesize gold nanorods by seed-mediated method with absorption peaks at 820 nm?",
        # message="Please synthesize CsPbBr3 nanocubes with a fluorescence emission wavelength of 520 nm at room temperature?",
        # message="Please design a new hybridized halide perovskite composite material that is biocompatible and water-stable",
        # message="please use phospholipid membrane as shell to encapsulate hybrid perovskite"
        # Now I want a novel bright perovskite composite based CH3NH3PbBr3 and phospholipid membrane(PM) to improve the stability and biocompatibility, please synthesis this materials under room temperature
    )
    print(outer_generate_agent.last_message(matagent_admin))
    print(matagent_admin)


if __name__ == "__main__":
    main()