import asyncio
from src.agents.assistantagent import TrackableAssistantAgent
from src.agents.userproxyagent import TrackableUserProxyAgent
import streamlit as st
import autogen


class MultiAgentCodeExecution:
    def __init__(self, assistant_name, user_proxy_name, llm_config, problem):
        self.coder = TrackableAssistantAgent(name=assistant_name[0],
                                                 system_message="""you are helpful assistant and efficient in writing code in python.""",
                                                 human_input_mode="NEVER",
                                                 llm_config=llm_config,
                                                 )
        
        self.pm = TrackableAssistantAgent(name=assistant_name[1],
                                                 system_message="""You are efficient in Creative in software product ideas. Reply "TERMINATE" in 
                                                 the end when everything is done """,
                                                 human_input_mode="NEVER",
                                                 llm_config=llm_config,
                                                 )
         
        self.user_proxy = TrackableUserProxyAgent(name=user_proxy_name,
                                                  system_message="You are human Admin",
                                                  human_input_mode="NEVER",
                                                  llm_config=llm_config,
                                                  code_execution_config={"last_n_messages":2,"work_dir" : "./codegen","use_docker":False},
                                                  is_termination_msg=lambda x: x.get("content", "").strip().endswith(
                                                      "TERMINATE"))
        
        self.groupchat = autogen.GroupChat(agents=[self.user_proxy,self.coder,self.pm],messages=[],max_round=2)
        self.manager = autogen.GroupChatManager(groupchat=self.groupchat,llm_config=llm_config)
        self.problem = problem
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        
    def _reset(self):
        self.coder.reset()
        self.pm.reset()
        
        

    async def initiate_chat(self):
        self._reset()
        await self.user_proxy.a_initiate_chat(self.manager,max_turns=2, message=self.problem, clear_history=st.session_state["chat_with_history"])

    def run(self):
        self.loop.run_until_complete(self.initiate_chat())