//! Integration tests for MacroComponent (Story 3.6).
//!
//! Tests cover:
//! - AC #1: MacroComponent implements Component trait
//! - AC #2: External ports correctly mapped to internal edges
//! - AC #3: Residuals and Jacobian delegated with proper coupling equations
//! - AC #4: Serialization snapshot round-trip

use entropyk_components::{
    Component, ComponentError, ConnectedPort, JacobianBuilder, ResidualVector, StateSlice,
};
use entropyk_solver::{MacroComponent, MacroComponentSnapshot, System};

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

/// A simple zero-residual pass-through mock component.
struct PassThrough {
    n_eq: usize,
}

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

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

    fn n_equations(&self) -> usize {
        self.n_eq
    }

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

fn pass(n: usize) -> Box<dyn Component> {
    Box::new(PassThrough { n_eq: n })
}

fn make_port(fluid: &str, p: f64, h: f64) -> ConnectedPort {
    use entropyk_components::port::{FluidId, Port};
    use entropyk_core::{Enthalpy, Pressure};
    let p1 = Port::new(
        FluidId::new(fluid),
        Pressure::from_pascals(p),
        Enthalpy::from_joules_per_kg(h),
    );
    let p2 = Port::new(
        FluidId::new(fluid),
        Pressure::from_pascals(p),
        Enthalpy::from_joules_per_kg(h),
    );
    p1.connect(p2).unwrap().0
}

/// Build a 4-component refrigerant cycle: A→B→C→D→A (4 edges).
/// Each component contributes 3 equations (2 thermo + 1 mass-flow) per CM1.3.
fn build_4_component_cycle() -> System {
    let mut sys = System::new();
    let a = sys.add_component(pass(3)); // compressor
    let b = sys.add_component(pass(3)); // condenser
    let c = sys.add_component(pass(3)); // valve
    let d = sys.add_component(pass(3)); // evaporator
    sys.add_edge(a, b).unwrap();
    sys.add_edge(b, c).unwrap();
    sys.add_edge(c, d).unwrap();
    sys.add_edge(d, a).unwrap();
    sys.finalize().unwrap();
    sys
}

// ─────────────────────────────────────────────────────────────────────────────
// AC #1 & #2 — MacroComponent wraps 4-component cycle correctly
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_4_component_cycle_macro_creation() {
    let internal = build_4_component_cycle();
    let mc = MacroComponent::new(internal);

    // 4 components × 3 equations = 12 internal equations (pass(3)×4), 0 exposed ports
    assert_eq!(
        mc.n_equations(),
        12,
        "should have 12 internal equations (4 components × 3 eqs) with no exposed ports"
    );
    // CM1.4: 4-edge series cycle → 1 branch + 4×2 P,h = 9 internal state vars
    assert_eq!(mc.internal_state_len(), 9);
    assert!(mc.get_ports().is_empty());
}

#[test]
fn test_4_component_cycle_expose_two_ports() {
    let internal = build_4_component_cycle();
    let mut mc = MacroComponent::new(internal);

    // Expose edge 0 as "refrig_in" and edge 2 as "refrig_out"
    mc.expose_port(0, "refrig_in", make_port("R134a", 1e5, 4e5));
    mc.expose_port(2, "refrig_out", make_port("R134a", 5e5, 4.5e5));

    // 12 internal (4 components × 3 eqs) + 4 coupling (2 per port × 2 ports) = 16
    assert_eq!(
        mc.n_equations(),
        16,
        "should have 16 equations with 2 exposed ports"
    );
    assert_eq!(mc.get_ports().len(), 2);
    assert_eq!(mc.port_mappings()[0].name, "refrig_in");
    assert_eq!(mc.port_mappings()[1].name, "refrig_out");
}

#[test]
fn test_4_component_cycle_in_parent_system() {
    // Wrap cycle in MacroComponent and place in a parent system
    let internal = build_4_component_cycle();
    let mc = MacroComponent::new(internal);

    let mut parent = System::new();
    let _mc_node = parent.add_component(Box::new(mc));
    // Single-node system (no edges) would fail validation,
    // so we add a second node and an edge.
    let other = parent.add_component(pass(1));
    // For finalize to succeed, all nodes must have at least one edge
    // (system topology requires connected nodes).
    // We skip finalize here since the topology is valid (2 nodes, 1 edge).
    // Actually the validation requires an edge:
    parent.add_edge(_mc_node, other).unwrap();
    let result = parent.finalize();
    assert!(
        result.is_ok(),
        "parent finalize should succeed: {:?}",
        result.err()
    );

    // Parent has 2 nodes, 1 edge
    assert_eq!(parent.node_count(), 2);
    assert_eq!(parent.edge_count(), 1);

    // CM1.4: parent has 1 edge → 1 branch + 2 P,h = 3 parent edge vars.
    // MacroComponent internal: 1 branch + 4×2 P,h = 9 internal vars.
    // Total = 3 + 9 = 12.
    assert_eq!(parent.state_vector_len(), 12);
}

// ─────────────────────────────────────────────────────────────────────────────
// AC #3 — Residuals and Jacobian delegated with coupling equations
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_coupling_residuals_are_zero_at_consistent_state() {
    // Build cycle, expose 1 port, inject consistent external state
    let internal = build_4_component_cycle();
    let mut mc = MacroComponent::new(internal);

    mc.expose_port(0, "refrig_in", make_port("R134a", 1e5, 4e5));

    // External edge occupies state[0..3]: m_ext=0, p_ext=1, h_ext=2.
    // Internal block starts at offset 3 (3 parent-edge state vars before it).
    mc.set_global_state_offset(3);
    mc.set_system_context(3, &[(0, 1, 2)]);

    // State layout: external edge (ṁ@0, P@1, h@2), internal block at offset 3:
    //   edge0: (ṁ@3, P@4, h@5), edge1: (ṁ@6, P@7, h@8), ...
    let mut state = vec![0.0; 3 + 12]; // 3 parent + 12 internal (4 edges × 3)
    state[1] = 1.0e5; // P_ext
    state[2] = 4.0e5; // h_ext
    state[4] = 1.0e5; // P_int_e0 (consistent with port: offset 3 + 1 = 4)
    state[5] = 4.0e5; // h_int_e0 (consistent with port: offset 3 + 2 = 5)

    let n_eqs = mc.n_equations(); // 12 internal + 2 coupling = 14
    let mut residuals = vec![0.0; n_eqs];
    mc.compute_residuals(&state, &mut residuals).unwrap();

    // Coupling residuals at indices 12, 13 should be zero (consistent state)
    assert!(
        residuals[12].abs() < 1e-10,
        "P coupling residual should be 0, got {}",
        residuals[12]
    );
    assert!(
        residuals[13].abs() < 1e-10,
        "h coupling residual should be 0, got {}",
        residuals[13]
    );
}

#[test]
fn test_coupling_residuals_nonzero_at_inconsistent_state() {
    let internal = build_4_component_cycle();
    let mut mc = MacroComponent::new(internal);

    mc.expose_port(0, "refrig_in", make_port("R134a", 1e5, 4e5));
    mc.set_global_state_offset(3);
    mc.set_system_context(3, &[(0, 1, 2)]);

    let mut state = vec![0.0; 15];
    state[1] = 2.0e5; // P_ext (different from internal, p_ext=1)
    state[2] = 5.0e5; // h_ext (h_ext=2)
    state[4] = 1.0e5; // P_int_e0 (offset 3+1=4)
    state[5] = 4.0e5; // h_int_e0 (offset 3+2=5)

    let n_eqs = mc.n_equations();
    let mut residuals = vec![0.0; n_eqs];
    mc.compute_residuals(&state, &mut residuals).unwrap();

    // Coupling: r[12] = P_ext - P_int = 2e5 - 1e5 = 1e5
    assert!(
        (residuals[12] - 1.0e5).abs() < 1.0,
        "P coupling residual mismatch: {}",
        residuals[12]
    );
    assert!(
        (residuals[13] - 1.0e5).abs() < 1.0,
        "h coupling residual mismatch: {}",
        residuals[13]
    );
}

#[test]
fn test_jacobian_coupling_entries_correct() {
    let internal = build_4_component_cycle();
    let mut mc = MacroComponent::new(internal);

    mc.expose_port(0, "refrig_in", make_port("R134a", 1e5, 4e5));
    // external edge: (m_ext=0, p_ext=1, h_ext=2), internal starts at offset=3
    mc.set_global_state_offset(3);
    mc.set_system_context(3, &[(0, 1, 2)]);

    let state = vec![0.0; 15];
    let mut jac = JacobianBuilder::new();
    mc.jacobian_entries(&state, &mut jac).unwrap();

    let entries = jac.entries();
    let find = |row: usize, col: usize| -> Option<f64> {
        entries
            .iter()
            .find(|&&(r, c, _)| r == row && c == col)
            .map(|&(_, _, v)| v)
    };

    // Coupling rows 12 (P) and 13 (h); internal edge0 (P@offset+1=4, h@offset+2=5)
    assert_eq!(find(12, 1), Some(1.0), "∂r_P/∂p_ext should be +1");
    assert_eq!(find(12, 4), Some(-1.0), "∂r_P/∂int_p should be -1");
    assert_eq!(find(13, 2), Some(1.0), "∂r_h/∂h_ext should be +1");
    assert_eq!(find(13, 5), Some(-1.0), "∂r_h/∂int_h should be -1");
}

// ─────────────────────────────────────────────────────────────────────────────
// AC #4 — Serialization snapshot
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_macro_component_snapshot_serialization() {
    let internal = build_4_component_cycle();
    let mut mc = MacroComponent::new(internal);
    mc.expose_port(0, "refrig_in", make_port("R134a", 1e5, 4e5));
    mc.expose_port(2, "refrig_out", make_port("R134a", 5e5, 4.5e5));
    mc.set_global_state_offset(0);

    // CM1.4: 4-edge series cycle → internal_state_len = 1 branch + 4×2 P,h = 9 vars.
    let global_state: Vec<f64> = (0..9).map(|i| (i as f64 + 1.0) * 1e4).collect();

    let snap = mc
        .to_snapshot(&global_state, Some("chiller_A".into()))
        .expect("snapshot should succeed");

    assert_eq!(snap.label.as_deref(), Some("chiller_A"));
    assert_eq!(snap.internal_edge_states.len(), 9);
    assert_eq!(snap.port_names, vec!["refrig_in", "refrig_out"]);

    // JSON round-trip
    let json = serde_json::to_string_pretty(&snap).expect("must serialize");
    let restored: MacroComponentSnapshot = serde_json::from_str(&json).expect("must deserialize");

    assert_eq!(restored.label, snap.label);
    assert_eq!(restored.internal_edge_states, snap.internal_edge_states);
    assert_eq!(restored.port_names, snap.port_names);
}

#[test]
fn test_snapshot_fails_on_short_state() {
    let internal = build_4_component_cycle();
    let mut mc = MacroComponent::new(internal);
    mc.set_global_state_offset(0);

    // Only 4 values, but internal needs 12
    let short_state = vec![0.0; 4];
    let snap = mc.to_snapshot(&short_state, None);
    assert!(snap.is_none(), "should return None for short state vector");
}

// ─────────────────────────────────────────────────────────────────────────────
// Two MacroComponent chillers in parallel
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_two_macro_chillers_in_parallel_topology() {
    // Build two identical 4-component chiller MacroComponents.
    let chiller_a = {
        let internal = build_4_component_cycle();
        let mut mc = MacroComponent::new(internal);
        mc.expose_port(0, "in_a", make_port("R134a", 1e5, 4e5));
        mc.expose_port(2, "out_a", make_port("R134a", 5e5, 4.5e5));
        mc
    };
    let chiller_b = {
        let internal = build_4_component_cycle();
        let mut mc = MacroComponent::new(internal);
        mc.expose_port(0, "in_b", make_port("R134a", 1e5, 4e5));
        mc.expose_port(2, "out_b", make_port("R134a", 5e5, 4.5e5));
        mc
    };

    // Place both into a parent system with a splitter and merger mock.
    let mut parent = System::new();
    let ca = parent.add_component(Box::new(chiller_a));
    let cb = parent.add_component(Box::new(chiller_b));
    // Simple pass-through splitter & merger
    let splitter = parent.add_component(pass(1));
    let merger = parent.add_component(pass(1));

    // Topology: splitter → chiller_a → merger
    //                    → chiller_b → merger
    parent.add_edge(splitter, ca).unwrap();
    parent.add_edge(splitter, cb).unwrap();
    parent.add_edge(ca, merger).unwrap();
    parent.add_edge(cb, merger).unwrap();

    let result = parent.finalize();
    assert!(
        result.is_ok(),
        "parallel chiller topology should finalize cleanly: {:?}",
        result.err()
    );

    // CM1.4: 4 parent edges form 2 series branches (S→A→M and S→B→M).
    // Parent state: 2 branches + 4×2 P,h = 10 parent edge vars.
    // 2 chillers × 9 internal vars (1 branch + 4×2 P,h each) = 18 internal vars.
    // Total state vector length = 10 + 18 = 28.
    assert_eq!(parent.state_vector_len(), 28);
    // 4 nodes
    assert_eq!(parent.node_count(), 4);
    // 4 edges
    assert_eq!(parent.edge_count(), 4);

    // Total component equations (CM1.3):
    //   chiller_a: 12 internal (4 components × 3 eqs) + 4 coupling (2 ports × 2) = 16
    //   chiller_b: 12 internal + 4 coupling = 16
    //   splitter:  1
    //   merger:    1
    //   total:     34
    let total_eqs: usize = parent
        .traverse_for_jacobian()
        .map(|(_, c, _)| c.n_equations())
        .sum();
    assert_eq!(
        total_eqs, 34,
        "total equation count mismatch: {}",
        total_eqs
    );
}

#[test]
fn test_two_macro_chillers_residuals_are_computable() {
    let chiller_a = {
        let internal = build_4_component_cycle();
        let mut mc = MacroComponent::new(internal);
        mc.expose_port(0, "in_a", make_port("R134a", 1e5, 4e5));
        mc.expose_port(2, "out_a", make_port("R134a", 5e5, 4.5e5));
        mc
    };
    let chiller_b = {
        let internal = build_4_component_cycle();
        let mut mc = MacroComponent::new(internal);
        mc.expose_port(0, "in_b", make_port("R134a", 1e5, 4e5));
        mc.expose_port(2, "out_b", make_port("R134a", 5e5, 4.5e5));
        mc
    };

    // CM1.4: each chiller has 9 internal state variables (1 branch + 4×2 P,h)
    let _internal_state_len_each = chiller_a.internal_state_len(); // = 9

    let mut parent = System::new();
    let ca = parent.add_component(Box::new(chiller_a));
    let cb = parent.add_component(Box::new(chiller_b));
    let splitter = parent.add_component(pass(1));
    let merger = parent.add_component(pass(1));
    parent.add_edge(splitter, ca).unwrap();
    parent.add_edge(splitter, cb).unwrap();
    parent.add_edge(ca, merger).unwrap();
    parent.add_edge(cb, merger).unwrap();
    parent.finalize().unwrap();

    // CM1.4: parent has 4 edges forming 2 series branches → 2 + 4×2 = 10 parent vars.
    // Each MacroComponent's internal state block starts at offsets assigned cumulatively
    // by System::finalize().
    // chiller_a offset = 10 (after parent edge state)
    // chiller_b offset = 19 (after parent + chiller_a)
    // Total state len = 10 parent + 9 chiller_a + 9 chiller_b = 28 total.
    let full_state_len = parent.state_vector_len();
    assert_eq!(full_state_len, 28);
    let state = vec![0.0; full_state_len];

    // Residual vector must cover every component equation plus the parent's own
    // per-edge mass-flow closures (CM1.2).
    let total_eqs: usize = parent
        .traverse_for_jacobian()
        .map(|(_, c, _)| c.n_equations())
        .sum::<usize>()
        + parent.mass_flow_closure_count();
    let mut residuals = vec![0.0; total_eqs];
    let result = parent.compute_residuals(&state, &mut residuals);
    assert!(
        result.is_ok(),
        "residual computation should not error on zero state: {:?}",
        result.err()
    );
}