from __future__ import annotations

import unittest

from app.main import (
    ReactFlowProjectPayload,
    build_reactflow_system_xml,
    compile_reactflow_network,
    run_system_xml_simulation,
)
from app.system_xml import validate_system_xml_document
from tests.test_generic_system_xml_simulation import component_node, physical_edge
from tests.test_system_xml_protocol import physical_port


def amesim_pn3node_project() -> ReactFlowProjectPayload:
    return ReactFlowProjectPayload(
        projectSchemaVersion=1,
        name="amesim-pn3node-smoke",
        nodes=[
            component_node(
                "cylinder_1",
                "cylinder",
                [physical_port("port_b", "outlet", "right")],
                {"volume": 0.01, "p0": 500000.0, "T0": 300.0},
            ),
            component_node(
                "orifice_1",
                "orifice",
                [
                    physical_port("port_a", "inlet", "left"),
                    physical_port("port_b", "outlet", "right"),
                ],
                {"K": 1e-5, "opening": 1.0},
            ),
            component_node(
                "node_1",
                "amesim_pn3node2",
                [
                    physical_port("port_1", "bidirectional", "left"),
                    physical_port("port_2", "bidirectional", "right"),
                    physical_port("port_3", "bidirectional", "right"),
                ],
            ),
            component_node(
                "tank_1",
                "tank",
                [physical_port("port_a", "inlet", "left")],
                {"volume": 0.1, "p0": 100000.0, "T0": 300.0},
            ),
            component_node(
                "pipe_2",
                "pipe",
                [
                    physical_port("port_a", "inlet", "left"),
                    physical_port("port_b", "outlet", "right"),
                ],
                {
                    "length": 1.0,
                    "diameter": 0.02,
                    "lambda_darcy": 0.02,
                    "p0": 100000.0,
                    "T0": 300.0,
                },
            ),
            component_node(
                "tank_2",
                "tank",
                [physical_port("port_a", "inlet", "left")],
                {"volume": 0.1, "p0": 100000.0, "T0": 300.0},
            ),
        ],
        edges=[
            physical_edge("edge-1", "cylinder_1", "port_b", "orifice_1", "port_a"),
            physical_edge("edge-2", "orifice_1", "port_b", "node_1", "port_1"),
            # PN3NODE2 takes its reference temperature/pressure from storage.
            physical_edge("edge-3", "node_1", "port_2", "tank_1", "port_a"),
            physical_edge("edge-5", "node_1", "port_3", "pipe_2", "port_a"),
            physical_edge("edge-6", "pipe_2", "port_b", "tank_2", "port_a"),
        ],
        simulation={
            "t_start": 0.0,
            "t_stop": 0.002,
            "step": 0.001,
            "max_step": 0.001,
            "method": "BDF",
        },
    )


class AmesimPneumaticNodeXmlTests(unittest.TestCase):
    def test_pn3node_reactflow_project_compiles(self) -> None:
        network = compile_reactflow_network(amesim_pn3node_project())

        self.assertIn("node_1", network.components)
        self.assertEqual(network.components["node_1"].model_type, "amesim_pn3node2")
        self.assertTrue(network.pressure_flow_structure_dict()["isSquare"])

    def test_pn3node_system_xml_validates_and_simulates(self) -> None:
        xml = build_reactflow_system_xml(amesim_pn3node_project())
        report = validate_system_xml_document(xml)

        self.assertTrue(report.valid, report.as_dict())

        result = run_system_xml_simulation(xml)

        self.assertTrue(result["success"], result["message"])
        self.assertEqual(result["series"]["time"], [0.0, 0.001, 0.002])
        self.assertIn("node_1.port_2.p", result["series"])


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