from __future__ import annotations

import unittest
from xml.etree import ElementTree as ET

from app.main import ReactFlowProjectPayload, build_reactflow_system_xml


def physical_port(
    name: str,
    role: str,
    side: str,
    *,
    domain: str = "pneumatic",
) -> dict[str, str]:
    return {
        "name": name,
        "kind": "physical",
        "domain": domain,
        "nominalRole": role,
        "positiveFlowDirection": "intoComponent",
        "side": side,
    }


def physical_connection_project(*, reverse_edge: bool = False) -> ReactFlowProjectPayload:
    source = {
        "id": "cylinder_1",
        "type": "simulationComponent",
        "position": {"x": 12.5, "y": 24.0},
        "data": {
            "label": "source-cylinder",
            "componentType": "cylinder",
            "modelType": "cylinder",
            "ports": [physical_port("port_b", "outlet", "right")],
            "parameters": {"volume": 0.01, "p0": 35000000, "T0": 300},
        },
    }
    target = {
        "id": "tank_1",
        "type": "simulationComponent",
        "position": {"x": 420.0, "y": 24.0},
        "data": {
            "label": "receiver-tank",
            "componentType": "tank",
            "modelType": "tank",
            "ports": [physical_port("port_a", "inlet", "left")],
            "parameters": {"volume": 0.1, "p0": 100000, "T0": 300},
        },
    }
    edge = (
        {
            "id": "edge-1",
            "source": "tank_1",
            "sourceHandle": "port_a",
            "target": "cylinder_1",
            "targetHandle": "port_b",
        }
        if reverse_edge
        else {
            "id": "edge-1",
            "source": "cylinder_1",
            "sourceHandle": "port_b",
            "target": "tank_1",
            "targetHandle": "port_a",
        }
    )
    return ReactFlowProjectPayload(
        name="protocol-test",
        nodes=[source, target],
        edges=[edge],
        simulation={
            "t_start": 1.0,
            "t_stop": 5.0,
            "step": 0.2,
            "max_step": 0.01,
            "method": "BDF",
        },
    )


class SystemXmlProtocolTests(unittest.TestCase):
    def test_v2_xml_contains_port_metadata_and_neutral_physical_endpoints(self) -> None:
        root = ET.fromstring(build_reactflow_system_xml(physical_connection_project()))

        self.assertEqual(root.tag, "System")
        self.assertEqual(root.attrib["name"], "protocol-test")
        self.assertEqual(root.attrib["schemaVersion"], "2")
        self.assertEqual(root.attrib["unitSystem"], "SI")
        self.assertEqual(root.attrib["mediumReferenceVersion"], "1")
        self.assertEqual(
            [child.tag for child in root],
            ["Simulation", "Components", "Connections"],
        )

        simulation = root.find("Simulation")
        assert simulation is not None
        self.assertEqual(
            simulation.attrib,
            {
                "tStart": "1.0",
                "tStop": "5.0",
                "step": "0.2",
                "maxStep": "0.01",
                "method": "BDF",
            },
        )

        port = root.find("./Components/Component[@id='cylinder_1']/Port")
        assert port is not None
        component = root.find("./Components/Component[@id='cylinder_1']")
        assert component is not None
        self.assertEqual(component.attrib["rotation"], "0")
        self.assertEqual(component.attrib["mirrored"], "false")
        self.assertEqual(
            port.attrib,
            {
                "name": "port_b",
                "kind": "physical",
                "domain": "pneumatic",
                "nominalRole": "outlet",
                "side": "right",
                "positiveFlowDirection": "intoComponent",
            },
        )

        connection = root.find("./Connections/Connection")
        assert connection is not None
        self.assertEqual(
            connection.attrib,
            {"id": "edge-1", "kind": "physical", "domain": "pneumatic"},
        )
        self.assertNotIn("source", connection.attrib)
        self.assertEqual(
            {
                (endpoint.attrib["component"], endpoint.attrib["port"])
                for endpoint in connection.findall("Endpoint")
            },
            {("cylinder_1", "port_b"), ("tank_1", "port_a")},
        )
        self.assertTrue(
            all("role" not in endpoint.attrib for endpoint in connection.findall("Endpoint"))
        )

    def test_physical_connection_meaning_does_not_depend_on_drag_direction(self) -> None:
        forward = ET.fromstring(
            build_reactflow_system_xml(physical_connection_project())
        )
        reverse = ET.fromstring(
            build_reactflow_system_xml(physical_connection_project(reverse_edge=True))
        )

        def endpoint_set(root: ET.Element) -> set[tuple[str, str]]:
            return {
                (endpoint.attrib["component"], endpoint.attrib["port"])
                for endpoint in root.findall("./Connections/Connection/Endpoint")
            }

        self.assertEqual(endpoint_set(forward), endpoint_set(reverse))

    def test_component_orientation_is_exported_as_layout_metadata(self) -> None:
        project = physical_connection_project()
        project.nodes[0].data.rotation = 90
        project.nodes[0].data.mirrored = True

        root = ET.fromstring(build_reactflow_system_xml(project))
        component = root.find("./Components/Component[@id='cylinder_1']")
        assert component is not None
        self.assertEqual(component.attrib["rotation"], "90")
        self.assertEqual(component.attrib["mirrored"], "true")

    def test_legacy_string_port_is_migrated_when_exporting_v2(self) -> None:
        project = ReactFlowProjectPayload(
            name="legacy-port",
            nodes=[
                {
                    "id": "tank_1",
                    "data": {
                        "label": "tank_1",
                        "componentType": "tank",
                        "modelType": "tank",
                        "ports": ["port_a"],
                    },
                }
            ],
        )

        root = ET.fromstring(build_reactflow_system_xml(project))
        port = root.find("./Components/Component/Port")
        assert port is not None
        self.assertEqual(port.attrib["kind"], "physical")
        self.assertEqual(port.attrib["domain"], "pneumatic")
        self.assertEqual(port.attrib["nominalRole"], "inlet")
        self.assertEqual(port.attrib["positiveFlowDirection"], "intoComponent")
        self.assertEqual(port.attrib["side"], "left")

    def test_legacy_tee_port_names_do_not_restore_direction_constraints(self) -> None:
        project = ReactFlowProjectPayload(
            name="legacy-tee",
            nodes=[
                {
                    "id": "tee_1",
                    "data": {
                        "label": "tee_1",
                        "componentType": "tee",
                        "modelType": "tee",
                        "ports": ["port_in", "port_out1", "port_out2"],
                    },
                }
            ],
        )

        root = ET.fromstring(build_reactflow_system_xml(project))
        ports = root.findall("./Components/Component/Port")
        self.assertEqual(
            [port.attrib["nominalRole"] for port in ports],
            ["bidirectional", "bidirectional", "bidirectional"],
        )

    def test_incompatible_physical_domains_are_rejected(self) -> None:
        project = physical_connection_project()
        tank_port = project.nodes[1].data.ports[0]
        assert not isinstance(tank_port, str)
        tank_port.domain = "hydraulic"

        with self.assertRaisesRegex(ValueError, "incompatible domains"):
            build_reactflow_system_xml(project)


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