from __future__ import annotations

from collections.abc import Mapping

from app.simulation.core.base import AlgebraicComponent
from app.simulation.core.catalog import ComponentDisplaySpec, PortDisplaySpec
from app.simulation.core.equations import EquationResidual
from app.simulation.core.medium import IdealGasMedium
from app.simulation.core.ports import PortDefinition, PortState


class _AmesimPneumaticNode(AlgebraicComponent):
    """Shared implementation for AMESim pneumatic junction submodels.

    PN3NODE2/P4NODE2 use port 2 as their pressure and temperature reference.
    Non-reference outlet ports use that reference temperature. When port 2 is
    an outlet, its enthalpy is the residual that closes the junction energy
    balance, matching the AMESim dh2 causality.
    """

    REFERENCE_PORT = "port_2"

    def __init__(self, name: str) -> None:
        super().__init__(name=name)
        self.set_parameter_values({})
        self.temperature_reference_h = 0.0
        for definition in self.PORTS:
            setattr(self, definition.name, self.register_declared_port(definition.name))

    def pressure_flow_equation_values(self) -> tuple[float, ...]:
        reference = self.get_port(self.REFERENCE_PORT)
        return tuple(
            self.get_port(definition.name).p - reference.p
            for definition in self.PORTS
            if definition.name != self.REFERENCE_PORT
        ) + (
            sum(
                self.get_port(definition.name).m_flow
                for definition in self.PORTS
            ),
        )

    def pressure_flow_equation_residuals(self) -> tuple[EquationResidual, ...]:
        reference = self.get_port(self.REFERENCE_PORT)
        residuals: list[EquationResidual] = []
        for definition in self.PORTS:
            if definition.name == self.REFERENCE_PORT:
                continue
            port = self.get_port(definition.name)
            residuals.append(
                EquationResidual(
                    id=f"{self.name}:{definition.name}_pressure_reference",
                    owner="component",
                    owner_id=self.name,
                    relation="equal",
                    variables=(
                        f"{self.name}.{definition.name}.p",
                        f"{self.name}.{self.REFERENCE_PORT}.p",
                    ),
                    role="effort",
                    value=port.p - reference.p,
                )
            )
        residuals.append(
            EquationResidual(
                id=f"{self.name}:mass_flow_balance",
                owner="component",
                owner_id=self.name,
                relation="sumToZero",
                variables=tuple(
                    f"{self.name}.{definition.name}.m_flow"
                    for definition in self.PORTS
                ),
                role="flow",
                value=sum(self.get_port(definition.name).m_flow for definition in self.PORTS),
            )
        )
        return tuple(residuals)

    def update_stream_outflows(self, connected_h: Mapping[str, float]) -> None:
        self.temperature_reference_h = connected_h.get(
            self.REFERENCE_PORT,
            sum(connected_h.values()) / len(connected_h) if connected_h else 0.0,
        )
        incoming = [
            (port.m_flow, connected_h[name])
            for name, port in self.ports.items()
            if port.m_flow > 1e-12
        ]
        total_flow = sum(m_flow for m_flow, _ in incoming)
        if total_flow > 1e-12:
            mixed_h = sum(m_flow * h for m_flow, h in incoming) / total_flow
        else:
            mixed_h = self.temperature_reference_h

        reference_port = self.get_port(self.REFERENCE_PORT)
        for name, port in self.ports.items():
            port.h_outflow = (
                mixed_h
                if name == self.REFERENCE_PORT
                else self.temperature_reference_h
            )

        if reference_port.m_flow < -1e-12:
            energy_without_reference = sum(
                port.m_flow
                * (
                    connected_h[name]
                    if port.m_flow > 1e-12
                    else self.temperature_reference_h
                )
                for name, port in self.ports.items()
                if name != self.REFERENCE_PORT
            )
            reference_port.h_outflow = (
                -energy_without_reference / reference_port.m_flow
            )


class AmesimPn3Node2(_AmesimPneumaticNode):
    """AMESim PN3NODE2 pneumatic three-port junction."""

    MODEL_TYPE = "amesim_pn3node2"
    MODEL_VERSION = "0.3.0"
    PORTS = (
        PortDefinition.pneumatic("port_1", nominal_role="bidirectional"),
        PortDefinition.pneumatic("port_2", nominal_role="bidirectional"),
        PortDefinition.pneumatic("port_3", nominal_role="bidirectional"),
    )
    PARAMETERS = ()
    RESULT_VARIABLES = ()
    DISPLAY = ComponentDisplaySpec(
        label="PN3NODE2 三端气动节点",
        library_id="amesim",
        category_id="junctions",
        symbol="amesim_pn3node2",
        ports=(
            PortDisplaySpec("port_1", "left", order=10),
            PortDisplaySpec("port_2", "right", order=20),
            PortDisplaySpec("port_3", "right", order=30),
        ),
        order=10,
    )

    @classmethod
    def create(
        cls,
        *,
        name: str,
        medium: IdealGasMedium,
        parameters: Mapping[str, float],
    ) -> AmesimPn3Node2:
        return cls(name=name)


class AmesimP4Node2(_AmesimPneumaticNode):
    """AMESim P4NODE2 pneumatic four-port junction."""

    MODEL_TYPE = "amesim_p4node2"
    MODEL_VERSION = "0.3.0"
    PORTS = (
        PortDefinition.pneumatic("port_1", nominal_role="bidirectional"),
        PortDefinition.pneumatic("port_2", nominal_role="bidirectional"),
        PortDefinition.pneumatic("port_3", nominal_role="bidirectional"),
        PortDefinition.pneumatic("port_4", nominal_role="bidirectional"),
    )
    PARAMETERS = ()
    RESULT_VARIABLES = ()
    DISPLAY = ComponentDisplaySpec(
        label="P4NODE2 四端气动节点",
        library_id="amesim",
        category_id="junctions",
        symbol="amesim_p4node2",
        ports=(
            PortDisplaySpec("port_1", "left", order=10),
            PortDisplaySpec("port_2", "right", order=20),
            PortDisplaySpec("port_3", "right", order=30),
            PortDisplaySpec("port_4", "right", order=40),
        ),
        order=20,
    )

    @classmethod
    def create(
        cls,
        *,
        name: str,
        medium: IdealGasMedium,
        parameters: Mapping[str, float],
    ) -> AmesimP4Node2:
        return cls(name=name)