from __future__ import annotations

from dataclasses import dataclass
from math import isfinite

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),
        }


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

    def __init__(self, network: SimulationNetwork) -> None:
        self.network = network
        self._connected_endpoint = self._build_connection_map()
        self.last_diagnostics: PneumaticVolumeDiagnostics | None = None

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

    def solve(self) -> PneumaticVolumeDiagnostics:
        for component in self.network.components.values():
            for definition in component.active_port_definitions:
                if definition.kind == "physical" and definition.domain == "pneumatic":
                    port = component.get_port(definition.name)
                    port.volume = 0.0
                    port.volume_flow = 0.0

        outputs: dict[Endpoint, tuple[float, float]] = {}
        for component in self.network.components.values():
            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():
            connected = self._connected_endpoint.get(endpoint)
            if connected is None:
                continue
            if connected in outputs:
                raise ValueError(
                    "A pneumatic connection cannot contain two external-volume "
                    f"sources: {endpoint} and {connected}."
                )
            connected_port = self.network.components[connected.component].get_port(
                connected.port
            )
            connected_port.volume, 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