import asyncio
from typing import Sequence
from autogen_core import CancellationToken
from autogen_agentchat.agents import AssistantAgent, SocietyOfMindAgent, UserProxyAgent
from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination, HandoffTermination, SourceMatchTermination, ExternalTermination
from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage, ToolCallExecutionEvent
from autogen_agentchat.teams import SelectorGroupChat, RoundRobinGroupChat
from autogen_agentchat.ui import Console
from autogen_agentchat.base import Handoff
from autogen_ext.models.openai import OpenAIChatCompletionClient
from _backend.constant import MODEL, OPENAI_API_KEY, OPENAI_BASE_URL
from _backend.scientist_team import create_scientist_team

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",
    },
)

async def _multiagent_with_rag_cot(task: str = "") -> dict:
    user = UserProxyAgent("user_agent", input_func=input)

    scientist_team = create_scientist_team()
    
    result = {}
    planning_agent = AssistantAgent(
        "PlanningAgent",
        description="An agent 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 planning agent.
        Your job is to break down complex Materials science research tasks into smaller, manageable subtasks.
        Assign these subtasks to the appropriate sub-teams; not all sub-teams are required to participate in every task.
        Your sub-teams are:
            1. Scientist: A professional team of material scientists who are mainly responsible for consulting on material synthesis, structure, application and properties.
                - The scientist team has the following members: 
                1.1 Synthesis Scientist: who is good at giving perfect and correct synthesis solutions.
                1.2 Structure Scientist: focusing on agents of structural topics in materials science.
                1.3 Property Scientist: focuses on physical and chemistry property topics in materials science.
                1.4 Application Scientist: Focus on practical applications of materials, such as devices, chips, etc.

        You only plan and delegate tasks - you do not execute them yourself.
        
        回答时你需要初始化/更新如下任务分配表和Mermaid流程图，并按顺序执行，使用如下格式并利用：
        | Team_name   | Member_name   | sub-task                             |
        | ----------- | ------------- | ------------------------------------ |
        | <team_name> | <member_name> | <status: brief sub-task description> |
        
        ```mermaid
        graph TD
        User[User]
        subgraph <team_name>
            A1[<member_name>]
        end
        style xxx # 推荐多样的风格
        ...
        User --> A1
        ...
        ```

        每次回答时，你需要清晰明确的指出已经完成的子任务下一步子任务，使用如下格式：
        **已完成子任务：**
        1. <team> : <subtask>
        **Next sub-task:**
        n. <team> : <subtask>
        
        Determine if all sub-teams have completed their tasks, and if so, summarize the findings and end with "TERMINATE".
        After all tasks of Scientist team are completed, ends with "TERMINATE".
        """,
        reflect_on_tool_use=False
    )

    # The termination condition is a combination of text mention termination and max message termination.
    text_mention_termination = TextMentionTermination("TERMINATE")
    max_messages_termination = MaxMessageTermination(max_messages=200)
    source_matched_termination = SourceMatchTermination(["scientist_team"])
    ext_termination = ExternalTermination()
    termination = text_mention_termination | max_messages_termination | source_matched_termination

    # The selector function is a function that takes the current message thread of the group chat
    # and returns the next speaker's name. If None is returned, the LLM-based selection method will be used.
    def selector_func(messages: Sequence[AgentEvent | ChatMessage]) -> str | None:
        if messages[-1].source != planning_agent.name:
            return planning_agent.name # Always return to the planning agent after the other agents have spoken.
        elif "HUMAN" in messages[-1].content:
            return user.name
        return None

    team = SelectorGroupChat(
        [planning_agent, user, scientist_team],
        model_client=model_client, # Use a smaller model for the selector.
        termination_condition=termination,
        selector_func=selector_func,
    )
    # team.run(task=task)
    # await Console(team.run_stream(task=task))
    result = ""

    async for message in team.run_stream(task=task):
        # if isinstance(message, TextMessage):
        #     print(f"----------------{message.source}----------------\n {message.content}")
        # elif isinstance(message, ToolCallExecutionEvent):
        #     print(f"----------------{message.source}----------------\n {message.content}")

        # if message.source == "Scientist_StructureAgent" or message.source == "Scientist_PropertyAgent" \
        #     or message.source == "Scientist_ApplicationAgent" or message.source == "Scientist_SynthesisAgent":
        #     return message.content
        if isinstance(message, TextMessage) and message.source == "scientist_team":
            message.content += "\nTERMINATE"
            result = message.content
            ext_termination.set()
            # break
    return result

# Example usage in another function
async def main_1(task: str):
    # result = await main(input("Enter your instructions below: \n"))
    result = await _multiagent_with_rag_cot(task)
    # result = await main("查一下CsPbBr3的晶体结构")

    return result

if __name__ == "__main__":
    asyncio.run(main_1("how to synthesize CsPbBr3 nanocubes at room temperature"))
    # result = asyncio.run(_multiagent_with_rag_cot("CsPbBr3 nanocubes 的结构是怎样的?"))
    # print(result)