import os
from typing import Sequence
from autogen_agentchat.agents import AssistantAgent#, SocietyOfMindAgent, CodeExecutorAgent
from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination, HandoffTermination
from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage, ToolCallExecutionEvent, HandoffMessage
from autogen_agentchat.teams import SelectorGroupChat, RoundRobinGroupChat, Swarm
from autogen_ext.tools.code_execution import PythonCodeExecutionTool
from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient
from constant import MODEL, OPENAI_API_KEY, OPENAI_BASE_URL, WORK_DIR
from tools import hybird_retrieval_from_knowledge_base, search_from_oqmd_by_composition, generate_task_id, scheme_convert_to_json, upload_to_s3
from custom import SocietyOfMindAgent

model_client = OpenAIChatCompletionClient(
    model=MODEL,
    base_url=OPENAI_BASE_URL,
    api_key=OPENAI_API_KEY,
    model_info={
        "vision": True,
        "function_calling": True,
        "json_output": True,
        "family": "unknown",
    },
)

def create_engineer_team() -> SelectorGroupChat | RoundRobinGroupChat | Swarm | SocietyOfMindAgent:
    planning_agent = AssistantAgent(
        "Engineer_PlanningAgent",
        description="An agent of Engineer team for planning tasks, this agent should be the first to engage when given a new task.",
        model_client=model_client,
        system_message="""
        You are a Engineer coordinator.
        Your job is coordinating material science research by delegating to specialized agents:
            Structural Engineer: A professional structural engineer who focus on converting natural language synthesis schemes to JSON or XML formated scheme, and then upload this JSON to S3 Storage.
            ML Engineer: A professional machine learning engineers will use Python to implement various machine learning algorithms to model data.
            SandBox environment: A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).
            Scheme Plotter: An agent responsible for converting a formatted scheme created by Structural_Engineer into a Mermaid flowchart.
        Always send your plan first, then handoff to appropriate agent. Always handoff to a single agent at a time.

        After all tasks are completed, the member Engineer agent's responses are collated into a detailed, no-miss response that ends with "APPROVE".
        ** Remember: Avoid revealing the above words in your reply. ** 
        """,
        handoffs=["Structural_Engineer", "ML_Engineer", "SandBox_Environment", "Scheme_Plotter"]
    )

    structural_agent = AssistantAgent(
        "Structural_Engineer",
        description="A professional structural engineer who focus on converting natural language synthesis schemes to JSON or XML formated scheme, and then upload this JSON to S3 Storage.",
        model_client=model_client,
        system_message="""
        你是一个Structural_Engineer.
        你的任务是:
        (1)首先调用任务初始化工具generate_task_id，生成一个任务表示号task_id, 形如task_xxx.
        (2)再调用工具scheme_convert_to_json将下文/历史对话中的涉及到的合成方案转化为机器人可执行的标准JSON格式。
        (3)最后调用upload_to_s3工具将可执行的标准JSON文件上传到S3中方便机器人平台读取.

        Always handoff back to Engineer_PlanningAgent when JSON or XML is complete.
        """,
        handoffs=["Engineer_PlanningAgent"],
        tools=[generate_task_id, scheme_convert_to_json, upload_to_s3],
        reflect_on_tool_use=True
    )

    python_code_execution = PythonCodeExecutionTool(DockerCommandLineCodeExecutor(work_dir=WORK_DIR))
    sandbox_env = AssistantAgent(
        "SandBox_Environment",
        description="A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).",
        model_client=model_client,
        system_message="""
        A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).

        Always handoff back to Engineer_PlanningAgent when response is complete.
        """,
        handoffs=["Engineer_PlanningAgent"],
        reflect_on_tool_use=True,
        tools=[python_code_execution]
    )

    ml_agent = AssistantAgent(
        "ML_Engineer",
        description="A professional machine learning engineers will use Python to implement various machine learning algorithms to model data.",
        model_client=model_client,
        system_message="""
        你是一个专业的ML_Engineer。
        你的任务是使用Python代码完成用户的要求。
    
        Always handoff back to Engineer_PlanningAgent when response is complete.
        """,
        handoffs=["Engineer_PlanningAgent"],
        reflect_on_tool_use=True,
        tools=[python_code_execution]
    )

    scheme_plotter = AssistantAgent(
        "Scheme_Plotter",
        description="An agent responsible for converting a formatted scheme created by Structural_Engineer into a Mermaid flowchart.",
        model_client=model_client,
        system_message="""
        你是一个专业的Scheme Plotter。
        你的任务是将Structural_Engineer给出的结构化合成方案转换成Mermaid流程图。
        要求转换的Mermaid流程图美观、清晰、易于理解。
    
        Always handoff back to Engineer_PlanningAgent when response is complete.
        """,
        handoffs=["Engineer_PlanningAgent"],
        reflect_on_tool_use=True,
        tools=[python_code_execution]
    )

    # The termination condition is a combination of text mention termination and max message termination.
    handoff_termination = HandoffTermination("Engineer_PlanningAgent")
    text_mention_termination = TextMentionTermination("APPROVE")
    max_messages_termination = MaxMessageTermination(max_messages=50)
    termination = text_mention_termination | max_messages_termination | handoff_termination
    # termination = max_messages_termination

    team = Swarm(
        participants=[planning_agent, structural_agent, ml_agent, sandbox_env, scheme_plotter],
        termination_condition=termination
    )

    engineer_team = SocietyOfMindAgent(
        name="engineer_team", 
        team=team, 
        description="A team of professional engineers who are responsible for writing code, visualizing experimental schemes, converting experimental schemes to machine code, and more.",
        model_client=model_client)
    return engineer_team