import unittest

from .block import BaseBlock, InputBlock, OutputBlock
from .graph import Edge, Graph
from .rule import (
    EndpointNotExistError,
    InputOutputCountError,
    NotDAGError,
    PortMismatchError,
    PortNotConnectedError,
    TypeNoStrMethodError,
)


class MockInput(InputBlock):
    def __init__(self, mock_output: str = ""):
        super(MockInput, self).__init__()
        self.mock_output = mock_output

    def input(self, messages):
        pass

    def __call__(self) -> str:
        return self.mock_output


class MockOutput(OutputBlock):
    def __init__(self, mock_output: str = ""):
        super(MockOutput, self).__init__()
        self.mock_output = mock_output

    def __call__(self, input) -> str:
        return self.mock_output


class MockNode(BaseBlock):
    def __init__(self, mock_output: str = ""):
        super(MockNode, self).__init__()
        self.mock_output = mock_output

    def __call__(self, input: str) -> str:
        return self.mock_output


class MockNode_int_str(BaseBlock):
    def __init__(self, mock_output: str = ""):
        super(MockNode_int_str, self).__init__()
        self.mock_output = mock_output

    def __call__(self, input: int) -> str:
        return self.mock_output


class MockNode_str_int(BaseBlock):
    def __init__(self, mock_output: int = 0):
        super(MockNode_str_int, self).__init__()
        self.mock_output = mock_output

    def __call__(self, input: str) -> int:
        return self.mock_output


class MockNode_any_str(BaseBlock):
    def __init__(self, mock_output: str = ""):
        super(MockNode_any_str, self).__init__()
        self.mock_output = mock_output

    def __call__(self, input) -> str:
        return self.mock_output


class MockTypeWithoutStrMethod:
    def __init__(self, messages: list):
        self.messages = messages

    def __iter__(self):
        return iter(self.messages)

    def __len__(self):
        return len(self.messages)


class MockNodeTypeWithoutStrMethod(BaseBlock):
    def __init__(self):
        super(MockNodeTypeWithoutStrMethod).__init__()

    def __call__(self, input: str) -> MockTypeWithoutStrMethod:
        return MockTypeWithoutStrMethod(["hello"])


class TestGraphRule(unittest.TestCase):
    def test_simple(self):
        nodes = {
            "node_1": MockInput(),
            "node_2": MockNode(),
            "output": MockOutput(),
        }
        edges = [
            Edge(
                source="node_1",
                sink="node_2",
                port="input",
                case=None,
            ),
            Edge(
                source="node_2",
                sink="output",
                port="input",
                case=None,
            ),
        ]
        g = Graph(nodes, edges)

    def test_dag(self):
        nodes = {
            "node_1": MockNode(),
            "node_2": MockNode(),
        }
        edges = [
            Edge(
                source="node_1",
                sink="node_2",
                port="input",
                case=None,
            ),
            Edge(
                source="node_2",
                sink="node_1",
                port="input",
                case=None,
            ),
        ]

        with self.assertRaises(NotDAGError):
            g = Graph(nodes, edges)

    def test_ports_match(self):
        nodes = {
            "node_1": MockInput(),
            "node_2": MockNode(),
            "output": MockOutput(),
        }
        edges = [
            Edge(
                source="node_1",
                sink="node_2",
                port="input",
                case=None,
            ),
            Edge(
                source="node_2",
                sink="output",
                port="input",
                case=None,
            ),
        ]
        g = Graph(nodes, edges)

        nodes = {
            "node_1": MockInput(),
            "node_2": MockNode_int_str(),
            "output": MockOutput(),
        }
        with self.assertRaises(PortMismatchError):
            g = Graph(nodes, edges)

    def test_match_any(self):
        nodes = {
            "node_1": MockInput(),
            "node_2": MockNode_str_int(),
            "node_3": MockNode_any_str(),
            "output": MockOutput(),
        }
        edges = [
            Edge(
                source="node_1",
                sink="node_2",
                port="input",
                case=None,
            ),
            Edge(
                source="node_2",
                sink="node_3",
                port="input",
                case=None,
            ),
            Edge(
                source="node_3",
                sink="output",
                port="input",
                case=None,
            ),
        ]
        g = Graph(nodes, edges)

        nodes = {
            "node_1": MockInput(),
            "node_2": MockNode(),
            "node_3": MockNode_any_str(),
            "output": MockOutput(),
        }
        g = Graph(nodes, edges)

    def test_endpoint_exist(self):
        nodes = {
            "node_1": MockInput(),
            "node_2": MockNode(),
            "output": MockOutput(),
        }
        edges = [
            Edge(
                source="node_1",
                sink="node_3",
                port="input",
                case=None,
            ),
            Edge(
                source="node_3",
                sink="output",
                port="input",
                case=None,
            ),
        ]
        with self.assertRaises(EndpointNotExistError):
            g = Graph(nodes, edges)

        edges = [
            Edge(
                source="node_1",
                sink="node_2",
                case=None,
                port="input",
            ),
            Edge(
                source="node_2",
                sink="output",
                case=None,
                port="input",
            ),
        ]
        g = Graph(nodes, edges)

        edges = [
            Edge(
                source="node_1",
                sink="node_2",
                case=None,
                port="not_exist_port",
            ),
            Edge(
                source="node_2",
                sink="output",
                case=None,
                port="input",
            ),
        ]
        with self.assertRaises(EndpointNotExistError):
            g = Graph(nodes, edges)

    def test_nodetype_has_str_method(self):
        nodes = {
            "node_1": MockInput(),
            "node_2": MockNodeTypeWithoutStrMethod(),
            "output": MockOutput(),
        }
        edges = [
            Edge(
                source="node_1",
                sink="node_2",
                port="input",
                case=None,
            ),
            Edge(
                source="node_2",
                sink="output",
                port="input",
                case=None,
            ),
        ]
        with self.assertRaises(TypeNoStrMethodError):
            g = Graph(nodes, edges)


class TestRunGraph(unittest.TestCase):
    def test_simple(self):
        nodes = {
            "input": MockInput("output of input"),
            "node_0": MockNode("output of node_0"),
            "output": MockOutput("mock output"),
        }
        edges = [
            Edge(
                source="input",
                sink="node_0",
                port="input",
                case=None,
            ),
            Edge(
                source="node_0",
                sink="output",
                port="input",
                case=None,
            ),
        ]
        g = Graph(nodes, edges)
        self.assertEqual(g.run("test"), "mock output")
        self.assertEqual(g.data["input"], "output of input")
        self.assertEqual(g.data["node_0"], "output of node_0")
        self.assertEqual(g.data["output"], "mock output")


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