from typing import List

from swarms.prompts.agent_judge_prompt import AGENT_JUDGE_PROMPT
from swarms.structs.agent import Agent
from swarms.structs.conversation import Conversation
from swarms.utils.any_to_str import any_to_str

from loguru import logger


class AgentJudge:
    """
    A class to represent an agent judge that processes tasks and generates responses.

    Attributes:
        agent_name (str): The name of the agent judge.
        system_prompt (str): The system prompt for the agent.
        model_name (str): The model name used for generating responses.
        conversation (Conversation): An instance of the Conversation class to manage conversation history.
        max_loops (int): The maximum number of iterations to run the tasks.
        agent (Agent): An instance of the Agent class that performs the task execution.

    Methods:
        step(tasks: List[str]) -> str:
            Processes a list of tasks and returns the agent's response.

        run(tasks: List[str]) -> List[str]:
            Executes the tasks in a loop, updating context and collecting responses.
    """

    def __init__(
        self,
        agent_name: str = "agent-judge-01",
        system_prompt: str = AGENT_JUDGE_PROMPT,
        model_name: str = "openai/o1",
        max_loops: int = 1,
    ) -> None:
        """
        Initializes the AgentJudge with the specified parameters.

        Args:
            agent_name (str): The name of the agent judge.
            system_prompt (str): The system prompt for the agent.
            model_name (str): The model name used for generating responses.
            max_loops (int): The maximum number of iterations to run the tasks.
        """
        self.agent_name = agent_name
        self.system_prompt = system_prompt
        self.model_name = model_name
        self.conversation = Conversation(time_enabled=False)
        self.max_loops = max_loops

        self.agent = Agent(
            agent_name=agent_name,
            agent_description="You're the agent judge",
            system_prompt=AGENT_JUDGE_PROMPT,
            model_name=model_name,
            max_loops=1,
        )

    def step(self, tasks: List[str]) -> str:
        """
        Processes a list of tasks and returns the agent's response.

        Args:
            tasks (List[str]): A list of tasks to be processed.

        Returns:
            str: The response generated by the agent.
        """
        prompt = any_to_str(tasks)
        logger.debug(f"Running step with prompt: {prompt}")

        print(prompt)

        response = self.agent.run(
            task=f"Evaluate the following output or outputs: {prompt}"
        )
        logger.debug(f"Received response: {response}")

        return response

    def run(self, tasks: List[str]) -> List[str]:
        """
        Executes the tasks in a loop, updating context and collecting responses.

        Args:
            tasks (List[str]): A list of tasks to be executed.

        Returns:
            List[str]: A list of responses generated by the agent for each iteration.
        """
        responses = []
        context = ""

        for _ in range(self.max_loops):
            # Add context to the tasks if available
            if context:
                contextualized_tasks = [
                    f"Previous context: {context}\nTask: {task}"
                    for task in tasks
                ]
            else:
                contextualized_tasks = tasks

            # Get response for current iteration
            current_response = self.step(contextualized_tasks)
            responses.append(current_response)
            logger.debug(
                f"Current response added: {current_response}"
            )

            # Update context for next iteration
            context = current_response

            # Add to conversation history
            logger.debug("Added message to conversation history.")

        return responses
