//! Brine Boundary Condition Components
//!
//! This module provides [`BrineSource`] and [`BrineSink`] components for
//! water-glycol (brine) heat transfer circuits with native glycol concentration
//! support via the [`Concentration`](entropyk_core::Concentration) NewType.
//!
//! ## Design
//!
//! Unlike the generic [`FlowSource`](crate::FlowSource)/[`FlowSink`](crate::FlowSink)
//! which use (Pressure, Enthalpy), these components use (Pressure, Temperature,
//! Concentration) for type-safe brine state specification.
//!
//! ## Equations
//!
//! **BrineSource** (Modelica-style Fixed/Free, dynamic equation count):
//! ```text
//! r = P_edge − P_set = 0                       (when pressure is Fixed)
//! r = h_edge − h(P_set, T_set, c) = 0           (when temperature is Fixed)
//! r = ṁ_edge − ṁ_set = 0                       (when mass flow is Fixed)
//! ```
//! Defaults match `Boundary_pT` + optional flow: Fixed P, Fixed T, optional Fixed ṁ.
//! Free ṁ + Fixed sink T_out is the usual ΔT-rating pattern (ṁ emerges from energy).
//!
//! **BrineSink** (1 or 2 depending on whether temperature is set):
//! ```text
//! r₀ = P_edge − P_back = 0                     (always)
//! r₁ = h_edge − h(P_back, T_back, c) = 0       (only when temperature is set)
//! ```
//!
//! ## Example
//!
//! ```ignore
//! use entropyk_components::{BrineSource, BrineSink};
//! use entropyk_core::{Pressure, Temperature, Concentration};
//! use std::sync::Arc;
//!
//! let backend = Arc::new(coolprop_backend);
//!
//! // Evaporator brine outlet: 3 bar, 20 °C, 30% MEG
//! let source = BrineSource::new(
//!     "MEG",
//!     Pressure::from_bar(3.0),
//!     Temperature::from_celsius(20.0),
//!     Concentration::from_percent(30.0),
//!     backend.clone(),
//!     outlet_port,
//! ).unwrap();
//!
//! // Free-enthalpy sink: only back-pressure constrained
//! let sink = BrineSink::new(
//!     "MEG",
//!     Pressure::from_bar(2.0),
//!     None,
//!     None,
//!     backend,
//!     inlet_port,
//! ).unwrap();
//! ```

use crate::flow_junction::is_incompressible;
use crate::{
    Component, ComponentError, ConnectedPort, JacobianBuilder, ResidualVector, StateSlice,
};
use entropyk_core::{Concentration, Enthalpy, MassFlow, Power, Pressure, Temperature};
use entropyk_fluids::{FluidBackend, FluidId, FluidState, Property};
use std::sync::Arc;

fn pt_concentration_to_enthalpy(
    backend: &dyn FluidBackend,
    fluid: &str,
    p: Pressure,
    t: Temperature,
    concentration: Concentration,
) -> Result<Enthalpy, ComponentError> {
    // Encode the glycol concentration into the fluid name using CoolProp's
    // incompressible mixture syntax: "INCOMP::MEG-30" for 30% MEG by mass.
    // Pure water (concentration ≈ 0) is passed as-is.
    let fluid_with_conc = if concentration.to_fraction() < 1e-10 {
        fluid.to_string()
    } else {
        format!("INCOMP::{}-{:.0}", fluid, concentration.to_percent())
    };

    let fluid_id = FluidId::new(&fluid_with_conc);
    let state = FluidState::from_pt(p, t);

    backend
        .property(fluid_id, Property::Enthalpy, state)
        .map(Enthalpy::from_joules_per_kg)
        .map_err(|e| {
            ComponentError::CalculationFailed(format!(
                "P-T-Concentration to enthalpy (fluid='{}', c={:.1}%): {}",
                fluid_with_conc,
                concentration.to_percent(),
                e
            ))
        })
}

/// A boundary source that imposes pressure, temperature and glycol concentration
/// on its outlet edge (each quantity optionally Fixed, Modelica-style).
///
/// Default equation count is **2** (Fixed P + Fixed T/h). Optional Fixed ṁ adds a
/// third equation. Freeing P or T drops the corresponding residual.
///
/// Only accepts incompressible fluids (MEG, PEG, Water, Glycol, Brine).
/// Use [`RefrigerantSource`](crate::RefrigerantSource) for refrigerant circuits.
pub struct BrineSource {
    fluid_id: String,
    p_set_pa: f64,
    t_set_k: f64,
    concentration: Concentration,
    h_set_jkg: f64,
    backend: Arc<dyn FluidBackend>,
    outlet: ConnectedPort,
    /// Outlet edge (P, h) state indices, wired by `set_system_context`.
    /// When present the residuals/Jacobian are state-driven (BOLT
    /// `BoundaryNode.Coolant.Source` semantics); otherwise the legacy
    /// port-frozen behaviour is kept for standalone unit tests.
    outlet_p_idx: Option<usize>,
    outlet_h_idx: Option<usize>,
    /// Outlet edge ṁ state index (for the optional flow constraint).
    outlet_m_idx: Option<usize>,
    /// Optional imposed mass flow [kg/s] (BOLT `Vd_fixed=true` equivalent):
    /// adds `r = ṁ_edge − ṁ_set` (+1 equation).
    m_flow_set_kg_s: Option<f64>,
    /// When true (default), impose `P_edge = P_set`.
    impose_pressure: bool,
    /// When true (default), impose `h_edge = h(P_set, T_set, c)`.
    impose_temperature: bool,
}

impl BrineSource {
    /// Creates a new `BrineSource`.
    ///
    /// # Arguments
    ///
    /// * `fluid`         — Fluid name accepted by the backend (e.g. `"MEG"`, `"Water"`)
    /// * `p_set`         — Set-point pressure
    /// * `t_set`         — Set-point temperature
    /// * `concentration` — Glycol mass fraction; `Concentration::from_percent(30.0)` = 30 % MEG
    /// * `backend`       — Shared fluid property backend
    /// * `outlet`        — Already-connected outlet port
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if:
    /// - `fluid` is not an incompressible brine fluid
    /// - `p_set` ≤ 0
    /// - The backend fails to evaluate enthalpy at (P, T, c)
    pub fn new(
        fluid: impl Into<String>,
        p_set: Pressure,
        t_set: Temperature,
        concentration: Concentration,
        backend: Arc<dyn FluidBackend>,
        outlet: ConnectedPort,
    ) -> Result<Self, ComponentError> {
        let fluid = fluid.into();

        if !is_incompressible(&fluid) {
            return Err(ComponentError::InvalidState(format!(
                "BrineSource: '{}' is not an incompressible fluid. Use RefrigerantSource instead.",
                fluid
            )));
        }

        let p_set_pa = p_set.to_pascals();
        if p_set_pa <= 0.0 {
            return Err(ComponentError::InvalidState(
                "BrineSource: set-point pressure must be positive".into(),
            ));
        }

        let h_set =
            pt_concentration_to_enthalpy(backend.as_ref(), &fluid, p_set, t_set, concentration)?;

        Ok(Self {
            fluid_id: fluid,
            p_set_pa,
            t_set_k: t_set.to_kelvin(),
            concentration,
            h_set_jkg: h_set.to_joules_per_kg(),
            backend,
            outlet,
            outlet_p_idx: None,
            outlet_h_idx: None,
            outlet_m_idx: None,
            m_flow_set_kg_s: None,
            impose_pressure: true,
            impose_temperature: true,
        })
    }

    /// Imposes the loop mass flow at the boundary (+1 equation,
    /// Modelica `MassFlowSource_T` / BOLT `Vd_fixed=true`).
    ///
    /// Also clears pressure imposition (Free P): a MassFlowSource must not
    /// prescribe both ṁ and P. Pair with a sink `Boundary_pT` and an isobaric
    /// (or frictional) secondary path on the HX. Call
    /// [`with_impose_pressure`](Self::with_impose_pressure) afterward only for
    /// an explicit non-Modelica override.
    ///
    /// Do NOT combine with another flow imposition in the same branch (e.g. a
    /// `ThermalLoad`/`Pump` design flow) or the loop becomes over-determined.
    pub fn with_imposed_mass_flow(mut self, m_flow_kg_s: f64) -> Result<Self, ComponentError> {
        if !(m_flow_kg_s.is_finite() && m_flow_kg_s > 0.0) {
            return Err(ComponentError::InvalidState(
                "BrineSource: imposed mass flow must be positive".into(),
            ));
        }
        self.m_flow_set_kg_s = Some(m_flow_kg_s);
        self.impose_pressure = false;
        Ok(self)
    }

    /// Clears an imposed mass flow (Free ṁ — typical for ΔT rating).
    pub fn clear_imposed_mass_flow(mut self) -> Self {
        self.m_flow_set_kg_s = None;
        self
    }

    /// Fixed (true, default) or Free pressure at the outlet edge.
    pub fn with_impose_pressure(mut self, impose: bool) -> Self {
        self.impose_pressure = impose;
        self
    }

    /// Fixed (true, default) or Free temperature/enthalpy at the outlet edge.
    pub fn with_impose_temperature(mut self, impose: bool) -> Self {
        self.impose_temperature = impose;
        self
    }

    /// Returns whether pressure is Fixed.
    pub fn imposes_pressure(&self) -> bool {
        self.impose_pressure
    }

    /// Returns whether temperature/enthalpy is Fixed.
    pub fn imposes_temperature(&self) -> bool {
        self.impose_temperature
    }

    /// Returns the optional imposed mass flow [kg/s].
    pub fn m_flow_set(&self) -> Option<f64> {
        self.m_flow_set_kg_s
    }

    /// Returns the fluid identifier string (e.g. `"MEG"`, `"Water"`).
    pub fn fluid_id(&self) -> &str {
        &self.fluid_id
    }

    /// Returns the set-point pressure.
    pub fn p_set(&self) -> Pressure {
        Pressure::from_pascals(self.p_set_pa)
    }

    /// Returns the set-point temperature.
    pub fn t_set(&self) -> Temperature {
        Temperature::from_kelvin(self.t_set_k)
    }

    /// Returns the glycol concentration.
    pub fn concentration(&self) -> Concentration {
        self.concentration
    }

    /// Returns the pre-computed set-point enthalpy.
    pub fn h_set(&self) -> Enthalpy {
        Enthalpy::from_joules_per_kg(self.h_set_jkg)
    }

    /// Returns a reference to the outlet connected port.
    pub fn outlet(&self) -> &ConnectedPort {
        &self.outlet
    }

    /// Updates the set-point pressure and recomputes the cached enthalpy.
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if `p` ≤ 0, or
    /// [`ComponentError::CalculationFailed`] if backend evaluation fails.
    pub fn set_pressure(&mut self, p: Pressure) -> Result<(), ComponentError> {
        let p_pa = p.to_pascals();
        if p_pa <= 0.0 {
            return Err(ComponentError::InvalidState(
                "BrineSource: pressure must be positive".into(),
            ));
        }
        self.p_set_pa = p_pa;
        self.recompute_enthalpy()
    }

    /// Updates the set-point temperature and recomputes the cached enthalpy.
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::CalculationFailed`] if backend evaluation fails.
    pub fn set_temperature(&mut self, t: Temperature) -> Result<(), ComponentError> {
        self.t_set_k = t.to_kelvin();
        self.recompute_enthalpy()
    }

    /// Updates the glycol concentration and recomputes the cached enthalpy.
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::CalculationFailed`] if backend evaluation fails.
    pub fn set_concentration(
        &mut self,
        concentration: Concentration,
    ) -> Result<(), ComponentError> {
        self.concentration = concentration;
        self.recompute_enthalpy()
    }

    fn recompute_enthalpy(&mut self) -> Result<(), ComponentError> {
        let h_set = pt_concentration_to_enthalpy(
            self.backend.as_ref(),
            &self.fluid_id,
            Pressure::from_pascals(self.p_set_pa),
            Temperature::from_kelvin(self.t_set_k),
            self.concentration,
        )?;
        self.h_set_jkg = h_set.to_joules_per_kg();
        Ok(())
    }
}

impl Component for BrineSource {
    fn n_equations(&self) -> usize {
        usize::from(self.impose_pressure)
            + usize::from(self.impose_temperature)
            + usize::from(self.m_flow_set_kg_s.is_some())
    }

    fn equation_roles(&self) -> Vec<crate::EquationRole> {
        let mut roles = Vec::with_capacity(self.n_equations());
        if self.impose_pressure {
            roles.push(crate::EquationRole::BoundaryDirichlet { quantity: "P" });
        }
        if self.impose_temperature {
            roles.push(crate::EquationRole::BoundaryDirichlet { quantity: "h" });
        }
        if self.m_flow_set_kg_s.is_some() {
            roles.push(crate::EquationRole::BoundaryDirichlet { quantity: "m" });
        }
        roles
    }

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize, usize)],
    ) {
        // A source has one incident edge: its outlet (outgoing). Incident
        // lists put incoming edges first, so take the LAST entry.
        if let Some(&(m, p, h)) = external_edge_state_indices.last() {
            self.outlet_m_idx = Some(m);
            self.outlet_p_idx = Some(p);
            self.outlet_h_idx = Some(h);
        }
    }

    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        let n = self.n_equations();
        if residuals.len() < n {
            return Err(ComponentError::InvalidResidualDimensions {
                expected: n,
                actual: residuals.len(),
            });
        }
        let mut row = 0;
        match (self.outlet_p_idx, self.outlet_h_idx) {
            (Some(p_idx), Some(h_idx)) if p_idx < state.len() && h_idx < state.len() => {
                if self.impose_pressure {
                    residuals[row] = state[p_idx] - self.p_set_pa;
                    row += 1;
                }
                if self.impose_temperature {
                    residuals[row] = state[h_idx] - self.h_set_jkg;
                    row += 1;
                }
            }
            _ if state.is_empty() => {
                if self.impose_pressure {
                    residuals[row] = self.outlet.pressure().to_pascals() - self.p_set_pa;
                    row += 1;
                }
                if self.impose_temperature {
                    residuals[row] = self.outlet.enthalpy().to_joules_per_kg() - self.h_set_jkg;
                    row += 1;
                }
            }
            _ => {
                return Err(ComponentError::InvalidState(
                    "BrineSource requires a live outlet edge state in solver context".to_string(),
                ));
            }
        }
        if let Some(m_set) = self.m_flow_set_kg_s {
            let Some(m_idx) = self.outlet_m_idx.filter(|&i| i < state.len()) else {
                return Err(ComponentError::InvalidState(
                    "BrineSource imposed mass flow requires a live outlet mass-flow index"
                        .to_string(),
                ));
            };
            residuals[row] = state[m_idx] - m_set;
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        // Exact unit entries on the live edge columns for each Fixed quantity.
        let (p_col, h_col) = match (self.outlet_p_idx, self.outlet_h_idx) {
            (Some(p), Some(h)) if p < _state.len() && h < _state.len() => (p, h),
            _ if _state.is_empty() => return Ok(()),
            _ => {
                return Err(ComponentError::InvalidState(
                    "BrineSource Jacobian requires live outlet pressure/enthalpy indices"
                        .to_string(),
                ));
            }
        };
        let mut row = 0;
        if self.impose_pressure {
            jacobian.add_entry(row, p_col, 1.0);
            row += 1;
        }
        if self.impose_temperature {
            jacobian.add_entry(row, h_col, 1.0);
            row += 1;
        }
        if self.m_flow_set_kg_s.is_some() {
            if let Some(m_col) = self.outlet_m_idx {
                jacobian.add_entry(row, m_col, 1.0);
            }
        }
        Ok(())
    }

    fn get_ports(&self) -> &[ConnectedPort] {
        &[]
    }

    fn port_mass_flows(&self, _state: &StateSlice) -> Result<Vec<MassFlow>, ComponentError> {
        Ok(vec![MassFlow::from_kg_per_s(0.0)])
    }

    fn port_enthalpies(&self, _state: &StateSlice) -> Result<Vec<Enthalpy>, ComponentError> {
        Ok(vec![self.outlet.enthalpy()])
    }

    fn energy_transfers(&self, _state: &StateSlice) -> Option<(Power, Power)> {
        Some((Power::from_watts(0.0), Power::from_watts(0.0)))
    }

    fn signature(&self) -> String {
        format!(
            "BrineSource({}:P={:.0}Pa,T={:.1}K,c={:.0}%)",
            self.fluid_id,
            self.p_set_pa,
            self.t_set_k,
            self.concentration.to_percent()
        )
    }

    fn to_params(&self) -> crate::ComponentParams {
        crate::ComponentParams::new("BrineSource")
            .with_param("fluid", self.fluid_id.as_str())
            .with_param("pSetPa", self.p_set_pa)
            .with_param("tSetK", self.t_set_k)
            .with_param("concentration", self.concentration.to_fraction())
    }

    fn set_fluid_backend_from_builder(&mut self, backend: Arc<dyn FluidBackend>) {
        self.backend = backend;
    }
}

/// A boundary sink that imposes back-pressure, and optionally a fixed enthalpy via
/// (temperature, concentration) on its inlet edge.
///
/// **Equation count is dynamic:**
/// - 1 equation (free enthalpy): `r₀ = P_edge − P_back = 0`
/// - 2 equations (temperature set): additionally `r₁ = h_edge − h(P_back, T_back, c) = 0`
///
/// Toggle with [`set_temperature`](Self::set_temperature) /
/// [`clear_temperature`](Self::clear_temperature) without rebuilding the component.
pub struct BrineSink {
    fluid_id: String,
    p_back_pa: f64,
    t_opt_k: Option<f64>,
    concentration_opt: Option<Concentration>,
    h_back_jkg: Option<f64>,
    backend: Arc<dyn FluidBackend>,
    inlet: ConnectedPort,
    /// Inlet edge (P, h) state indices, wired by `set_system_context`
    /// (state-driven BOLT `BoundaryNode.Coolant.Sink` semantics).
    inlet_p_idx: Option<usize>,
    inlet_h_idx: Option<usize>,
    /// Inlet edge ṁ state index (for the optional flow constraint).
    inlet_m_idx: Option<usize>,
    /// Optional imposed mass flow [kg/s] (BOLT `Vd_fixed=true` equivalent).
    m_flow_set_kg_s: Option<f64>,
}

impl BrineSink {
    /// Creates a new `BrineSink`.
    ///
    /// # Arguments
    ///
    /// * `fluid`             — Fluid name accepted by the backend (e.g. `"MEG"`, `"Water"`)
    /// * `p_back`            — Back-pressure imposed on the inlet edge
    /// * `t_opt`             — Optional set-point temperature; `None` = free enthalpy (1 equation)
    /// * `concentration_opt` — Required when `t_opt` is `Some`; ignored otherwise
    /// * `backend`           — Shared fluid property backend
    /// * `inlet`             — Already-connected inlet port
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if:
    /// - `fluid` is not an incompressible brine fluid
    /// - `p_back` ≤ 0
    /// - `t_opt` is `Some` but `concentration_opt` is `None`
    /// - The backend fails to evaluate enthalpy
    pub fn new(
        fluid: impl Into<String>,
        p_back: Pressure,
        t_opt: Option<Temperature>,
        concentration_opt: Option<Concentration>,
        backend: Arc<dyn FluidBackend>,
        inlet: ConnectedPort,
    ) -> Result<Self, ComponentError> {
        let fluid = fluid.into();

        if !is_incompressible(&fluid) {
            return Err(ComponentError::InvalidState(format!(
                "BrineSink: '{}' is not an incompressible fluid. Use RefrigerantSink instead.",
                fluid
            )));
        }

        let p_back_pa = p_back.to_pascals();
        if p_back_pa <= 0.0 {
            return Err(ComponentError::InvalidState(
                "BrineSink: back-pressure must be positive".into(),
            ));
        }

        if t_opt.is_some() && concentration_opt.is_none() {
            return Err(ComponentError::InvalidState(
                "BrineSink: concentration must be specified when temperature is set".into(),
            ));
        }

        let (t_opt_k, h_back_jkg) = if let (Some(t), Some(c)) = (t_opt, concentration_opt) {
            let h = pt_concentration_to_enthalpy(backend.as_ref(), &fluid, p_back, t, c)?;
            (Some(t.to_kelvin()), Some(h.to_joules_per_kg()))
        } else {
            (None, None)
        };

        Ok(Self {
            fluid_id: fluid,
            p_back_pa,
            t_opt_k,
            concentration_opt,
            h_back_jkg,
            backend,
            inlet,
            inlet_p_idx: None,
            inlet_h_idx: None,
            inlet_m_idx: None,
            m_flow_set_kg_s: None,
        })
    }

    /// Imposes the loop mass flow at the boundary (+1 equation,
    /// BOLT `Vd_fixed=true` equivalent). Do NOT combine with another flow
    /// imposition in the same branch or the loop becomes over-determined.
    pub fn with_imposed_mass_flow(mut self, m_flow_kg_s: f64) -> Result<Self, ComponentError> {
        if !(m_flow_kg_s.is_finite() && m_flow_kg_s > 0.0) {
            return Err(ComponentError::InvalidState(
                "BrineSink: imposed mass flow must be positive".into(),
            ));
        }
        self.m_flow_set_kg_s = Some(m_flow_kg_s);
        Ok(self)
    }

    /// Returns the optional imposed mass flow [kg/s].
    pub fn m_flow_set(&self) -> Option<f64> {
        self.m_flow_set_kg_s
    }

    /// Returns the fluid identifier string (e.g. `"MEG"`, `"Water"`).
    pub fn fluid_id(&self) -> &str {
        &self.fluid_id
    }

    /// Returns the back-pressure.
    pub fn p_back(&self) -> Pressure {
        Pressure::from_pascals(self.p_back_pa)
    }

    /// Returns the optional set-point temperature.
    pub fn t_opt(&self) -> Option<Temperature> {
        self.t_opt_k.map(Temperature::from_kelvin)
    }

    /// Returns the optional glycol concentration.
    pub fn concentration(&self) -> Option<Concentration> {
        self.concentration_opt
    }

    /// Returns the optional pre-computed back enthalpy.
    pub fn h_back(&self) -> Option<Enthalpy> {
        self.h_back_jkg.map(Enthalpy::from_joules_per_kg)
    }

    /// Returns a reference to the inlet connected port.
    pub fn inlet(&self) -> &ConnectedPort {
        &self.inlet
    }

    /// Updates the back-pressure and recomputes enthalpy if temperature is set.
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if `p` ≤ 0, or
    /// [`ComponentError::CalculationFailed`] if backend evaluation fails.
    pub fn set_pressure(&mut self, p: Pressure) -> Result<(), ComponentError> {
        let p_pa = p.to_pascals();
        if p_pa <= 0.0 {
            return Err(ComponentError::InvalidState(
                "BrineSink: back-pressure must be positive".into(),
            ));
        }
        self.p_back_pa = p_pa;
        if let (Some(t_k), Some(c)) = (self.t_opt_k, self.concentration_opt) {
            self.h_back_jkg = Some(
                pt_concentration_to_enthalpy(
                    self.backend.as_ref(),
                    &self.fluid_id,
                    p,
                    Temperature::from_kelvin(t_k),
                    c,
                )?
                .to_joules_per_kg(),
            );
        }
        Ok(())
    }

    /// Sets a temperature (and required concentration) constraint, switching the sink to
    /// 2-equation mode.  Recomputes the back enthalpy from (P, T, c).
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::CalculationFailed`] if backend evaluation fails.
    pub fn set_temperature(
        &mut self,
        t: Temperature,
        concentration: Concentration,
    ) -> Result<(), ComponentError> {
        self.t_opt_k = Some(t.to_kelvin());
        self.concentration_opt = Some(concentration);
        self.h_back_jkg = Some(
            pt_concentration_to_enthalpy(
                self.backend.as_ref(),
                &self.fluid_id,
                Pressure::from_pascals(self.p_back_pa),
                t,
                concentration,
            )?
            .to_joules_per_kg(),
        );
        Ok(())
    }

    /// Removes the temperature constraint, switching the sink back to 1-equation mode
    /// (free enthalpy).
    pub fn clear_temperature(&mut self) {
        self.t_opt_k = None;
        self.concentration_opt = None;
        self.h_back_jkg = None;
    }
}

impl Component for BrineSink {
    fn n_equations(&self) -> usize {
        let base = if self.h_back_jkg.is_some() { 2 } else { 1 };
        base + usize::from(self.m_flow_set_kg_s.is_some())
    }

    fn equation_roles(&self) -> Vec<crate::EquationRole> {
        let mut roles = vec![crate::EquationRole::BoundaryDirichlet { quantity: "P" }];
        if self.h_back_jkg.is_some() {
            // Optional T/h fix — consumes DoF (do not also free h).
            roles.push(crate::EquationRole::BoundaryDirichlet { quantity: "h" });
        }
        if self.m_flow_set_kg_s.is_some() {
            roles.push(crate::EquationRole::BoundaryDirichlet { quantity: "m" });
        }
        roles
    }

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize, usize)],
    ) {
        // A sink has one incident edge: its inlet (incoming, listed first).
        if let Some(&(m, p, h)) = external_edge_state_indices.first() {
            self.inlet_m_idx = Some(m);
            self.inlet_p_idx = Some(p);
            self.inlet_h_idx = Some(h);
        }
    }

    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        let n = self.n_equations();
        if residuals.len() < n {
            return Err(ComponentError::InvalidResidualDimensions {
                expected: n,
                actual: residuals.len(),
            });
        }
        let mut row = 1;
        match (self.inlet_p_idx, self.inlet_h_idx) {
            (Some(p_idx), Some(h_idx)) if p_idx < state.len() && h_idx < state.len() => {
                residuals[0] = state[p_idx] - self.p_back_pa;
                if let Some(h_back) = self.h_back_jkg {
                    residuals[row] = state[h_idx] - h_back;
                    row += 1;
                }
            }
            _ if state.is_empty() => {
                residuals[0] = self.inlet.pressure().to_pascals() - self.p_back_pa;
                if let Some(h_back) = self.h_back_jkg {
                    residuals[row] = self.inlet.enthalpy().to_joules_per_kg() - h_back;
                    row += 1;
                }
            }
            _ => {
                return Err(ComponentError::InvalidState(
                    "BrineSink requires a live inlet edge state in solver context".to_string(),
                ));
            }
        }
        if let Some(m_set) = self.m_flow_set_kg_s {
            let Some(m_idx) = self.inlet_m_idx.filter(|&i| i < state.len()) else {
                return Err(ComponentError::InvalidState(
                    "BrineSink imposed mass flow requires a live inlet mass-flow index".to_string(),
                ));
            };
            let m = state[m_idx];
            residuals[row] = m - m_set;
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        let (p_col, h_col) = match (self.inlet_p_idx, self.inlet_h_idx) {
            (Some(p), Some(h)) if p < _state.len() && h < _state.len() => (p, h),
            _ if _state.is_empty() => return Ok(()),
            _ => {
                return Err(ComponentError::InvalidState(
                    "BrineSink Jacobian requires live inlet pressure/enthalpy indices".to_string(),
                ));
            }
        };
        jacobian.add_entry(0, p_col, 1.0);
        let mut row = 1;
        if self.h_back_jkg.is_some() {
            jacobian.add_entry(row, h_col, 1.0);
            row += 1;
        }
        if self.m_flow_set_kg_s.is_some() {
            if let Some(m_col) = self.inlet_m_idx {
                jacobian.add_entry(row, m_col, 1.0);
            }
        }
        Ok(())
    }

    fn get_ports(&self) -> &[ConnectedPort] {
        &[]
    }

    fn port_mass_flows(&self, _state: &StateSlice) -> Result<Vec<MassFlow>, ComponentError> {
        Ok(vec![MassFlow::from_kg_per_s(0.0)])
    }

    fn port_enthalpies(&self, _state: &StateSlice) -> Result<Vec<Enthalpy>, ComponentError> {
        Ok(vec![self.inlet.enthalpy()])
    }

    fn energy_transfers(&self, _state: &StateSlice) -> Option<(Power, Power)> {
        Some((Power::from_watts(0.0), Power::from_watts(0.0)))
    }

    fn signature(&self) -> String {
        match self.t_opt_k {
            Some(t_k) => format!(
                "BrineSink({}:P={:.0}Pa,T={:.1}K,c={:.0}%)",
                self.fluid_id,
                self.p_back_pa,
                t_k,
                self.concentration_opt.map_or(0.0, |c| c.to_percent())
            ),
            None => format!(
                "BrineSink({}:P={:.0}Pa,T=free)",
                self.fluid_id, self.p_back_pa
            ),
        }
    }

    fn to_params(&self) -> crate::ComponentParams {
        let mut params = crate::ComponentParams::new("BrineSink")
            .with_param("fluid", self.fluid_id.as_str())
            .with_param("pBackPa", self.p_back_pa);
        if let Some(t_k) = self.t_opt_k {
            params = params.with_param("tBackK", t_k);
        }
        if let Some(c) = self.concentration_opt {
            params = params.with_param("concentration", c.to_fraction());
        }
        params
    }

    fn set_fluid_backend_from_builder(&mut self, backend: Arc<dyn FluidBackend>) {
        self.backend = backend;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::port::{FluidId, Port};
    use entropyk_fluids::{
        CriticalPoint, FluidBackend, FluidError, FluidResult, FluidState, Phase, Property,
    };

    struct MockBrineBackend;

    impl MockBrineBackend {
        fn new() -> Self {
            Self
        }
    }

    impl FluidBackend for MockBrineBackend {
        fn property(
            &self,
            _fluid: FluidId,
            property: Property,
            state: FluidState,
        ) -> FluidResult<f64> {
            match state {
                FluidState::PressureTemperature(p, t) => {
                    let _p_pa = p.to_pascals();
                    let t_k = t.to_kelvin();
                    match property {
                        Property::Enthalpy => {
                            let h = 3500.0 * (t_k - 273.15);
                            Ok(h)
                        }
                        Property::Temperature => Ok(t_k),
                        Property::Pressure => Ok(_p_pa),
                        _ => Err(FluidError::UnsupportedProperty {
                            property: property.to_string(),
                        }),
                    }
                }
                _ => Err(FluidError::InvalidState {
                    reason: "MockBrineBackend only supports P-T state".to_string(),
                }),
            }
        }

        fn critical_point(&self, _fluid: FluidId) -> FluidResult<CriticalPoint> {
            Ok(CriticalPoint::new(
                Temperature::from_kelvin(373.0),
                Pressure::from_pascals(22.06e6),
                322.0,
            ))
        }

        fn is_fluid_available(&self, fluid: &FluidId) -> bool {
            let s = fluid.as_str();
            s == "Water" || s == "Glycol" || s == "MEG" || s == "PEG" || s.starts_with("INCOMP::")
        }

        fn phase(&self, _fluid: FluidId, _state: FluidState) -> FluidResult<Phase> {
            Ok(Phase::Liquid)
        }

        fn full_state(
            &self,
            _fluid: FluidId,
            _p: Pressure,
            _h: entropyk_core::Enthalpy,
        ) -> FluidResult<entropyk_fluids::ThermoState> {
            Err(FluidError::UnsupportedProperty {
                property: "full_state".to_string(),
            })
        }

        fn list_fluids(&self) -> Vec<FluidId> {
            vec![
                FluidId::new("Glycol"),
                FluidId::new("MEG"),
                FluidId::new("PEG"),
                FluidId::new("Water"),
            ]
        }
    }

    fn make_port(fluid: &str, p_pa: f64, h_jkg: f64) -> ConnectedPort {
        let a = Port::new(
            FluidId::new(fluid),
            Pressure::from_pascals(p_pa),
            Enthalpy::from_joules_per_kg(h_jkg),
        );
        let b = Port::new(
            FluidId::new(fluid),
            Pressure::from_pascals(p_pa),
            Enthalpy::from_joules_per_kg(h_jkg),
        );
        a.connect(b).unwrap().0
    }

    #[test]
    fn test_brine_source_creation() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 3.0e5, 70_000.0);
        let source = BrineSource::new(
            "Glycol",
            Pressure::from_pascals(3.0e5),
            Temperature::from_celsius(20.0),
            Concentration::from_percent(30.0),
            backend,
            port,
        )
        .unwrap();

        assert_eq!(source.n_equations(), 2);
        assert_eq!(source.fluid_id(), "Glycol");
    }

    #[test]
    fn test_brine_source_rejects_refrigerant() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("R410A", 8.5e5, 260_000.0);
        let result = BrineSource::new(
            "R410A",
            Pressure::from_pascals(8.5e5),
            Temperature::from_celsius(10.0),
            Concentration::from_percent(30.0),
            backend,
            port,
        );

        assert!(result.is_err());
    }

    #[test]
    fn test_brine_sink_creation() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 2.0e5, 60_000.0);
        let sink = BrineSink::new(
            "Glycol",
            Pressure::from_pascals(2.0e5),
            None,
            None,
            backend,
            port,
        )
        .unwrap();

        assert_eq!(sink.n_equations(), 1);
        assert_eq!(sink.fluid_id(), "Glycol");
    }

    #[test]
    fn test_brine_sink_dynamic_toggle() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 2.0e5, 60_000.0);
        let mut sink = BrineSink::new(
            "Glycol",
            Pressure::from_pascals(2.0e5),
            None,
            None,
            backend,
            port,
        )
        .unwrap();

        assert_eq!(sink.n_equations(), 1);

        sink.set_temperature(
            Temperature::from_celsius(15.0),
            Concentration::from_percent(30.0),
        )
        .unwrap();
        assert_eq!(sink.n_equations(), 2);

        sink.clear_temperature();
        assert_eq!(sink.n_equations(), 1);
    }

    /// Modelica MassFlowSource_T: Fixed ṁ + T, Free P (not both ṁ and P).
    #[test]
    fn test_brine_source_imposed_mass_flow() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 2.0e5, 60_000.0);
        let mut source = BrineSource::new(
            "Glycol",
            Pressure::from_pascals(2.0e5),
            Temperature::from_celsius(15.0),
            Concentration::from_percent(30.0),
            backend,
            port,
        )
        .unwrap()
        .with_imposed_mass_flow(0.8)
        .unwrap();
        assert!(!source.imposes_pressure());
        assert_eq!(source.n_equations(), 2); // h + m

        // Wire the outlet edge at (m=0, p=1, h=2).
        source.set_system_context(0, &[(0, 1, 2)]);
        let h_set = 3500.0 * 15.0;
        let state = vec![0.8, 2.0e5, h_set];
        let mut r = vec![0.0; 2];
        source.compute_residuals(&state, &mut r).unwrap();
        assert!(r[1].abs() < 1e-12, "flow residual at set-point: {}", r[1]);

        let state2 = vec![0.5, 2.0e5, h_set];
        source.compute_residuals(&state2, &mut r).unwrap();
        assert!((r[1] + 0.3).abs() < 1e-12, "flow residual: {}", r[1]);

        // Jacobian: unit entry on the ṁ column (row 1 after h).
        let mut jac = JacobianBuilder::new();
        source.jacobian_entries(&state, &mut jac).unwrap();
        assert!(jac
            .entries()
            .iter()
            .any(|&(row, col, v)| row == 1 && col == 0 && (v - 1.0).abs() < 1e-12));
    }

    #[test]
    fn test_brine_boundary_rejects_invalid_imposed_flow() {
        let make_source = || {
            BrineSource::new(
                "Glycol",
                Pressure::from_pascals(2.0e5),
                Temperature::from_celsius(15.0),
                Concentration::from_percent(30.0),
                Arc::new(MockBrineBackend::new()),
                make_port("Glycol", 2.0e5, 60_000.0),
            )
            .unwrap()
        };
        assert!(make_source().with_imposed_mass_flow(0.0).is_err());
        assert!(make_source().with_imposed_mass_flow(f64::NAN).is_err());
    }

    /// Modelica Boundary_pT: Fixed P+T, Free ṁ (ΔT rating via sink T_out).
    #[test]
    fn test_brine_source_free_mass_flow_and_pressure() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 2.0e5, 60_000.0);
        let mut source = BrineSource::new(
            "Glycol",
            Pressure::from_pascals(2.0e5),
            Temperature::from_celsius(15.0),
            Concentration::from_percent(30.0),
            backend,
            port,
        )
        .unwrap()
        .with_impose_pressure(false);
        // Free P, Fixed T → 1 equation
        assert_eq!(source.n_equations(), 1);
        assert!(!source.imposes_pressure());
        assert!(source.imposes_temperature());

        source.set_system_context(0, &[(0, 1, 2)]);
        let h_set = 3500.0 * 15.0;
        let state = vec![0.5, 1.5e5, h_set];
        let mut r = vec![0.0; 1];
        source.compute_residuals(&state, &mut r).unwrap();
        assert!(r[0].abs() < 1e-12, "only h residual expected: {}", r[0]);

        // Boundary_pT + free flow stays at 2 eqs
        let source_pt = BrineSource::new(
            "Glycol",
            Pressure::from_pascals(2.0e5),
            Temperature::from_celsius(15.0),
            Concentration::from_percent(30.0),
            Arc::new(MockBrineBackend::new()),
            make_port("Glycol", 2.0e5, 60_000.0),
        )
        .unwrap();
        assert_eq!(source_pt.n_equations(), 2);
        assert!(source_pt.m_flow_set().is_none());
    }

    #[test]
    fn test_brine_sink_imposed_mass_flow_free_temperature() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 2.0e5, 60_000.0);
        let mut sink = BrineSink::new(
            "Glycol",
            Pressure::from_pascals(2.0e5),
            None,
            None,
            backend,
            port,
        )
        .unwrap()
        .with_imposed_mass_flow(0.8)
        .unwrap();
        // 1 (back-pressure) + 1 (flow); temperature stays free.
        assert_eq!(sink.n_equations(), 2);

        sink.set_system_context(0, &[(0, 1, 2)]);
        let state = vec![0.8, 2.0e5, 60_000.0];
        let mut r = vec![0.0; 2];
        sink.compute_residuals(&state, &mut r).unwrap();
        assert!(r[0].abs() < 1e-12 && r[1].abs() < 1e-12, "{r:?}");

        // Flow row occupies the row AFTER the optional h row: here row 1.
        let mut jac = JacobianBuilder::new();
        sink.jacobian_entries(&state, &mut jac).unwrap();
        assert!(jac
            .entries()
            .iter()
            .any(|&(row, col, v)| row == 1 && col == 0 && (v - 1.0).abs() < 1e-12));
    }

    // Task 4.3 — Residual validation: residuals must be zero at set-point.
    #[test]
    fn test_brine_source_residuals_zero_at_setpoint() {
        let p_pa = 3.0e5_f64;
        let t_c = 20.0_f64;
        let backend = Arc::new(MockBrineBackend::new());

        // The mock computes h = 3500 * (T_celsius) J/kg
        let h_expected = 3500.0 * t_c;
        let port = make_port("Glycol", p_pa, h_expected);

        let source = BrineSource::new(
            "Glycol",
            Pressure::from_pascals(p_pa),
            Temperature::from_celsius(t_c),
            Concentration::from_percent(30.0),
            backend,
            port,
        )
        .unwrap();

        let state: Vec<f64> = vec![];
        let mut residuals = vec![0.0_f64; 2];
        source.compute_residuals(&state, &mut residuals).unwrap();

        // Port is initialised at set-point → residuals must be zero
        assert!(
            residuals[0].abs() < 1e-6,
            "Pressure residual should be zero at set-point, got {}",
            residuals[0]
        );
        assert!(
            residuals[1].abs() < 1e-6,
            "Enthalpy residual should be zero at set-point, got {}",
            residuals[1]
        );
    }

    #[test]
    fn test_brine_sink_residuals_zero_at_setpoint() {
        let p_pa = 2.0e5_f64;
        let t_c = 15.0_f64;
        let backend = Arc::new(MockBrineBackend::new());

        let h_expected = 3500.0 * t_c;
        let port = make_port("Glycol", p_pa, h_expected);

        let mut sink = BrineSink::new(
            "Glycol",
            Pressure::from_pascals(p_pa),
            None,
            None,
            backend,
            port,
        )
        .unwrap();

        let state: Vec<f64> = vec![];

        // 1-equation mode: only pressure residual
        let mut residuals_1 = vec![0.0_f64; 1];
        sink.compute_residuals(&state, &mut residuals_1).unwrap();
        assert!(
            residuals_1[0].abs() < 1e-6,
            "Pressure residual (1-eq) should be zero at set-point, got {}",
            residuals_1[0]
        );

        // 2-equation mode after setting temperature
        sink.set_temperature(
            Temperature::from_celsius(t_c),
            Concentration::from_percent(30.0),
        )
        .unwrap();

        let mut residuals_2 = vec![0.0_f64; 2];
        sink.compute_residuals(&state, &mut residuals_2).unwrap();
        assert!(
            residuals_2[0].abs() < 1e-6,
            "Pressure residual (2-eq) should be zero, got {}",
            residuals_2[0]
        );
        assert!(
            residuals_2[1].abs() < 1e-6,
            "Enthalpy residual (2-eq) should be zero, got {}",
            residuals_2[1]
        );
    }

    // Task 4.4 — Trait object compatibility: both components usable as Box<dyn Component>.
    #[test]
    fn test_brine_source_trait_object() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 3.0e5, 70_000.0);
        let source = BrineSource::new(
            "Glycol",
            Pressure::from_pascals(3.0e5),
            Temperature::from_celsius(20.0),
            Concentration::from_percent(30.0),
            backend,
            port,
        )
        .unwrap();

        let boxed: Box<dyn Component> = Box::new(source);
        assert_eq!(boxed.n_equations(), 2);
        assert!(boxed.get_ports().is_empty());
        assert_eq!(
            boxed.energy_transfers(&[]),
            Some((Power::from_watts(0.0), Power::from_watts(0.0)))
        );
    }

    #[test]
    fn test_brine_sink_trait_object() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 2.0e5, 60_000.0);
        let sink = BrineSink::new(
            "Glycol",
            Pressure::from_pascals(2.0e5),
            None,
            None,
            backend,
            port,
        )
        .unwrap();

        let boxed: Box<dyn Component> = Box::new(sink);
        assert_eq!(boxed.n_equations(), 1);
        assert!(boxed.get_ports().is_empty());
        assert_eq!(
            boxed.energy_transfers(&[]),
            Some((Power::from_watts(0.0), Power::from_watts(0.0)))
        );
    }

    // Task 4.5 — Energy methods: Q=0, W=0 for boundary components.
    #[test]
    fn test_brine_source_energy_transfers_zero() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("Glycol", 3.0e5, 70_000.0);
        let source = BrineSource::new(
            "Glycol",
            Pressure::from_pascals(3.0e5),
            Temperature::from_celsius(20.0),
            Concentration::from_percent(0.0),
            backend,
            port,
        )
        .unwrap();

        let transfers = source.energy_transfers(&[]);
        assert_eq!(
            transfers,
            Some((Power::from_watts(0.0), Power::from_watts(0.0)))
        );
    }

    #[test]
    fn test_brine_source_accepts_meg_fluid() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("MEG", 3.0e5, 70_000.0);
        let result = BrineSource::new(
            "MEG",
            Pressure::from_pascals(3.0e5),
            Temperature::from_celsius(20.0),
            Concentration::from_percent(30.0),
            backend,
            port,
        );
        assert!(
            result.is_ok(),
            "MEG should be accepted as an incompressible fluid"
        );
    }

    #[test]
    fn test_brine_source_accepts_peg_fluid() {
        let backend = Arc::new(MockBrineBackend::new());
        let port = make_port("PEG", 3.0e5, 70_000.0);
        let result = BrineSource::new(
            "PEG",
            Pressure::from_pascals(3.0e5),
            Temperature::from_celsius(20.0),
            Concentration::from_percent(40.0),
            backend,
            port,
        );
        assert!(
            result.is_ok(),
            "PEG should be accepted as an incompressible fluid"
        );
    }
}
