//! Air Boundary Condition Components
//!
//! This module provides [`AirSource`] and [`AirSink`] components for moist-air
//! (humid air) heat transfer circuits with native psychrometric property support
//! via the [`RelativeHumidity`](entropyk_core::RelativeHumidity) NewType.
//!
//! ## Design
//!
//! Unlike the generic [`FlowSource`](crate::FlowSource)/[`FlowSink`](crate::FlowSink)
//! which use (Pressure, Enthalpy), these components use psychrometric state
//! specifications for type-safe air state definition:
//!
//! - **`AirSource`**: (T_dry, RelativeHumidity, Pressure) or (T_dry, T_wet_bulb, Pressure)
//! - **`AirSink`**: (Pressure, optional T_back)
//!
//! ## Psychrometric Formulas
//!
//! **Saturation vapour pressure** (Magnus-Tetens approximation):
//! ```text
//! P_sat = 610.78 · exp(17.27 · T_celsius / (T_celsius + 237.3))   [Pa]
//! ```
//!
//! **Humidity ratio** (kg water vapour / kg dry air):
//! ```text
//! W = 0.622 · P_v / (P_atm − P_v)
//! where P_v = RH · P_sat
//! ```
//!
//! **Specific enthalpy of moist air** (J/kg dry air):
//! ```text
//! h = (1006 · T_celsius + W · (2_501_000 + 1860 · T_celsius))   [J/kg]
//! ```
//!
//! ## Equations injected into the solver
//!
//! **AirSource** (always 2):
//! ```text
//! r₀ = P_edge − P_set = 0
//! r₁ = h_edge − h(T_dry, RH, P_set) = 0
//! ```
//!
//! **AirSink** (1 or 2 depending on whether a return temperature is set):
//! ```text
//! r₀ = P_edge − P_back = 0                     (always)
//! r₁ = h_edge − h(T_back, RH_back, P_back) = 0 (only when return temperature is set)
//! ```
//!
//! ## Example
//!
//! ```ignore
//! use entropyk_components::{AirSource, AirSink};
//! use entropyk_core::{Pressure, Temperature, RelativeHumidity};
//!
//! // Outdoor air at 35 °C / 50% RH / 101.325 kPa
//! let source = AirSource::from_dry_bulb_rh(
//!     Temperature::from_celsius(35.0),
//!     RelativeHumidity::from_percent(50.0),
//!     Pressure::from_pascals(101_325.0),
//!     outlet_port,
//! ).unwrap();
//!
//! // Free-enthalpy air sink
//! let sink = AirSink::new(Pressure::from_pascals(101_325.0), inlet_port).unwrap();
//! ```

use crate::{
    Component, ComponentError, ConnectedPort, JacobianBuilder, ResidualVector, StateSlice,
};
use entropyk_core::{Enthalpy, MassFlow, Power, Pressure, RelativeHumidity, Temperature};

// ─────────────────────────────────────────────────────────────────────────────
// Psychrometric helper functions
// ─────────────────────────────────────────────────────────────────────────────

/// Saturation vapour pressure using the Magnus-Tetens approximation.
///
/// Returns pressure in Pa.
///
/// ## Formula
/// ```text
/// P_sat = 610.78 · exp(17.27 · T_c / (T_c + 237.3))
/// ```
fn saturation_vapor_pressure(t: Temperature) -> Pressure {
    let t_c = t.to_celsius();
    let p_sat = 610.78 * ((17.27 * t_c) / (t_c + 237.3)).exp();
    Pressure::from_pascals(p_sat)
}

/// Humidity ratio (kg water vapour / kg dry air) from relative humidity.
///
/// ## Formula
/// ```text
/// W = 0.622 · P_v / (P_atm − P_v)
/// where P_v = RH · P_sat
/// ```
///
/// # Errors
///
/// Returns [`ComponentError::InvalidState`] if P_v ≥ P_atm (physically impossible condition).
fn humidity_ratio_from_rh(
    rh: RelativeHumidity,
    t_dry: Temperature,
    p_atm: Pressure,
) -> Result<f64, ComponentError> {
    let p_sat = saturation_vapor_pressure(t_dry);
    let p_v = p_sat.to_pascals() * rh.to_fraction();
    let p_atm_pa = p_atm.to_pascals();
    if p_v >= p_atm_pa {
        return Err(ComponentError::InvalidState(format!(
            "Vapor pressure ({:.2} Pa) exceeds atmospheric pressure ({:.2} Pa). Check temperature and RH values.",
            p_v, p_atm_pa
        )));
    }
    Ok(0.622 * p_v / (p_atm_pa - p_v))
}

/// Specific enthalpy of moist air (J per kg dry air).
///
/// ## Formula
/// ```text
/// h = 1006 · T_c + W · (2_501_000 + 1860 · T_c)   [J/kg_da]
/// ```
fn specific_enthalpy_from_w(t_dry: Temperature, w: f64) -> Enthalpy {
    let t_c = t_dry.to_celsius();
    let h = 1006.0 * t_c + w * (2_501_000.0 + 1860.0 * t_c);
    Enthalpy::from_joules_per_kg(h)
}

/// Relative humidity computed from dry-bulb and wet-bulb temperatures.
///
/// Uses the Sprung psychrometric formula for **unventilated (sling) psychrometers**:
/// ```text
/// e  = e_wb − A_psy · (T_dry − T_wb) · P_atm
/// RH = e / e_sat(T_dry)
/// where A_psy = 6.6e-4 /°C  (Sprung constant for unventilated psychrometer)
/// ```
///
/// # Important
/// This formula assumes an **unventilated (sling) psychrometer**. For ventilated psychrometers
/// (e.g., Assmann type), the psychrometric constant differs (A_psy ≈ 6.0e-4 /°C).
fn rh_from_wet_bulb(
    t_dry: Temperature,
    t_wet: Temperature,
    p_atm: Pressure,
) -> Result<RelativeHumidity, ComponentError> {
    let t_dry_c = t_dry.to_celsius();
    let t_wet_c = t_wet.to_celsius();

    if t_wet_c > t_dry_c + 1e-9 {
        return Err(ComponentError::InvalidState(
            "AirSource: wet-bulb temperature cannot exceed dry-bulb temperature".into(),
        ));
    }

    let e_sat_dry = saturation_vapor_pressure(t_dry).to_pascals();
    let e_sat_wet = saturation_vapor_pressure(t_wet).to_pascals();

    // Sprung psychrometric formula
    let a_psy = 6.6e-4; // 1/°C
    let e = e_sat_wet - a_psy * (t_dry_c - t_wet_c) * p_atm.to_pascals();

    let rh = (e / e_sat_dry).clamp(0.0, 1.0);
    Ok(RelativeHumidity::from_fraction(rh))
}

// ─────────────────────────────────────────────────────────────────────────────
// AirSource
// ─────────────────────────────────────────────────────────────────────────────

/// A boundary source that imposes fixed psychrometric state on its outlet edge
/// (dry-bulb temperature + relative humidity + atmospheric pressure).
///
/// Contributes **2 equations** to the system:
/// - `r₀ = P_edge − P_set = 0`
/// - `r₁ = h_edge − h(T_dry, RH, P_set) = 0`
///
/// Only models the *dry-air side* of heat exchangers (evaporator coils, condenser
/// coils, air-to-water heat pumps, etc.).
pub struct AirSource {
    /// Dry-bulb temperature [K]
    t_dry_k: f64,
    /// Relative humidity [0, 1]
    rh: RelativeHumidity,
    /// Atmospheric pressure [Pa]
    p_set_pa: f64,
    /// Pre-computed humidity ratio [kg_vap / kg_da]
    w: f64,
    /// Pre-computed specific enthalpy [J / kg_da]
    h_set_jkg: f64,
    /// Connected outlet port
    outlet: ConnectedPort,
    /// Outlet edge (ṁ, P, h) state indices, wired by `set_system_context`.
    /// When present the residuals/Jacobian are state-driven (same semantics as
    /// [`BrineSource`](crate::BrineSource)); otherwise the legacy port-frozen
    /// behaviour is kept for standalone unit tests.
    outlet_m_idx: Option<usize>,
    outlet_p_idx: Option<usize>,
    outlet_h_idx: Option<usize>,
    /// Optional imposed air mass flow [kg/s]: adds `r = ṁ_edge − ṁ_set`.
    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(T_dry, w)`.
    impose_temperature: bool,
}

impl AirSource {
    /// Creates an `AirSource` from dry-bulb temperature and relative humidity.
    ///
    /// # Arguments
    ///
    /// * `t_dry`     — Dry-bulb temperature
    /// * `rh`        — Relative humidity (automatically clamped to [0, 100 %])
    /// * `p_atm`     — Atmospheric (set-point) pressure
    /// * `outlet`    — Already-connected outlet port
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if `p_atm` ≤ 0.
    pub fn from_dry_bulb_rh(
        t_dry: Temperature,
        rh: RelativeHumidity,
        p_atm: Pressure,
        outlet: ConnectedPort,
    ) -> Result<Self, ComponentError> {
        let p_set_pa = p_atm.to_pascals();
        if p_set_pa <= 0.0 {
            return Err(ComponentError::InvalidState(
                "AirSource: atmospheric pressure must be positive".into(),
            ));
        }

        let w = humidity_ratio_from_rh(rh, t_dry, p_atm)?;
        let h_set_jkg = specific_enthalpy_from_w(t_dry, w).to_joules_per_kg();

        Ok(Self {
            t_dry_k: t_dry.to_kelvin(),
            rh,
            p_set_pa,
            w,
            h_set_jkg,
            outlet,
            outlet_m_idx: None,
            outlet_p_idx: None,
            outlet_h_idx: None,
            m_flow_set_kg_s: None,
            impose_pressure: true,
            impose_temperature: true,
        })
    }

    /// Creates an `AirSource` from dry-bulb and wet-bulb temperatures.
    ///
    /// The relative humidity is computed automatically via the Sprung psychrometric
    /// formula.  This constructor is convenient when measurements come from a
    /// psychrometer (sling hygrometer) rather than a capacitive RH sensor.
    ///
    /// # Arguments
    ///
    /// * `t_dry`     — Dry-bulb temperature
    /// * `t_wet`     — Wet-bulb temperature (must be ≤ `t_dry`)
    /// * `p_atm`     — Atmospheric (set-point) pressure
    /// * `outlet`    — Already-connected outlet port
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if:
    /// - `p_atm` ≤ 0
    /// - `t_wet` > `t_dry`
    pub fn from_dry_and_wet_bulb(
        t_dry: Temperature,
        t_wet: Temperature,
        p_atm: Pressure,
        outlet: ConnectedPort,
    ) -> Result<Self, ComponentError> {
        let p_set_pa = p_atm.to_pascals();
        if p_set_pa <= 0.0 {
            return Err(ComponentError::InvalidState(
                "AirSource: atmospheric pressure must be positive".into(),
            ));
        }

        let rh = rh_from_wet_bulb(t_dry, t_wet, p_atm)?;
        let w = humidity_ratio_from_rh(rh, t_dry, p_atm)?;
        let h_set_jkg = specific_enthalpy_from_w(t_dry, w).to_joules_per_kg();

        Ok(Self {
            t_dry_k: t_dry.to_kelvin(),
            rh,
            p_set_pa,
            w,
            h_set_jkg,
            outlet,
            outlet_m_idx: None,
            outlet_p_idx: None,
            outlet_h_idx: None,
            m_flow_set_kg_s: None,
            impose_pressure: true,
            impose_temperature: true,
        })
    }

    /// Imposes the air loop mass flow at the boundary (+1 equation).
    ///
    /// Do NOT combine with another flow imposition in the same branch
    /// (e.g. a `Fan` design flow) or the loop becomes over-determined.
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if the flow is not finite and positive.
    /// Modelica `MassFlowSource_T`: Fixed ṁ clears pressure imposition (Free P).
    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(
                "AirSource: 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 ṁ).
    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 dry-bulb temperature.
    pub fn t_dry(&self) -> Temperature {
        Temperature::from_kelvin(self.t_dry_k)
    }

    /// Returns the relative humidity.
    pub fn rh(&self) -> RelativeHumidity {
        self.rh
    }

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

    /// Returns the humidity ratio (kg water vapour / kg dry air).
    pub fn humidity_ratio(&self) -> f64 {
        self.w
    }

    /// Returns the pre-computed specific enthalpy of moist air.
    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 dry-bulb temperature and recomputes cached psychrometric properties.
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if the pressure becomes invalid during recomputation.
    pub fn set_temperature(&mut self, t_dry: Temperature) -> Result<(), ComponentError> {
        self.t_dry_k = t_dry.to_kelvin();
        self.recompute()
    }

    /// Updates the relative humidity and recomputes cached psychrometric properties.
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if the pressure becomes invalid during recomputation.
    pub fn set_rh(&mut self, rh: RelativeHumidity) -> Result<(), ComponentError> {
        self.rh = rh;
        self.recompute()
    }

    fn recompute(&mut self) -> Result<(), ComponentError> {
        if self.p_set_pa <= 0.0 {
            return Err(ComponentError::InvalidState(
                "AirSource: atmospheric pressure must be positive for psychrometric calculations"
                    .into(),
            ));
        }
        let t_dry = Temperature::from_kelvin(self.t_dry_k);
        let p_atm = Pressure::from_pascals(self.p_set_pa);
        self.w = humidity_ratio_from_rh(self.rh, t_dry, p_atm)?;
        self.h_set_jkg = specific_enthalpy_from_w(t_dry, self.w).to_joules_per_kg();
        Ok(())
    }
}

impl Component for AirSource {
    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(
                    "AirSource 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(
                    "AirSource 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> {
        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(
                    "AirSource 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!(
            "AirSource(P={:.0}Pa,T={:.1}K,RH={:.0}%)",
            self.p_set_pa,
            self.t_dry_k,
            self.rh.to_percent()
        )
    }

    fn to_params(&self) -> crate::ComponentParams {
        crate::ComponentParams::new("AirSource")
            .with_param("pSetPa", self.p_set_pa)
            .with_param("tDryK", self.t_dry_k)
            .with_param("rh", self.rh.to_fraction())
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// AirSink
// ─────────────────────────────────────────────────────────────────────────────

/// A boundary sink that imposes back-pressure, and optionally a fixed enthalpy
/// via (return dry-bulb temperature + relative humidity) on its inlet edge.
///
/// **Equation count is dynamic:**
/// - 1 equation (free enthalpy): `r₀ = P_edge − P_back = 0`
/// - 2 equations (return temperature set):
///   additionally `r₁ = h_edge − h(T_back, RH_back, P_back) = 0`
///
/// Toggle with [`set_return_temperature`](Self::set_return_temperature) /
/// [`clear_return_temperature`](Self::clear_return_temperature) without
/// rebuilding the component.
pub struct AirSink {
    /// Atmospheric back-pressure [Pa]
    p_back_pa: f64,
    /// Optional return dry-bulb temperature [K]
    t_back_k: Option<f64>,
    /// Optional return relative humidity
    rh_back: Option<RelativeHumidity>,
    /// Optional pre-computed return enthalpy [J/kg_da]
    h_back_jkg: Option<f64>,
    /// Connected inlet port
    inlet: ConnectedPort,
    /// Inlet edge (ṁ, P, h) state indices, wired by `set_system_context`
    /// (state-driven, same semantics as [`BrineSink`](crate::BrineSink)).
    inlet_m_idx: Option<usize>,
    inlet_p_idx: Option<usize>,
    inlet_h_idx: Option<usize>,
    /// Optional imposed air mass flow [kg/s].
    m_flow_set_kg_s: Option<f64>,
}

impl AirSink {
    /// Creates a new `AirSink` at the given back-pressure with free (unconstrained) enthalpy.
    ///
    /// # Arguments
    ///
    /// * `p_back`  — Back-pressure imposed on the inlet edge
    /// * `inlet`   — Already-connected inlet port
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if `p_back` ≤ 0.
    pub fn new(p_back: Pressure, inlet: ConnectedPort) -> Result<Self, ComponentError> {
        let p_back_pa = p_back.to_pascals();
        if p_back_pa <= 0.0 {
            return Err(ComponentError::InvalidState(
                "AirSink: back-pressure must be positive".into(),
            ));
        }
        Ok(Self {
            p_back_pa,
            t_back_k: None,
            rh_back: None,
            h_back_jkg: None,
            inlet,
            inlet_m_idx: None,
            inlet_p_idx: None,
            inlet_h_idx: None,
            m_flow_set_kg_s: None,
        })
    }

    /// Imposes the air loop mass flow at the boundary (+1 equation).
    ///
    /// Do NOT combine with another flow imposition in the same branch or the
    /// loop becomes over-determined.
    ///
    /// # Errors
    ///
    /// Returns [`ComponentError::InvalidState`] if the flow is not finite and positive.
    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(
                "AirSink: 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 back-pressure.
    pub fn p_back(&self) -> Pressure {
        Pressure::from_pascals(self.p_back_pa)
    }

    /// Returns the optional return dry-bulb temperature.
    pub fn t_back(&self) -> Option<Temperature> {
        self.t_back_k.map(Temperature::from_kelvin)
    }

    /// Returns the optional return relative humidity.
    pub fn rh_back(&self) -> Option<RelativeHumidity> {
        self.rh_back
    }

    /// Returns the optional pre-computed return 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
    }

    /// Sets a return temperature and relative humidity constraint, switching the
    /// sink to 2-equation mode.  The return enthalpy is recomputed immediately.
    ///
    /// # Arguments
    ///
    /// * `t_back` — Return dry-bulb temperature
    /// * `rh`     — Return relative humidity
    ///
    /// # Errors
    ///
    /// Currently infallible; returns `Ok(())`. Reserved for future validation.
    pub fn set_return_temperature(
        &mut self,
        t_back: Temperature,
        rh: RelativeHumidity,
    ) -> Result<(), ComponentError> {
        let p_atm = Pressure::from_pascals(self.p_back_pa);
        let w = humidity_ratio_from_rh(rh, t_back, p_atm)?;
        let h = specific_enthalpy_from_w(t_back, w);

        self.t_back_k = Some(t_back.to_kelvin());
        self.rh_back = Some(rh);
        self.h_back_jkg = Some(h.to_joules_per_kg());
        Ok(())
    }

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

impl Component for AirSink {
    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::with_capacity(self.n_equations());
        roles.push(crate::EquationRole::BoundaryDirichlet { quantity: "P" });
        if self.h_back_jkg.is_some() {
            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(
                    "AirSink 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(
                    "AirSink 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(
                    "AirSink 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_back_k {
            Some(t_k) => format!(
                "AirSink(P={:.0}Pa,T={:.1}K,RH={:.0}%)",
                self.p_back_pa,
                t_k,
                self.rh_back.map_or(0.0, |rh| rh.to_percent())
            ),
            None => format!("AirSink(P={:.0}Pa,T=free)", self.p_back_pa),
        }
    }

    fn to_params(&self) -> crate::ComponentParams {
        let mut params =
            crate::ComponentParams::new("AirSink").with_param("pBackPa", self.p_back_pa);
        if let Some(t_k) = self.t_back_k {
            params = params.with_param("tBackK", t_k);
        }
        if let Some(rh) = self.rh_back {
            params = params.with_param("rhBack", rh.to_fraction());
        }
        params
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests
// ─────────────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::port::{FluidId, Port};

    // ── Helper ────────────────────────────────────────────────────────────────

    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
    }

    // Standard atmospheric pressure used throughout tests
    const P_ATM_PA: f64 = 101_325.0;

    // ── Acceptance Criteria #1: AirSource::from_dry_bulb_rh ──────────────────

    #[test]
    fn test_air_source_from_dry_bulb_rh() {
        // AC#1: AirSource::from_dry_bulb_rh() creates a source with T_dry and RH
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let result = AirSource::from_dry_bulb_rh(
            Temperature::from_celsius(25.0),
            RelativeHumidity::from_percent(60.0),
            Pressure::from_pascals(P_ATM_PA),
            port,
        );
        assert!(result.is_ok(), "Should create AirSource from dry-bulb + RH");
        let source = result.unwrap();
        assert_eq!(source.n_equations(), 2);
        assert!((source.rh().to_percent() - 60.0).abs() < 1e-9);
        assert!((source.t_dry().to_celsius() - 25.0).abs() < 1e-9);
    }

    // ── Acceptance Criteria #2: AirSource::from_dry_and_wet_bulb ─────────────

    #[test]
    fn test_air_source_from_wet_bulb() {
        // AC#2: from_dry_and_wet_bulb() computes RH from wet-bulb temperature
        // At 25 °C dry / 18 °C wet, expected RH ≈ 50-55 %
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let result = AirSource::from_dry_and_wet_bulb(
            Temperature::from_celsius(25.0),
            Temperature::from_celsius(18.0),
            Pressure::from_pascals(P_ATM_PA),
            port,
        );
        assert!(result.is_ok(), "Should create AirSource from dry/wet bulb");
        let source = result.unwrap();
        // RH must be in a physically plausible range
        let rh_pct = source.rh().to_percent();
        assert!(
            rh_pct > 30.0 && rh_pct < 80.0,
            "RH = {rh_pct:.1}% should be ~50-55%"
        );
    }

    #[test]
    fn test_air_source_wet_bulb_exceeds_dry_bulb_returns_error() {
        // Wet-bulb > dry-bulb is physically impossible → must return error
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let result = AirSource::from_dry_and_wet_bulb(
            Temperature::from_celsius(20.0),
            Temperature::from_celsius(25.0), // wet > dry
            Pressure::from_pascals(P_ATM_PA),
            port,
        );
        assert!(result.is_err(), "Should reject wet-bulb > dry-bulb");
    }

    // ── Acceptance Criteria #3: specific_enthalpy() ───────────────────────────

    #[test]
    fn test_specific_enthalpy_calculation() {
        // AC#3: specific_enthalpy() returns moist air enthalpy
        // ASHRAE reference: at 25°C, W ≈ 0.01 kg/kg_da → h ≈ 50.5 kJ/kg_da
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let source = AirSource::from_dry_bulb_rh(
            Temperature::from_celsius(25.0),
            RelativeHumidity::from_percent(50.0),
            Pressure::from_pascals(P_ATM_PA),
            port,
        )
        .unwrap();
        // Stricter range: ASHRAE suggests h ≈ 50.5 kJ/kg for 25°C/50%RH
        // Tolerance: ±5 kJ/kg (10% relative error)
        let h_jkg = source.h_set().to_joules_per_kg();
        assert!(
            h_jkg > 45_000.0 && h_jkg < 56_000.0,
            "h = {:.0} J/kg should be ~50,500 J/kg for 25°C/50%RH (tolerance: ±5,500)",
            h_jkg
        );
    }

    #[test]
    fn test_specific_enthalpy_dry_air() {
        // At 0°C and 0% RH, h = 1006*0 + 0*(…) = 0 J/kg_da
        let port = make_port("Air", P_ATM_PA, 0.0);
        let source = AirSource::from_dry_bulb_rh(
            Temperature::from_celsius(0.0),
            RelativeHumidity::from_percent(0.0),
            Pressure::from_pascals(P_ATM_PA),
            port,
        )
        .unwrap();
        let h = source.h_set().to_joules_per_kg();
        assert!(
            (h).abs() < 1e-6,
            "h should be 0 J/kg for 0°C dry air, got {h}"
        );
    }

    // ── Acceptance Criteria #4: humidity_ratio() ─────────────────────────────

    #[test]
    fn test_humidity_ratio_calculation() {
        // AC#4: humidity_ratio() returns W in kg_vap / kg_da
        // ASHRAE: at 25°C, 60% RH → W ≈ 0.0119 kg/kg_da
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let source = AirSource::from_dry_bulb_rh(
            Temperature::from_celsius(25.0),
            RelativeHumidity::from_percent(60.0),
            Pressure::from_pascals(P_ATM_PA),
            port,
        )
        .unwrap();
        let w = source.humidity_ratio();
        // Expect W in [0.009, 0.016] for 25°C / 60% RH at sea level
        assert!(
            w > 0.009 && w < 0.016,
            "W = {w:.5} should be ~0.0119 for 25°C/60%RH"
        );
    }

    // ── Acceptance Criteria #5: AirSink::new() ───────────────────────────────

    #[test]
    fn test_air_sink_creation() {
        // AC#5: AirSink::new() creates a sink at atmospheric pressure
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let sink = AirSink::new(Pressure::from_pascals(P_ATM_PA), port).unwrap();
        assert_eq!(
            sink.n_equations(),
            1,
            "Free-enthalpy sink should have 1 equation"
        );
        assert!((sink.p_back().to_pascals() - P_ATM_PA).abs() < 1e-6);
    }

    #[test]
    fn test_air_sink_invalid_pressure() {
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let result = AirSink::new(Pressure::from_pascals(-1.0), port);
        assert!(result.is_err(), "Should reject non-positive pressure");
    }

    // ── Acceptance Criteria #6: energy_transfers() ────────────────────────────

    #[test]
    fn test_air_source_energy_transfers_zero() {
        // AC#6: energy_transfers() returns (Power(0), Power(0))
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let source = AirSource::from_dry_bulb_rh(
            Temperature::from_celsius(20.0),
            RelativeHumidity::from_percent(50.0),
            Pressure::from_pascals(P_ATM_PA),
            port,
        )
        .unwrap();
        assert_eq!(
            source.energy_transfers(&[]),
            Some((Power::from_watts(0.0), Power::from_watts(0.0)))
        );
    }

    #[test]
    fn test_air_sink_energy_transfers_zero() {
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let sink = AirSink::new(Pressure::from_pascals(P_ATM_PA), port).unwrap();
        assert_eq!(
            sink.energy_transfers(&[]),
            Some((Power::from_watts(0.0), Power::from_watts(0.0)))
        );
    }

    // ── Acceptance Criteria #7: RH validation ────────────────────────────────

    #[test]
    fn test_rh_clamped_to_valid_range() {
        // AC#7: RelativeHumidity type clamps to [0, 100%] automatically
        let rh_over = RelativeHumidity::from_percent(150.0);
        assert!(
            (rh_over.to_percent() - 100.0).abs() < 1e-9,
            "RH clamped to 100%"
        );

        let rh_under = RelativeHumidity::from_percent(-10.0);
        assert!((rh_under.to_percent()).abs() < 1e-9, "RH clamped to 0%");

        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let source = AirSource::from_dry_bulb_rh(
            Temperature::from_celsius(25.0),
            rh_over,
            Pressure::from_pascals(P_ATM_PA),
            port,
        );
        assert!(source.is_ok(), "Clamped RH = 100% should still be valid");
    }

    // ── Acceptance Criteria #8: ASHRAE reference values ──────────────────────

    #[test]
    fn test_saturation_vapor_pressure_ashrae() {
        // AC#8: Test psychrometric calculations against ASHRAE reference values
        // ASHRAE Fundamentals 2021, Table 2: P_sat at 25°C = 3169 Pa
        let p_sat = saturation_vapor_pressure(Temperature::from_celsius(25.0));
        let p_sat_pa = p_sat.to_pascals();
        // Stricter tolerance: 0.5% relative error (Magnus-Tetens accuracy)
        assert!(
            (p_sat_pa - 3169.0).abs() < 15.845,
            "P_sat at 25°C = {p_sat_pa:.1} Pa, expected ~3169 Pa (ASHRAE, tolerance: ±15.8 Pa)"
        );
    }

    #[test]
    fn test_humidity_ratio_ashrae_reference() {
        // ASHRAE: at 20°C, 50% RH → W ≈ 0.00726 kg/kg_da
        let t = Temperature::from_celsius(20.0);
        let rh = RelativeHumidity::from_percent(50.0);
        let p_atm = Pressure::from_pascals(101_325.0);
        let w = humidity_ratio_from_rh(rh, t, p_atm).unwrap();
        // Stricter tolerance: 1% relative error (ASHRAE standard precision)
        assert!(
            (w - 0.00726).abs() < 0.000073,
            "W at 20°C/50%RH = {w:.5}, expected ~0.00726 (tolerance: ±0.000073)"
        );
    }

    #[test]
    fn test_specific_enthalpy_ashrae_reference() {
        // ASHRAE: at 20°C, W = 0.00726 → h = 20.12 + 0.00726*(2501 + 1.86*20)*1000 ≈ 38.6 kJ/kg
        let t = Temperature::from_celsius(20.0);
        let w = 0.00726;
        let h = specific_enthalpy_from_w(t, w);
        let h_kjkg = h.to_joules_per_kg() / 1000.0;
        // Stricter tolerance: 1% relative error
        assert!(
            (h_kjkg - 38.6).abs() < 0.386,
            "h at 20°C/W=0.00726 = {h_kjkg:.3} kJ/kg, expected ~38.6 kJ/kg (tolerance: ±0.386)"
        );
    }

    // ── Residual validation ───────────────────────────────────────────────────

    #[test]
    fn test_air_source_residuals_zero_at_setpoint() {
        let t_c = 25.0_f64;
        let rh = RelativeHumidity::from_percent(50.0);
        let p_pa = P_ATM_PA;
        let p_atm = Pressure::from_pascals(p_pa);

        // Compute expected h using same helper
        let w = humidity_ratio_from_rh(rh, Temperature::from_celsius(t_c), p_atm).unwrap();
        let h_jkg = specific_enthalpy_from_w(Temperature::from_celsius(t_c), w).to_joules_per_kg();

        let port = make_port("Air", p_pa, h_jkg);
        let source =
            AirSource::from_dry_bulb_rh(Temperature::from_celsius(t_c), rh, p_atm, port).unwrap();

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

        assert!(
            residuals[0].abs() < 1e-6,
            "Pressure residual should be zero, got {}",
            residuals[0]
        );
        assert!(
            residuals[1].abs() < 1e-6,
            "Enthalpy residual should be zero, got {}",
            residuals[1]
        );
    }

    #[test]
    fn test_air_sink_residuals_zero_at_setpoint_free_mode() {
        let p_pa = P_ATM_PA;
        let port = make_port("Air", p_pa, 50_000.0);
        let sink = AirSink::new(Pressure::from_pascals(p_pa), port).unwrap();

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

        assert!(
            residuals[0].abs() < 1e-6,
            "Pressure residual (1-eq) should be zero, got {}",
            residuals[0]
        );
    }

    #[test]
    fn test_air_sink_residuals_zero_at_setpoint_constrained_mode() {
        let t_c = 15.0_f64;
        let rh = RelativeHumidity::from_percent(60.0);
        let p_pa = P_ATM_PA;
        let p_atm = Pressure::from_pascals(p_pa);

        let w = humidity_ratio_from_rh(rh, Temperature::from_celsius(t_c), p_atm).unwrap();
        let h_jkg = specific_enthalpy_from_w(Temperature::from_celsius(t_c), w).to_joules_per_kg();

        let port = make_port("Air", p_pa, h_jkg);
        let mut sink = AirSink::new(p_atm, port).unwrap();
        sink.set_return_temperature(Temperature::from_celsius(t_c), rh)
            .unwrap();

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

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

        assert!(
            residuals[0].abs() < 1e-6,
            "Pressure residual should be zero, got {}",
            residuals[0]
        );
        assert!(
            residuals[1].abs() < 1e-6,
            "Enthalpy residual should be zero, got {}",
            residuals[1]
        );
    }

    // ── Trait object compatibility ────────────────────────────────────────────

    #[test]
    fn test_air_source_trait_object() {
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let source = AirSource::from_dry_bulb_rh(
            Temperature::from_celsius(25.0),
            RelativeHumidity::from_percent(50.0),
            Pressure::from_pascals(P_ATM_PA),
            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_air_sink_trait_object() {
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let sink = AirSink::new(Pressure::from_pascals(P_ATM_PA), 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)))
        );
    }

    // ── Dynamic toggle ────────────────────────────────────────────────────────

    #[test]
    fn test_air_sink_dynamic_temperature_toggle() {
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let mut sink = AirSink::new(Pressure::from_pascals(P_ATM_PA), port).unwrap();

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

        sink.set_return_temperature(
            Temperature::from_celsius(15.0),
            RelativeHumidity::from_percent(60.0),
        )
        .unwrap();
        assert_eq!(sink.n_equations(), 2);

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

    // ── Signature ─────────────────────────────────────────────────────────────

    #[test]
    fn test_air_source_signature() {
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let source = AirSource::from_dry_bulb_rh(
            Temperature::from_celsius(25.0),
            RelativeHumidity::from_percent(60.0),
            Pressure::from_pascals(P_ATM_PA),
            port,
        )
        .unwrap();
        let sig = source.signature();
        assert!(
            sig.contains("AirSource"),
            "Signature should contain 'AirSource'"
        );
        assert!(sig.contains("P="), "Signature should contain pressure");
        assert!(sig.contains("T="), "Signature should contain temperature");
        assert!(sig.contains("RH="), "Signature should contain RH");
    }

    #[test]
    fn test_air_sink_signature_free() {
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let sink = AirSink::new(Pressure::from_pascals(P_ATM_PA), port).unwrap();
        let sig = sink.signature();
        assert!(sig.contains("AirSink"));
        assert!(sig.contains("T=free"));
    }

    #[test]
    fn test_air_sink_signature_constrained() {
        let port = make_port("Air", P_ATM_PA, 50_000.0);
        let mut sink = AirSink::new(Pressure::from_pascals(P_ATM_PA), port).unwrap();
        sink.set_return_temperature(
            Temperature::from_celsius(20.0),
            RelativeHumidity::from_percent(40.0),
        )
        .unwrap();
        let sig = sink.signature();
        assert!(sig.contains("AirSink"));
        assert!(sig.contains("T="));
        assert!(!sig.contains("T=free"));
    }
}
