//! Integration test for mass balance validation with multiple components.
//!
//! This test verifies that the mass balance validation works correctly
//! across a multi-component system simulating a refrigeration cycle.

use entropyk_components::port::{FluidId, Port};
use entropyk_components::{
    Component, ComponentError, ConnectedPort, JacobianBuilder, ResidualVector, StateSlice,
};
use entropyk_core::{Enthalpy, MassFlow, Pressure};
use entropyk_solver::system::System;

// ─────────────────────────────────────────────────────────────────────────────
// Mock components for testing
// ─────────────────────────────────────────────────────────────────────────────

/// A mock component that simulates balanced mass flow (like a pipe or heat exchanger).
struct BalancedComponent {
    ports: Vec<ConnectedPort>,
    mass_flow_in: f64,
}

impl BalancedComponent {
    fn new(ports: Vec<ConnectedPort>, mass_flow: f64) -> Self {
        Self {
            ports,
            mass_flow_in: mass_flow,
        }
    }
}

impl Component for BalancedComponent {
    fn compute_residuals(
        &self,
        _state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        for r in residuals.iter_mut() {
            *r = 0.0;
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        for i in 0..self.n_equations() {
            jacobian.add_entry(i, i, 1.0);
        }
        Ok(())
    }

    fn n_equations(&self) -> usize {
        2
    }

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

    fn port_mass_flows(&self, _state: &StateSlice) -> Result<Vec<MassFlow>, ComponentError> {
        // Balanced: inlet positive, outlet negative
        Ok(vec![
            MassFlow::from_kg_per_s(self.mass_flow_in),
            MassFlow::from_kg_per_s(-self.mass_flow_in),
        ])
    }
}

/// A mock component with imbalanced mass flow (for testing violation detection).
struct ImbalancedComponent {
    ports: Vec<ConnectedPort>,
}

impl ImbalancedComponent {
    fn new(ports: Vec<ConnectedPort>) -> Self {
        Self { ports }
    }
}

impl Component for ImbalancedComponent {
    fn compute_residuals(
        &self,
        _state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        for r in residuals.iter_mut() {
            *r = 0.0;
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        for i in 0..self.n_equations() {
            jacobian.add_entry(i, i, 1.0);
        }
        Ok(())
    }

    fn n_equations(&self) -> usize {
        2
    }

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

    fn port_mass_flows(&self, _state: &StateSlice) -> Result<Vec<MassFlow>, ComponentError> {
        // Imbalanced: inlet 1.0 kg/s, outlet -0.5 kg/s (sum = 0.5 kg/s violation)
        Ok(vec![
            MassFlow::from_kg_per_s(1.0),
            MassFlow::from_kg_per_s(-0.5),
        ])
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Test helpers
// ─────────────────────────────────────────────────────────────────────────────

fn make_connected_port_pair(
    fluid: &str,
    p_bar: f64,
    h_j_kg: f64,
) -> (ConnectedPort, ConnectedPort) {
    let p1 = Port::new(
        FluidId::new(fluid),
        Pressure::from_bar(p_bar),
        Enthalpy::from_joules_per_kg(h_j_kg),
    );
    let p2 = Port::new(
        FluidId::new(fluid),
        Pressure::from_bar(p_bar),
        Enthalpy::from_joules_per_kg(h_j_kg),
    );
    let (c1, c2) = p1.connect(p2).unwrap();
    (c1, c2)
}

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

#[test]
fn test_mass_balance_4_component_cycle() {
    // Simulate a 4-component refrigeration cycle: Compressor → Condenser → Valve → Evaporator
    let mut system = System::new();

    // Create 4 pairs of connected ports for 4 components
    let (p1a, p1b) = make_connected_port_pair("R134a", 5.0, 400_000.0);
    let (p2a, p2b) = make_connected_port_pair("R134a", 5.0, 400_000.0);
    let (p3a, p3b) = make_connected_port_pair("R134a", 5.0, 400_000.0);
    let (p4a, p4b) = make_connected_port_pair("R134a", 5.0, 400_000.0);

    // Create 4 balanced components (simulating compressor, condenser, valve, evaporator)
    let mass_flow = 0.1; // kg/s
    let comp1 = Box::new(BalancedComponent::new(vec![p1a, p1b], mass_flow));
    let comp2 = Box::new(BalancedComponent::new(vec![p2a, p2b], mass_flow));
    let comp3 = Box::new(BalancedComponent::new(vec![p3a, p3b], mass_flow));
    let comp4 = Box::new(BalancedComponent::new(vec![p4a, p4b], mass_flow));

    // Add components to system
    let n1 = system.add_component(comp1);
    let n2 = system.add_component(comp2);
    let n3 = system.add_component(comp3);
    let n4 = system.add_component(comp4);

    // Connect in a cycle
    system.add_edge(n1, n2).unwrap();
    system.add_edge(n2, n3).unwrap();
    system.add_edge(n3, n4).unwrap();
    system.add_edge(n4, n1).unwrap();

    system.finalize().unwrap();

    // Test with zero state vector
    let state = vec![0.0; system.full_state_vector_len()];
    let result = system.check_mass_balance(&state);

    assert!(
        result.is_ok(),
        "Mass balance should pass for balanced 4-component cycle"
    );
}

#[test]
fn test_mass_balance_detects_imbalance_in_cycle() {
    // Create a cycle with one imbalanced component
    let mut system = System::new();

    let (p1a, p1b) = make_connected_port_pair("R134a", 5.0, 400_000.0);
    let (p2a, p2b) = make_connected_port_pair("R134a", 5.0, 400_000.0);
    let (p3a, p3b) = make_connected_port_pair("R134a", 5.0, 400_000.0);

    // Two balanced components
    let comp1 = Box::new(BalancedComponent::new(vec![p1a, p1b], 0.1));
    let comp3 = Box::new(BalancedComponent::new(vec![p3a, p3b], 0.1));

    // One imbalanced component
    let comp2 = Box::new(ImbalancedComponent::new(vec![p2a, p2b]));

    let n1 = system.add_component(comp1);
    let n2 = system.add_component(comp2);
    let n3 = system.add_component(comp3);

    system.add_edge(n1, n2).unwrap();
    system.add_edge(n2, n3).unwrap();
    system.add_edge(n3, n1).unwrap();

    system.finalize().unwrap();

    let state = vec![0.0; system.full_state_vector_len()];
    let result = system.check_mass_balance(&state);

    assert!(
        result.is_err(),
        "Mass balance should fail when one component is imbalanced"
    );
}

#[test]
fn test_mass_balance_multiple_components_same_flow() {
    // Test that multiple components with the same mass flow pass validation
    let mut system = System::new();

    // Create 6 components in a chain
    let mut ports = Vec::new();
    for _ in 0..6 {
        let (pa, pb) = make_connected_port_pair("R134a", 5.0, 400_000.0);
        ports.push((pa, pb));
    }

    let mass_flow = 0.5; // kg/s
    let components: Vec<_> = ports
        .into_iter()
        .map(|(pa, pb)| Box::new(BalancedComponent::new(vec![pa, pb], mass_flow)))
        .collect();

    let nodes: Vec<_> = components
        .into_iter()
        .map(|c| system.add_component(c))
        .collect();

    // Connect in a cycle
    for i in 0..nodes.len() {
        let next = (i + 1) % nodes.len();
        system.add_edge(nodes[i], nodes[next]).unwrap();
    }

    system.finalize().unwrap();

    let state = vec![0.0; system.full_state_vector_len()];
    let result = system.check_mass_balance(&state);

    assert!(
        result.is_ok(),
        "Mass balance should pass for multiple balanced components"
    );
}

#[test]
fn test_mass_balance_tolerance_constant_accessible() {
    // Verify the tolerance constant is accessible
    assert_eq!(System::MASS_BALANCE_TOLERANCE_KG_S, 1e-9);
}