from __future__ import annotations

from dataclasses import dataclass
from math import isfinite

from app.simulation.core.base import Component
from app.simulation.core.ports import PortState
from app.simulation.performance import profile_phase
from app.simulation.systems.network import Endpoint, SimulationNetwork


@dataclass(frozen=True)
class PneumaticVolumeDiagnostics:
    propagated: int
    output_ports: tuple[str, ...]

    def as_dict(self) -> dict[str, object]:
        return {
            "propagated": self.propagated,
            "outputPorts": list(self.output_ports),
        }


@dataclass(frozen=True)
class _PneumaticVolumeConnectionBinding:
    connected_endpoint: Endpoint
    connected_port: PortState


class PneumaticVolumeResolver:
    """Propagate AMESim pneumatic external-volume connector variables."""

    def __init__(self, network: SimulationNetwork) -> None:
        self.network = network
        self._pneumatic_ports = tuple(
            component.get_port(definition.name)
            for component in network.components.values()
            for definition in component.active_port_definitions
            if definition.kind == "physical" and definition.domain == "pneumatic"
        )
        self._output_components = tuple(
            component
            for component in network.components.values()
            if type(component).pneumatic_volume_outputs
            is not Component.pneumatic_volume_outputs
        )
        self._connected_endpoint = self._build_connection_map()
        self.last_diagnostics: PneumaticVolumeDiagnostics | None = None

    def _build_connection_map(
        self,
    ) -> dict[Endpoint, _PneumaticVolumeConnectionBinding]:
        result: dict[Endpoint, _PneumaticVolumeConnectionBinding] = {}
        for connection in self.network.connections:
            if connection.kind != "physical" or connection.domain != "pneumatic":
                continue
            first, second = connection.endpoints
            result[first] = _PneumaticVolumeConnectionBinding(
                connected_endpoint=second,
                connected_port=self.network.components[second.component].get_port(
                    second.port
                ),
            )
            result[second] = _PneumaticVolumeConnectionBinding(
                connected_endpoint=first,
                connected_port=self.network.components[first.component].get_port(
                    first.port
                ),
            )
        return result

    @profile_phase("simulation.pneumatic_volume", minimum_mode="audit")
    def solve(self) -> PneumaticVolumeDiagnostics:
        for port in self._pneumatic_ports:
            port.volume = 0.0
            port.volume_flow = 0.0

        outputs: dict[Endpoint, tuple[float, float]] = {}
        for component in self._output_components:
            for port_name, raw_values in component.pneumatic_volume_outputs().items():
                port = component.get_port(port_name)
                definition = port.definition
                if (
                    definition is None
                    or definition.kind != "physical"
                    or definition.domain != "pneumatic"
                ):
                    raise ValueError(
                        f"Component {component.name} declares pneumatic volume output "
                        f"on non-pneumatic port {port_name}."
                    )
                volume, volume_flow = (float(raw_values[0]), float(raw_values[1]))
                if not isfinite(volume) or not isfinite(volume_flow):
                    raise ValueError(
                        f"Component {component.name}.{port_name} produced a non-finite "
                        "pneumatic volume value."
                    )
                endpoint = Endpoint(component.name, port_name)
                outputs[endpoint] = (volume, volume_flow)
                port.volume = volume
                port.volume_flow = volume_flow

        propagated = 0
        for endpoint, values in outputs.items():
            binding = self._connected_endpoint.get(endpoint)
            if binding is None:
                continue
            if binding.connected_endpoint in outputs:
                raise ValueError(
                    "A pneumatic connection cannot contain two external-volume "
                    f"sources: {endpoint} and {binding.connected_endpoint}."
                )
            binding.connected_port.volume, binding.connected_port.volume_flow = values
            propagated += 1

        diagnostics = PneumaticVolumeDiagnostics(
            propagated=propagated,
            output_ports=tuple(sorted(str(endpoint) for endpoint in outputs)),
        )
        self.last_diagnostics = diagnostics
        return diagnostics