# -*- coding: utf-8 -*-
"""
Unit tests for pipeline classes and functions
"""

import unittest
import random

from agentscope.pipelines import (
    SequentialPipeline,
    IfElsePipeline,
    SwitchPipeline,
    ForLoopPipeline,
    WhileLoopPipeline,
    sequentialpipeline,
    ifelsepipeline,
)

from agentscope.agents import AgentBase


class Add(AgentBase):
    """Operator for adding a value"""

    def __init__(self, name: str, value: int) -> None:
        self.name = name
        self.value = value
        super().__init__(name=name)

    def __call__(self, x: dict = None) -> dict:
        x["value"] += self.value
        return x


class Mult(AgentBase):
    """Operator for multiplying a value"""

    def __init__(self, name: str, value: int) -> None:
        self.name = name
        self.value = value
        super().__init__(name=name)

    def __call__(self, x: dict = None) -> dict:
        x["value"] *= self.value
        return x


class If_agent(AgentBase):
    """Operator for If condition"""

    def __init__(self, name: str) -> None:
        self.name = name
        super().__init__(name=name)

    def __call__(self, _: dict = None) -> dict:
        return {"operation": "A"}


class Else_agent(AgentBase):
    """Operator for Else condition"""

    def __init__(self, name: str) -> None:
        self.name = name
        super().__init__(name=name)

    def __call__(self, _: dict = None) -> dict:
        return {"operation": "B"}


class Case_agent(AgentBase):
    """Operator for Switch condition"""

    def __init__(self, name: str) -> None:
        self.name = name
        super().__init__(name=name)

    def __call__(self, x: dict = None) -> dict:
        return {"operation": x["text"].strip()}


class Default_agent(AgentBase):
    """Operator for Switch default condition"""

    def __init__(self, name: str) -> None:
        self.name = name
        super().__init__(name=name)

    def __call__(self, x: dict = None) -> dict:
        return {"operation": "IDLE"}


class Loop_for_agent(AgentBase):
    """Operator for Loop condition"""

    def __init__(self, name: str) -> None:
        self.name = name
        super().__init__(name=name)

    def __call__(self, x: dict = None) -> dict:
        x["value"] += 1
        return x


class Loop_while_agent(AgentBase):
    """Operator for Loop condition"""

    def __init__(self, name: str) -> None:
        self.name = name
        super().__init__(name=name)

    def __call__(self, x: dict = None) -> dict:
        x["token_num"] += random.randint(1, 100)
        x["round"] += 1
        return x


class BasicPipelineTest(unittest.TestCase):
    """Test cases for Basic Pipelines"""

    def test_sequential_pipeline(self) -> None:
        """Test SequentialPipeline executes agents sequentially"""

        add1 = Add("add1", 1)
        add2 = Add("add2", 2)
        mult3 = Mult("mult3", 3)

        x = {"value": 0}
        pipeline = SequentialPipeline([add1, add2, mult3])
        self.assertEqual(pipeline(x)["value"], 9)
        x = {"value": 0}
        pipeline = SequentialPipeline([add1, mult3, add2])
        self.assertEqual(pipeline(x)["value"], 5)
        x = {"value": 0}
        pipeline = SequentialPipeline([mult3, add1, add2])
        self.assertEqual(pipeline(x)["value"], 3)

    def test_if_else_pipeline(self) -> None:
        """Test IfElsePipeline executes if or else agent based on condition"""
        if_x = {"text": "xxxx [PASS]"}
        else_x = {"text": "xxxx"}

        if_agent = If_agent("if_agent")
        else_agent = Else_agent("else_agent")

        p = IfElsePipeline(
            condition_func=lambda x: "[PASS]" in x["text"],
            if_body_operators=if_agent,
            else_body_operators=else_agent,
        )
        x = p(if_x)
        self.assertEqual(x["operation"], "A")
        x = p(else_x)
        self.assertEqual(x["operation"], "B")

    def test_switch_pipeline(self) -> None:
        """Test SwitchPipeline executes one of the case agents or the default
        agent
        """
        tool_types = ["A", "B", "C", "D"]

        case_agent = Case_agent(name="case_agent")

        default_agent = Default_agent("default_agent")

        case_agents = {k: case_agent for k in tool_types}

        p = SwitchPipeline(
            condition_func=lambda x: x["text"].strip(),
            case_operators=case_agents,
            default_operators=default_agent,
        )
        for tool in tool_types:
            x = {"text": f"\n\n{tool}\n\n"}
            x = p(x)
            self.assertEqual(x["operation"], tool)
        x = p({"text": "hello"})
        self.assertEqual(x["operation"], "IDLE")

    def test_for_pipeline(self) -> None:
        """Test ForLoopPipeline"""

        loop_agent = Loop_for_agent("loop_agent")

        # test max loop
        x = {"value": 0}
        p = ForLoopPipeline(loop_body_operators=loop_agent, max_loop=10)
        x = p(x)
        self.assertEqual(x["value"], 10)

        x = {"value": 0}
        p = ForLoopPipeline(
            loop_body_operators=loop_agent,
            max_loop=10,
            break_func=lambda x: x["value"] > 5,
        )
        x = p(x)
        self.assertEqual(x["value"], 6)

    def test_while_pipeline(self) -> None:
        """Test WhileLoopPipeline"""

        loop_agent = Loop_while_agent("loop_agent")

        p = WhileLoopPipeline(
            loop_body_operators=loop_agent,
            condition_func=lambda i, x: i < 10 and not x["token_num"] > 500,
        )
        for _ in range(50):
            x = {"token_num": 0, "round": 0}
            x = p(x)
            self.assertTrue(x["round"] >= 10 or x["token_num"] > 500)


class FunctionalPipelineTest(unittest.TestCase):
    """Test cases for Functional Pipelines"""

    def test_sequential_pipeline(self) -> None:
        """Test SequentialPipeline executes agents sequentially"""

        add1 = Add("add1", 1)
        add2 = Add("add2", 2)
        mult3 = Mult("mult3", 3)

        x = {"value": 0}
        x = sequentialpipeline(x=x, operators=[add1, add2, mult3])
        self.assertEqual(x["value"], 9)
        x = {"value": 0}
        x = sequentialpipeline(x=x, operators=[add1, mult3, add2])
        self.assertEqual(x["value"], 5)
        x = {"value": 0}
        x = sequentialpipeline(x=x, operators=[mult3, add1, add2])
        self.assertEqual(x["value"], 3)

    def test_if_else_pipeline(self) -> None:
        """Test ifelsepipeline executes if or else agent based on condition"""
        if_x = {"text": "xxxx [PASS]"}
        else_x = {"text": "xxxx"}

        if_agent = If_agent("if_agent")
        else_agent = Else_agent("else_agent")

        if_x = ifelsepipeline(
            x=if_x,
            condition_func=lambda x: "[PASS]" in x["text"],
            if_body_operators=if_agent,
            else_body_operators=else_agent,
        )
        else_x = ifelsepipeline(
            x=else_x,
            condition_func=lambda x: "[PASS]" in x["text"],
            if_body_operators=if_agent,
            else_body_operators=else_agent,
        )
        self.assertEqual(if_x["operation"], "A")
        self.assertEqual(else_x["operation"], "B")


if __name__ == "__main__":
    unittest.main()