from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

from PythonModels.components.amesim_pneumatic import (
    AmesimPneumaticOrifice,
    AmesimPneumaticVolume,
)
from PythonModels.core.medium import ThermodynamicProperties
from PythonModels.core.state import VolumeState


@dataclass(frozen=True)
class TestMqlPneumaticBranchSpec:
    name: str
    upstream_volume_alias: str
    orifice_alias: str
    downstream_volume_alias: str
    source: str = "manual"


@dataclass(frozen=True)
class TestMqlPneumaticBranchComponents:
    name: str
    upstream_volume: AmesimPneumaticVolume
    orifice: AmesimPneumaticOrifice
    downstream_volume: AmesimPneumaticVolume
    spec: TestMqlPneumaticBranchSpec | None = None


@dataclass(frozen=True)
class TestMqlPneumaticBranchState:
    name: str
    upstream: ThermodynamicProperties
    downstream: ThermodynamicProperties
    flow: float
    upstream_inlet_h: float
    downstream_inlet_h: float


@dataclass(frozen=True)
class TestMqlPneumaticSnapshot:
    branch: TestMqlPneumaticBranchState

    @property
    def flow(self) -> float:
        return self.branch.flow

    @property
    def upstream(self) -> ThermodynamicProperties:
        return self.branch.upstream

    @property
    def downstream(self) -> ThermodynamicProperties:
        return self.branch.downstream


class TestMqlPneumaticClosure:
    """Minimal test_mql pneumatic closure following the existing Testmodel pattern.

    The canonical branch flow is positive from ``upstream_volume`` through the
    orifice port_a/port_b into ``downstream_volume``. Port ``m_flow`` values are
    written with Modelica-style signs: positive means flow into that component.
    """

    def __init__(
        self,
        *,
        components: TestMqlPneumaticBranchComponents,
        initial_state_vector: Callable[[], list[float]],
        apply_state_vector: Callable[[list[float]], None],
    ) -> None:
        self.components = components
        self._initial_state_vector = initial_state_vector
        self._apply_state_vector = apply_state_vector

    def initial_state_vector(self) -> list[float]:
        return self._initial_state_vector()

    def apply_state_vector(self, values: list[float]) -> None:
        self._apply_state_vector(values)

    def snapshot(
        self,
        state_vector: list[float] | None = None,
    ) -> TestMqlPneumaticSnapshot:
        if state_vector is not None:
            self._apply_state_vector(state_vector)
        upstream = self.components.upstream_volume.properties()
        downstream = self.components.downstream_volume.properties()
        flow = self._solve_branch_flow(upstream, downstream)
        upstream_inlet_h = self.components.upstream_volume.connection_inlet_enthalpy(
            port_m_flow=-flow,
            connected_h=downstream.h,
            internal_h=upstream.h,
        )
        downstream_inlet_h = self.components.downstream_volume.connection_inlet_enthalpy(
            port_m_flow=flow,
            connected_h=upstream.h,
            internal_h=downstream.h,
        )
        branch = TestMqlPneumaticBranchState(
            name=self.components.name,
            upstream=upstream,
            downstream=downstream,
            flow=flow,
            upstream_inlet_h=upstream_inlet_h,
            downstream_inlet_h=downstream_inlet_h,
        )
        self._write_port_states(branch)
        return TestMqlPneumaticSnapshot(branch=branch)

    def _solve_branch_flow(
        self,
        upstream: ThermodynamicProperties,
        downstream: ThermodynamicProperties,
    ) -> float:
        upstream_temperature = upstream.T if upstream.p >= downstream.p else downstream.T
        return self.components.orifice.mass_flow(
            upstream.p,
            downstream.p,
            upstream_temperature,
        )

    def _write_port_states(self, branch: TestMqlPneumaticBranchState) -> None:
        upstream_volume = self.components.upstream_volume
        downstream_volume = self.components.downstream_volume
        orifice = self.components.orifice

        upstream_volume.port_a.p = branch.upstream.p
        upstream_volume.port_a.m_flow = -branch.flow
        upstream_volume.port_a.h_outflow = branch.upstream.h

        orifice.port_a.p = branch.upstream.p
        orifice.port_a.m_flow = branch.flow
        orifice.port_a.h_outflow = branch.upstream.h
        orifice.port_b.p = branch.downstream.p
        orifice.port_b.m_flow = -branch.flow
        orifice.port_b.h_outflow = branch.downstream.h

        downstream_volume.port_a.p = branch.downstream.p
        downstream_volume.port_a.m_flow = branch.flow
        downstream_volume.port_a.h_outflow = branch.downstream.h

    def branch_derivatives(
        self,
        snapshot: TestMqlPneumaticSnapshot,
    ) -> tuple[VolumeState, VolumeState]:
        flow = snapshot.flow
        upstream_derivative = self.components.upstream_volume.derivatives(
            inlet_h=snapshot.branch.upstream_inlet_h,
            m_flow=-flow,
        )
        downstream_derivative = self.components.downstream_volume.derivatives(
            inlet_h=snapshot.branch.downstream_inlet_h,
            m_flow=flow,
        )
        return upstream_derivative, downstream_derivative

    def rhs(self, state_vector: list[float]) -> list[float]:
        snapshot = self.snapshot(state_vector)
        upstream_derivative, downstream_derivative = self.branch_derivatives(snapshot)
        return [
            *upstream_derivative.as_vector(),
            *downstream_derivative.as_vector(),
        ]