from __future__ import annotations

import asyncio
import json
import os
import threading
import time
import unittest
from unittest.mock import patch
from uuid import uuid4

from fastapi import HTTPException, Request

from app.main import (
    ReactFlowEdgePayload,
    ReactFlowProjectPayload,
    SimulationCancellationPayload,
    _mark_simulation_task_result,
    _register_simulation_task,
    build_reactflow_system_xml,
    cancel_system_xml_simulation,
    compile_reactflow_network,
    get_system_xml_simulation,
    simulation_event_stream,
    simulate_system_xml,
)
from app.simulation.components.experimental.flow.resistive_pipe import ResistivePipe
from app.simulation.registry import get_component_model_spec
from app.simulation.config import (
    SolveIVPConfig,
    SolverActivityTracker,
)
from app.simulation.native_codegen.runner import simulate_native
from tests.test_system_xml_protocol import physical_port


def component_node(
    component_id: str,
    model_type: str,
    ports: list[dict[str, str]],
    parameters: dict[str, float] | None = None,
) -> dict[str, object]:
    return {
        "id": component_id,
        "type": "simulationComponent",
        "position": {"x": 0.0, "y": 0.0},
        "data": {
            "label": component_id,
            "componentType": model_type,
            "modelType": model_type,
            "modelVersion": get_component_model_spec(model_type).model_version,
            "ports": ports,
            "parameters": parameters or {},
        },
    }


def physical_edge(
    edge_id: str,
    first_component: str,
    first_port: str,
    second_component: str,
    second_port: str,
) -> dict[str, str]:
    return {
        "id": edge_id,
        "source": first_component,
        "sourceHandle": first_port,
        "target": second_component,
        "targetHandle": second_port,
    }


def two_port_definitions() -> list[dict[str, str]]:
    return [
        physical_port("port_a", "inlet", "left"),
        physical_port("port_b", "outlet", "right"),
    ]


def tee_port_definitions() -> list[dict[str, str]]:
    return [
        physical_port("port_in", "bidirectional", "left"),
        physical_port("port_out1", "bidirectional", "right"),
        physical_port("port_out2", "bidirectional", "right"),
    ]


def chain_project(*, reverse_edges: bool = False) -> ReactFlowProjectPayload:
    edges = [
        physical_edge("edge-1", "cylinder_1", "port_b", "orifice_1", "port_a"),
        physical_edge("edge-2", "orifice_1", "port_b", "pipe_1", "port_a"),
        physical_edge("edge-3", "pipe_1", "port_b", "tank_1", "port_a"),
    ]
    if reverse_edges:
        edges = [
            physical_edge(
                edge["id"],
                edge["target"],
                edge["targetHandle"],
                edge["source"],
                edge["sourceHandle"],
            )
            for edge in edges
        ]
    return ReactFlowProjectPayload(
        projectSchemaVersion=1,
        name="generic-chain",
        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",
                two_port_definitions(),
                {"K": 1e-5, "opening": 1.0},
            ),
            component_node(
                "pipe_1",
                "pipe",
                two_port_definitions(),
                {
                    "length": 1.0,
                    "diameter": 0.02,
                    "lambda_darcy": 0.02,
                    "p0": 100000.0,
                    "T0": 300.0,
                },
            ),
            component_node(
                "tank_1",
                "tank",
                [physical_port("port_a", "inlet", "left")],
                {"volume": 0.1, "p0": 100000.0, "T0": 300.0},
            ),
        ],
        edges=edges,
        simulation={
            "t_start": 0.0,
            "t_stop": 0.01,
            "step": 0.005,
            "max_step": 0.001,
            "method": "BDF",
        },
    )


def branched_project() -> ReactFlowProjectPayload:
    return ReactFlowProjectPayload(
        projectSchemaVersion=1,
        name="generic-branch",
        nodes=[
            component_node(
                "cylinder_1",
                "cylinder",
                [physical_port("port_b", "outlet", "right")],
                {"volume": 0.01, "p0": 500000.0, "T0": 300.0},
            ),
            component_node("tee_1", "tee", tee_port_definitions()),
            component_node(
                "orifice_1",
                "orifice",
                two_port_definitions(),
                {"K": 1e-5, "opening": 1.0},
            ),
            component_node(
                "pipe_1",
                "pipe",
                two_port_definitions(),
                {
                    "length": 1.0,
                    "diameter": 0.02,
                    "lambda_darcy": 0.02,
                    "p0": 100000.0,
                    "T0": 300.0,
                },
            ),
            component_node(
                "orifice_2",
                "orifice",
                two_port_definitions(),
                {"K": 1e-5, "opening": 1.0},
            ),
            component_node(
                "pipe_2",
                "pipe",
                two_port_definitions(),
                {
                    "length": 1.0,
                    "diameter": 0.02,
                    "lambda_darcy": 0.02,
                    "p0": 100000.0,
                    "T0": 300.0,
                },
            ),
            component_node("tee_2", "tee", tee_port_definitions()),
            component_node(
                "tank_1",
                "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", "tee_1", "port_in"),
            physical_edge("edge-2", "tee_1", "port_out1", "orifice_1", "port_a"),
            physical_edge("edge-3", "orifice_1", "port_b", "pipe_1", "port_a"),
            physical_edge("edge-4", "pipe_1", "port_b", "tee_2", "port_out1"),
            physical_edge("edge-5", "tee_1", "port_out2", "orifice_2", "port_a"),
            physical_edge("edge-6", "orifice_2", "port_b", "pipe_2", "port_a"),
            physical_edge("edge-7", "pipe_2", "port_b", "tee_2", "port_out2"),
            physical_edge("edge-8", "tee_2", "port_in", "tank_1", "port_a"),
        ],
        simulation={
            "t_start": 0.0,
            "t_stop": 0.005,
            "step": 0.005,
            "max_step": 0.001,
            "method": "BDF",
        },
    )


def xml_request(body: bytes) -> Request:
    delivered = False

    async def receive():
        nonlocal delivered
        if delivered:
            return {"type": "http.disconnect"}
        delivered = True
        return {"type": "http.request", "body": body, "more_body": False}

    return Request(
        {
            "type": "http",
            "method": "POST",
            "path": "/api/system-xml/simulate",
            "headers": [(b"content-type", b"application/xml")],
        },
        receive,
    )


class GenericSystemXmlSimulationTests(unittest.TestCase):
    def test_native_chain_conserves_mass_energy_and_edge_order(self):
        results = []
        for reverse in (False, True):
            result = simulate_native(compile_reactflow_network(chain_project(reverse_edges=reverse)), SolveIVPConfig(t_stop=.01, max_step=.001), sample_step=.005)
            self.assertTrue(result.success)
            self.assertEqual(result.diagnostics['backend'], 'native-c')
            self.assertLess(result.final['cylinder_1.p'], 500000)
            self.assertGreater(result.final['tank_1.p'], 100000)
            for field, tolerance in (('m', 1e-12), ('U', 1e-6)):
                totals = [a+b for a,b in zip(result.series['cylinder_1.'+field], result.series['tank_1.'+field])]
                self.assertLess(max(totals)-min(totals), tolerance)
            results.append(result)
        self.assertAlmostEqual(results[0].final['tank_1.p'], results[1].final['tank_1.p'], places=7)

    def test_native_branched_topology_conserves_flow(self):
        result = simulate_native(compile_reactflow_network(branched_project()), SolveIVPConfig(t_stop=.005, max_step=.001), sample_step=.005)
        self.assertTrue(result.success)
        self.assertAlmostEqual(result.final['pipe_1.port_a.m_flow'], result.final['pipe_2.port_a.m_flow'], places=10)
        self.assertGreater(result.final['tank_1.p'], 100000)

    def test_xml_pipe_compiles_to_quasi_steady_resistance(self) -> None:
        network = compile_reactflow_network(chain_project())

        self.assertIsInstance(network.components["pipe_1"], ResistivePipe)
        structure = network.pressure_flow_structure_dict()
        self.assertEqual(structure["unknownCount"], 12)
        self.assertEqual(structure["equationCount"], 12)
        self.assertTrue(structure["isSquare"])




    def test_task_registry_distinguishes_user_stop_and_stalled_stop(self) -> None:
        for reason, expected_status in (("user", "stopped"), ("stalled", "stalled")):
            simulation_id = f"test-{uuid4().hex}"
            task = _register_simulation_task(simulation_id)
            response = cancel_system_xml_simulation(
                simulation_id,
                SimulationCancellationPayload(reason=reason),
            )
            self.assertTrue(response["accepted"])

            result = _mark_simulation_task_result(
                task,
                {
                    "success": False,
                    "status": "cancelled",
                    "partial": True,
                    "series": {"time": [0.0, 0.1]},
                },
            )
            snapshot = get_system_xml_simulation(simulation_id)

            self.assertEqual(result["status"], expected_status)
            self.assertEqual(snapshot["status"], expected_status)
            self.assertEqual(snapshot["cancelReason"], reason)
            self.assertEqual(snapshot["result"]["series"]["time"], [0.0, 0.1])




    def test_raw_system_xml_runs_through_generic_simulation_endpoint(self) -> None:
        xml = build_reactflow_system_xml(chain_project())

        response = asyncio.run(simulate_system_xml(xml_request(xml)))

        self.assertTrue(response["success"])
        self.assertEqual(response["simulation"]["method"], "BDF")
        self.assertEqual(response["model"]["pressureFlowSystem"]["unknownCount"], 12)
        self.assertEqual(response["diagnostics"]["stateCount"], 4)
        self.assertGreater(response["final"]["tank_1.p"], 100000.0)
        variables = {variable["key"]: variable for variable in response["variables"]}
        self.assertEqual(set(variables), set(response["series"]) - {"time"})
        self.assertEqual(variables["cylinder_1.p"]["componentId"], "cylinder_1")
        self.assertEqual(variables["cylinder_1.p"]["label"], "压力")
        self.assertEqual(variables["cylinder_1.p"]["unit"], "Pa")
        self.assertEqual(variables["pipe_1.port_a.m_flow"]["scope"], "port")
        self.assertEqual(
            variables["pipe_1.port_a.m_flow"]["portName"],
            "port_a",
        )

    def test_raw_system_xml_serializes_more_than_legacy_sample_limit(self) -> None:
        project = chain_project()
        project.simulation.step = 0.0000005
        xml = build_reactflow_system_xml(project)

        response = asyncio.run(simulate_system_xml(xml_request(xml)))
        round_tripped = json.loads(json.dumps(response))

        self.assertTrue(round_tripped["success"])
        self.assertEqual(round_tripped["diagnostics"]["sampleCount"], 20001)
        self.assertEqual(len(round_tripped["series"]["time"]), 20001)
        self.assertEqual(round_tripped["series"]["time"][0], 0.0)
        self.assertEqual(round_tripped["series"]["time"][-1], 0.01)

    def test_streaming_endpoint_events_have_monotonic_progress_and_result(self) -> None:
        project = chain_project()
        xml = build_reactflow_system_xml(project)

        events = [json.loads(line) for line in simulation_event_stream(xml)]

        progress_events = [
            event for event in events if event["event"] == "progress"
        ]
        progress = [event["progress"] for event in progress_events]
        self.assertGreater(len(progress), 3)
        self.assertTrue(
            all(current <= following for current, following in zip(progress, progress[1:]))
        )
        preparation_events = [
            event
            for event in progress_events
            if event["phase"] in {"validation", "compilation", "initialization"}
        ]
        self.assertTrue(preparation_events)
        self.assertTrue(all(event["progress"] == 0 for event in preparation_events))
        timed_events = [
            event
            for event in progress_events
            if event.get("simulatedTime") is not None
        ]
        self.assertGreater(len(timed_events), 1)
        for event in timed_events:
            self.assertEqual(event["totalTime"], project.simulation.t_stop)
            expected_progress = round(
                100
                * (event["simulatedTime"] - project.simulation.t_start)
                / (project.simulation.t_stop - project.simulation.t_start)
            )
            self.assertEqual(event["progress"], expected_progress)
        self.assertEqual(events[-1]["event"], "result")
        self.assertEqual(events[-1]["progress"], 100)
        self.assertTrue(events[-1]["result"]["success"])

    def test_streaming_task_stop_returns_partial_result(self) -> None:
        project = branched_project()
        project.simulation.t_stop = 2.0
        project.simulation.step = 0.02
        project.simulation.max_step = 0.000001
        xml = build_reactflow_system_xml(project)
        simulation_id = f"test-{uuid4().hex}"
        task = _register_simulation_task(simulation_id)
        events: list[dict[str, object]] = []
        stop_requested = False

        for line in simulation_event_stream(xml, task=task):
            event = json.loads(line)
            events.append(event)
            if (
                not stop_requested
                and event["event"] == "progress"
                and event["phase"] == "integrating"
                and event["progress"] >= 20
            ):
                response = cancel_system_xml_simulation(
                    simulation_id,
                    SimulationCancellationPayload(reason="user"),
                )
                self.assertTrue(response["accepted"])
                stop_requested = True

        result_event = next(event for event in events if event["event"] == "result")
        result = result_event["result"]
        snapshot = get_system_xml_simulation(simulation_id)

        self.assertTrue(stop_requested)
        self.assertEqual(result["status"], "stopped")
        self.assertTrue(result["partial"])
        self.assertGreaterEqual(result["diagnostics"]["sampleCount"], 2)
        self.assertGreater(result["simulatedUntil"], 0.0)
        self.assertLess(result["simulatedUntil"], 2.0)
        self.assertEqual(snapshot["status"], "stopped")
        self.assertEqual(snapshot["result"]["status"], "stopped")

    def test_streaming_endpoint_keeps_quiet_solver_connection_alive(self) -> None:
        heartbeat_observed = threading.Event()

        def delayed_simulation(
            _xml_bytes,
            progress_callback,
            _cancel_check=None,
            activity_tracker=None,
        ):
            progress_callback(49, "integrating", "正在进行时间积分与压力流量求解")
            self.assertIsNotNone(activity_tracker)
            activity_tracker.start_integration(0.0)
            for trial_time in (0.0487, 0.0488, 0.0489, 0.0490):
                heartbeat_observed.clear()
                activity_tracker.record_rhs(trial_time)
                # Wait for the consumer, independent of machine scheduling.
                self.assertTrue(heartbeat_observed.wait(timeout=2.0))
            return {
                "success": True,
                "status": "completed",
                "simulatedUntil": 0.2,
                "requestedStopTime": 0.2,
            }

        with (
            patch("app.main.SIMULATION_STREAM_HEARTBEAT_SECONDS", 0.005),
            patch(
                "app.main.run_system_xml_simulation",
                side_effect=delayed_simulation,
            ),
        ):
            events = []
            for line in simulation_event_stream(b"<System />"):
                event = json.loads(line)
                events.append(event)
                if event.get("heartbeat"):
                    heartbeat_observed.set()

        heartbeats = [event for event in events if event.get("heartbeat") is True]
        self.assertGreaterEqual(len(heartbeats), 1)
        self.assertTrue(all(event["progress"] == 49 for event in heartbeats))
        self.assertTrue(all(event["phase"] == "integrating" for event in heartbeats))
        self.assertTrue(
            all(event["acceptedStepSequence"] == 0 for event in heartbeats)
        )
        self.assertTrue(all(event["acceptedTime"] == 0.0 for event in heartbeats))
        self.assertGreater(
            len({event["activitySequence"] for event in heartbeats}),
            1,
        )
        self.assertTrue(
            all(event["activityKind"] == "rhs" for event in heartbeats)
        )
        self.assertEqual(events[-1]["event"], "result")
        self.assertEqual(events[-1]["activityKind"], "complete")
        self.assertEqual(events[-1]["acceptedStepSequence"], 0)
        self.assertEqual(events[-1]["rhsCallCount"], 4)

    def test_streaming_endpoint_returns_structured_validation_error(self) -> None:
        events = [
            json.loads(line)
            for line in simulation_event_stream(b"<System schemaVersion='3'>")
        ]

        self.assertEqual(events[-1]["event"], "error")
        self.assertEqual(events[-1]["status"], 422)
        self.assertIn("issues", events[-1]["detail"])

    def test_simulation_endpoint_returns_422_for_ideal_storage_coupling(self) -> None:
        project = chain_project()
        project.nodes = [project.nodes[0], project.nodes[-1]]
        project.edges = [
            ReactFlowEdgePayload(
                **physical_edge(
                    "edge-1",
                    "cylinder_1",
                    "port_b",
                    "tank_1",
                    "port_a",
                )
            )
        ]

        with self.assertRaises(HTTPException) as caught:
            asyncio.run(
                simulate_system_xml(
                    xml_request(build_reactflow_system_xml(project))
                )
            )

        self.assertEqual(caught.exception.status_code, 422)
        self.assertIn(
            "SIMULATION_EXECUTION_FAILED",
            {
                issue["code"]
                for issue in caught.exception.detail["issues"]
            },
        )


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