from __future__ import annotations

from dataclasses import dataclass

from app.simulation.core.base import DynamicComponent
from app.simulation.systems.network import Endpoint, SimulationNetwork


class StreamSolveError(RuntimeError):
    def __init__(self, message: str, diagnostics: "StreamSolveDiagnostics") -> None:
        super().__init__(message)
        self.diagnostics = diagnostics


@dataclass(frozen=True)
class StreamSolveDiagnostics:
    converged: bool
    iterations: int
    max_delta: float

    def as_dict(self) -> dict[str, object]:
        return {
            "converged": self.converged,
            "iterations": self.iterations,
            "maxDelta": self.max_delta,
        }


class StreamResolver:
    """Resolve outflow enthalpy propagation after pressure and flow are known."""

    def __init__(
        self,
        network: SimulationNetwork,
        *,
        relative_tolerance: float = 1e-9,
        max_iterations: int = 100,
    ) -> None:
        self.network = network
        self.relative_tolerance = relative_tolerance
        self.max_iterations = max_iterations
        self._connected_endpoint = self._build_connection_map()
        self.last_diagnostics: StreamSolveDiagnostics | None = None

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

    def connected_enthalpies(self) -> dict[str, dict[str, float]]:
        values: dict[str, dict[str, float]] = {
            component.name: {} for component in self.network.components.values()
        }
        for endpoint, connected in self._connected_endpoint.items():
            connected_port = self.network.components[connected.component].get_port(
                connected.port
            )
            values[endpoint.component][endpoint.port] = connected_port.h_outflow
        return values

    def solve(self) -> tuple[StreamSolveDiagnostics, dict[str, dict[str, float]]]:
        dynamic_components = [
            component
            for component in self.network.components.values()
            if isinstance(component, DynamicComponent)
        ]
        for component in dynamic_components:
            component.refresh_thermodynamic_ports()

        max_delta = 0.0
        for iteration in range(1, self.max_iterations + 1):
            previous = {
                (component.name, port_name): port.h_outflow
                for component in self.network.components.values()
                for port_name, port in component.ports.items()
            }
            connected = self.connected_enthalpies()
            for component in self.network.components.values():
                if isinstance(component, DynamicComponent):
                    component.refresh_thermodynamic_ports()
                else:
                    component.update_stream_outflows(connected[component.name])

            deltas = [
                abs(port.h_outflow - previous[(component.name, port_name)])
                for component in self.network.components.values()
                for port_name, port in component.ports.items()
            ]
            magnitudes = [
                abs(port.h_outflow)
                for component in self.network.components.values()
                for port in component.ports.values()
            ]
            max_delta = max(deltas, default=0.0)
            scale = max(magnitudes + [1.0])
            if max_delta <= self.relative_tolerance * scale:
                diagnostics = StreamSolveDiagnostics(
                    converged=True,
                    iterations=iteration,
                    max_delta=max_delta,
                )
                self.last_diagnostics = diagnostics
                return diagnostics, self.connected_enthalpies()

        diagnostics = StreamSolveDiagnostics(
            converged=False,
            iterations=self.max_iterations,
            max_delta=max_delta,
        )
        self.last_diagnostics = diagnostics
        raise StreamSolveError(
            "Stream enthalpy propagation did not converge.",
            diagnostics,
        )