from __future__ import annotations

from dataclasses import dataclass

from app.simulation.core.base import Component, DynamicComponent
from app.simulation.core.equations import EquationResidual
from app.simulation.core.metadata import ResultVariableMetadata
from app.simulation.core.ports import PortState


@dataclass(frozen=True)
class Endpoint:
    component: str
    port: str

    @property
    def key(self) -> tuple[str, str]:
        return self.component, self.port

    def __str__(self) -> str:
        return f"{self.component}.{self.port}"


@dataclass(frozen=True)
class Connection:
    id: str
    kind: str
    domain: str
    endpoint_a: Endpoint
    endpoint_b: Endpoint

    @property
    def endpoints(self) -> tuple[Endpoint, Endpoint]:
        return self.endpoint_a, self.endpoint_b

    @property
    def undirected_key(self) -> tuple[tuple[str, str], tuple[str, str]]:
        first, second = sorted((self.endpoint_a.key, self.endpoint_b.key))
        return first, second

    # Compatibility accessors for existing reports. They do not imply physical flow.
    @property
    def source_component(self) -> str:
        return self.endpoint_a.component

    @property
    def source_port(self) -> str:
        return self.endpoint_a.port

    @property
    def target_component(self) -> str:
        return self.endpoint_b.component

    @property
    def target_port(self) -> str:
        return self.endpoint_b.port

    def as_interface_dict(self) -> dict[str, object]:
        return {
            "id": self.id,
            "kind": self.kind,
            "domain": self.domain,
            "endpoints": [
                {"component": endpoint.component, "port": endpoint.port}
                for endpoint in self.endpoints
            ],
        }


class SimulationNetwork:
    """Container for components, topology, and state-vector bookkeeping."""

    def __init__(self, name: str) -> None:
        self.name = name
        self.components: dict[str, Component] = {}
        self.connections: list[Connection] = []

    def add_component(self, component: Component) -> None:
        if component.name in self.components:
            raise ValueError(f"Duplicate component name: {component.name}")
        self.components[component.name] = component

    def connect(
        self,
        endpoint_a_component: str,
        endpoint_a_port: str,
        endpoint_b_component: str,
        endpoint_b_port: str,
        *,
        connection_id: str | None = None,
    ) -> Connection:
        endpoint_a = Endpoint(endpoint_a_component, endpoint_a_port)
        endpoint_b = Endpoint(endpoint_b_component, endpoint_b_port)
        if endpoint_a.component == endpoint_b.component:
            raise ValueError(
                f"Cannot connect component {endpoint_a.component} to itself."
            )

        first_port = self._active_port_for(endpoint_a)
        second_port = self._active_port_for(endpoint_b)
        first_definition = first_port.definition
        second_definition = second_port.definition
        if first_definition is None or second_definition is None:
            raise ValueError("Connected ports must expose interface definitions.")
        if first_definition.kind != second_definition.kind:
            raise ValueError(f"Connection mixes physical and signal ports: {endpoint_a}, {endpoint_b}.")
        if first_definition.domain != second_definition.domain:
            raise ValueError(f"Connection domains do not match: {endpoint_a}, {endpoint_b}.")
        if first_definition.variables != second_definition.variables:
            raise ValueError(
                f"Connection variable contracts do not match: {endpoint_a}, {endpoint_b}."
            )
        if first_definition.kind == "signal" and {
            first_definition.nominal_role,
            second_definition.nominal_role,
        } != {"input", "output"}:
            raise ValueError("A signal connection must contain one output and one input.")

        occupied_endpoints = {
            endpoint
            for item in self.connections
            for endpoint in item.endpoints
        }
        if first_definition.kind == "physical":
            occupied = [
                str(endpoint)
                for endpoint in (endpoint_a, endpoint_b)
                if endpoint in occupied_endpoints
            ]
            if occupied:
                raise ValueError(
                    "Physical ports accept one connection; already connected: "
                    + ", ".join(occupied)
                    + ". Use a junction component for branching."
                )
        else:
            signal_input = (
                endpoint_a
                if first_definition.nominal_role == "input"
                else endpoint_b
            )
            if signal_input in occupied_endpoints:
                raise ValueError(
                    f"Signal input {signal_input} already has a driver."
                )

        if first_definition.kind == "physical" and endpoint_b.key < endpoint_a.key:
            endpoint_a, endpoint_b = endpoint_b, endpoint_a

        connection = Connection(
            id=connection_id or f"connection_{len(self.connections) + 1}",
            kind=first_definition.kind,
            domain=first_definition.domain,
            endpoint_a=endpoint_a,
            endpoint_b=endpoint_b,
        )
        if any(item.undirected_key == connection.undirected_key for item in self.connections):
            raise ValueError(f"Duplicate connection between {endpoint_a} and {endpoint_b}.")
        if any(item.id == connection.id for item in self.connections):
            raise ValueError(f"Duplicate connection id: {connection.id}.")
        self.connections.append(connection)
        return connection

    def _port_for(self, endpoint: Endpoint) -> PortState:
        try:
            component = self.components[endpoint.component]
        except KeyError as exc:
            raise ValueError(f"Unknown component: {endpoint.component}.") from exc
        return component.get_port(endpoint.port)

    def _active_port_for(self, endpoint: Endpoint) -> PortState:
        try:
            component = self.components[endpoint.component]
        except KeyError as exc:
            raise ValueError(f"Unknown component: {endpoint.component}.") from exc
        active_names = {
            definition.name for definition in component.active_port_definitions
        }
        if endpoint.port not in active_names:
            raise ValueError(
                f"Port {endpoint} is not active for the component's current parameters."
            )
        return component.get_port(endpoint.port)

    def connection_equation_residuals(self) -> tuple[EquationResidual, ...]:
        """Evaluate connector equations that have a direct scalar residual.

        Stream variables are resolved by the stream-mixing layer and therefore do
        not incorrectly appear here as an equality between outflow properties.
        """

        residuals: list[EquationResidual] = []
        for connection in self.connections:
            if connection.kind != "physical":
                continue

            first_port = self._port_for(connection.endpoint_a)
            second_port = self._port_for(connection.endpoint_b)
            definition = first_port.definition
            if definition is None:
                raise ValueError(
                    f"Connected port {connection.endpoint_a} has no interface definition."
                )

            for variable in definition.variables:
                if variable.connection_rule == "equal":
                    value = float(getattr(first_port, variable.name)) - float(
                        getattr(second_port, variable.name)
                    )
                elif variable.connection_rule == "sumToZero":
                    value = float(getattr(first_port, variable.name)) + float(
                        getattr(second_port, variable.name)
                    )
                else:
                    continue
                residuals.append(
                    EquationResidual(
                        id=f"{connection.id}:{variable.name}",
                        owner="connection",
                        owner_id=connection.id,
                        relation=variable.connection_rule,
                        variables=(
                            f"{connection.endpoint_a}.{variable.name}",
                            f"{connection.endpoint_b}.{variable.name}",
                        ),
                        role=variable.role,
                        value=value,
                    )
                )
        return tuple(residuals)

    def pressure_flow_equation_residuals(self) -> tuple[EquationResidual, ...]:
        """Evaluate the complete algebraic pressure-flow equation subsystem."""

        component_residuals = tuple(
            residual
            for component in self.components.values()
            for residual in component.pressure_flow_equation_residuals()
        )
        return component_residuals + self.connection_equation_residuals()

    def pressure_flow_unknowns(self) -> tuple[str, ...]:
        return tuple(
            f"{component.name}.{definition.name}.{variable.name}"
            for component in self.components.values()
            for definition in component.active_port_definitions
            if definition.kind == "physical"
            for variable in definition.variables
            if variable.role in {"effort", "flow"}
        )

    def pressure_flow_structure_dict(self) -> dict[str, object]:
        unknowns = self.pressure_flow_unknowns()
        equations = self.pressure_flow_equation_residuals()
        return {
            "unknownCount": len(unknowns),
            "equationCount": len(equations),
            "isSquare": len(unknowns) == len(equations),
            "unknowns": list(unknowns),
            "equations": [
                equation.as_definition_dict() for equation in equations
            ],
        }

    def dynamic_components(self) -> list[DynamicComponent]:
        return [
            component
            for component in self.components.values()
            if isinstance(component, DynamicComponent)
        ]

    def initial_state_vector(self) -> list[float]:
        values: list[float] = []
        for component in self.dynamic_components():
            values.extend(component.get_state_vector())
        return values

    def apply_state_vector(self, values: list[float]) -> None:
        cursor = 0
        for component in self.dynamic_components():
            next_cursor = cursor + component.state_size
            component.set_state_vector(values[cursor:next_cursor])
            cursor = next_cursor
        if cursor != len(values):
            raise ValueError("State vector length does not match dynamic components.")

    def result_variable_metadata(self) -> tuple[ResultVariableMetadata, ...]:
        return tuple(
            variable
            for component in self.components.values()
            for variable in component.result_variable_metadata()
        )

    def summary(self) -> str:
        lines = [f"Network: {self.name}", "Components:"]
        for name, component in self.components.items():
            lines.append(f"  - {name}: {component.__class__.__name__}")
        lines.append("Connections:")
        for conn in self.connections:
            connector = "<->" if conn.kind == "physical" else "->"
            lines.append(
                f"  - {conn.endpoint_a} {connector} {conn.endpoint_b}"
            )
        return "\n".join(lines)

    def as_interface_dict(self) -> dict[str, object]:
        connected_endpoints = {
            endpoint.key
            for connection in self.connections
            for endpoint in connection.endpoints
        }
        return {
            "name": self.name,
            "components": [
                {
                    "id": component.name,
                    "type": component.model_type,
                    "parameters": component.parameter_interface_dicts(),
                    "ports": [
                        definition.as_interface_dict()
                        for definition in component.active_port_definitions
                    ],
                    "resultVariables": [
                        variable.as_dict()
                        for variable in component.result_variable_metadata()
                    ],
                }
                for component in self.components.values()
            ],
            "connections": [
                connection.as_interface_dict() for connection in self.connections
            ],
            "pressureFlowSystem": self.pressure_flow_structure_dict(),
            "unconnectedPorts": [
                {"component": component.name, "port": definition.name}
                for component in self.components.values()
                for definition in component.active_port_definitions
                if (component.name, definition.name) not in connected_endpoints
            ],
        }