//! 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, SystemState,
};
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: &SystemState,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        for r in residuals.iter_mut().take(self.n_eq) {
            *r = 0.0;
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &SystemState,
        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).
fn build_4_component_cycle() -> System {
    let mut sys = System::new();
    let a = sys.add_component(pass(2)); // compressor
    let b = sys.add_component(pass(2)); // condenser
    let c = sys.add_component(pass(2)); // valve
    let d = sys.add_component(pass(2)); // 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 × 2 eqs = 8 internal equations, 0 exposed ports
    assert_eq!(mc.n_equations(), 8,
        "should have 8 internal equations with no exposed ports");
    // 4 edges × 2 vars = 8 internal state vars
    assert_eq!(mc.internal_state_len(), 8);
    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));

    // 8 internal + 4 coupling (2 per port) = 12 equations
    assert_eq!(mc.n_equations(), 12,
        "should have 12 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);

    // Parent state vector: 1 edge × 2 = 2 state vars
    assert_eq!(parent.state_vector_len(), 2);
}

// ─────────────────────────────────────────────────────────────────────────────
// 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));

    // Internal block starts at offset 2 (2 parent-edge state vars before it).
    // External edge for port 0 is at (p=0, h=1).
    mc.set_global_state_offset(2);
    mc.set_system_context(2, &[(0, 1)]);

    // State layout: [P_ext=1e5, h_ext=4e5, P_int_e0=1e5, h_int_e0=4e5, ...]
    //                indices:    0        1         2         3
    let mut state = vec![0.0; 2 + 8]; // 2 parent + 8 internal
    state[0] = 1.0e5; // P_ext
    state[1] = 4.0e5; // h_ext
    state[2] = 1.0e5; // P_int_e0 (consistent with port)
    state[3] = 4.0e5; // h_int_e0

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

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

#[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(2);
    mc.set_system_context(2, &[(0, 1)]);

    let mut state = vec![0.0; 10];
    state[0] = 2.0e5; // P_ext (different from internal)
    state[1] = 5.0e5; // h_ext
    state[2] = 1.0e5; // P_int_e0
    state[3] = 4.0e5; // h_int_e0

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

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

#[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: (p_ext=0, h_ext=1), internal starts at offset=2
    mc.set_global_state_offset(2);
    mc.set_system_context(2, &[(0, 1)]);

    let state = vec![0.0; 10];
    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 8 (P) and 9 (h)
    assert_eq!(find(8, 0), Some(1.0),  "∂r_P/∂p_ext should be +1");
    assert_eq!(find(8, 2), Some(-1.0), "∂r_P/∂int_p should be -1");
    assert_eq!(find(9, 1), Some(1.0),  "∂r_h/∂h_ext should be +1");
    assert_eq!(find(9, 3), 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);

    // Simulate a converged global state (8 internal vars, all nonzero)
    let global_state: Vec<f64> = (0..8).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(), 8);
    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 8
    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());

    // 4 parent edges × 2 = 8 state variables in the parent
    // 2 chillers × 8 internal variables = 16 internal variables
    // Total state vector length = 24
    assert_eq!(parent.state_vector_len(), 24);
    // 4 nodes
    assert_eq!(parent.node_count(), 4);
    // 4 edges
    assert_eq!(parent.edge_count(), 4);

    // Total equations:
    //   chiller_a: 8 internal + 4 coupling (2 ports) = 12
    //   chiller_b: 8 internal + 4 coupling (2 ports) = 12
    //   splitter:  1
    //   merger:    1
    //   total:     26
    let total_eqs: usize = parent
        .traverse_for_jacobian()
        .map(|(_, c, _)| c.n_equations())
        .sum();
    assert_eq!(total_eqs, 26, "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
    };

    // Each chiller has 8 internal state variables (4 edges × 2)
    let internal_state_len_each = chiller_a.internal_state_len(); // = 8

    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();

    // The parent's own state vector covers its 4 edges (8 vars).
    // Each MacroComponent's internal state block starts at offsets assigned cumulatively
    // by System::finalize().
    // chiller_a offset = 8
    // chiller_b offset = 16
    // Total state len = 8 parent + 8 chiller_a + 8 chiller_b = 24 total.
    let full_state_len = parent.state_vector_len();
    assert_eq!(full_state_len, 24);
    let state = vec![0.0; full_state_len];

    let total_eqs: usize = parent
        .traverse_for_jacobian()
        .map(|(_, c, _)| c.n_equations())
        .sum();
    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()
    );
}