from autogen.agentchat.groupchat import GroupChat
from autogen.agentchat.agent import Agent
from autogen.agentchat.assistant_agent import AssistantAgent
import json

class InputParserGroupChat(GroupChat):
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None  # Keep track of the previous speaker
        self.counter = 0

    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        # Check if last message suggests a next speaker or termination
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("Interior_designer")
        elif last_speaker_name == "Interior_designer":
            return self.agent_by_name("Admin")

class RequirementGroupChat(GroupChat):
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None  # Keep track of the previous speaker
        self.counter = 0

    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        # Check if last message suggests a next speaker or termination
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("requirements_analyzer")
        elif last_speaker_name == "requirements_analyzer":
            return self.agent_by_name("substructure_analyzer")
        elif last_speaker_name == "substructure_analyzer":
            return self.agent_by_name("Interior_designer")
        elif last_speaker_name == "Interior_designer":
            return self.agent_by_name("Designer_checker")
        elif last_speaker_name == "Designer_checker":
            if "SUCCESS" not in last_message["content"]:
                return self.agent_by_name("Interior_designer")
            else:
                return self.agent_by_name("requirements_analyzer")

class LanguageGroupChat(GroupChat):
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None  # Keep track of the previous speaker
        self.counter = 0

    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        # Check if last message suggests a next speaker or termination
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("Language_architect")
        elif last_speaker_name == "Language_architect":
            return self.agent_by_name("Language_architect")

class CalculationGroupChat(GroupChat):
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None  # Keep track of the previous speaker
        self.counter = 0

    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        # Check if last message suggests a next speaker or termination
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("Calculation_architect")
        elif last_speaker_name == "Calculation_architect":
            return self.agent_by_name("Calculation_architect")

class SceneGraphGroupChat(GroupChat):
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None  # Keep track of the previous speaker
        self.counter = 0

    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        # Check if last message suggests a next speaker or termination
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("Interior_architect")
        elif last_speaker_name == "Interior_architect":
            return self.agent_by_name("Schema_engineer")
        elif last_speaker_name == "Schema_engineer":
            return self.agent_by_name("Interior_architect")

            
class SchemaGroupChat(GroupChat): 
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None  # Keep track of the previous speaker
        self.counter = 0

    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        # Check if last message suggests a next speaker or termination
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("Schema_engineer")
        elif last_speaker_name == "Schema_engineer":
            return self.agent_by_name("Schema_debugger")
        elif last_speaker_name == "Schema_debugger":
            return self.agent_by_name("Schema_engineer")

class LayoutCorrectorGroupChat(GroupChat):
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None
    
    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("Spatial_corrector_agent")
        elif last_speaker_name == "Spatial_corrector_agent":
            return self.agent_by_name("Json_schema_debugger")
        elif last_speaker_name == "Json_schema_debugger":
            if "SUCCESS" not in last_message["content"]:
                return self.agent_by_name("Spatial_corrector_agent")
            else:
                return self.agent_by_name("Admin")

class ObjectDeletionGroupChat(GroupChat):
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None
    
    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("Object_deletion_agent")
        elif last_speaker_name == "Object_deletion_agent":
            return self.agent_by_name("Object_deletion_agent")

class LayoutRefinerGroupChat(GroupChat):
    def __init__(self, agents, messages, max_round=15):
        super().__init__(agents, messages, max_round)
        self.previous_speaker = None  # Keep track of the previous speaker
        self.counter = 0

    def select_speaker(self, last_speaker: Agent, selector: AssistantAgent):
        # Check if last message suggests a next speaker or termination
        last_message = self.messages[-1] if self.messages else None
        last_speaker_name = last_speaker.name if last_speaker else None

        self.previous_speaker = last_speaker

        if last_speaker_name == "Admin":
            return self.agent_by_name("Layout_refiner")
        elif last_speaker_name == "Layout_refiner":
            return self.agent_by_name("Json_schema_debugger")
        elif last_speaker_name == "Json_schema_debugger":
            if "SUCCESS" not in last_message["content"]:
                return self.agent_by_name("Layout_refiner")
            else:
                return self.agent_by_name("Admin")