//! Integration tests for Story 4.5: Time-Budgeted Solving
//!
//! Tests the timeout behavior with best-state return:
//! - Timeout returns best state instead of error
//! - Best state is the lowest residual encountered
//! - ZOH (Zero-Order Hold) fallback for HIL scenarios
//! - Configurable timeout behavior
//! - Timeout across fallback switches preserves best state

use entropyk_components::{
    Component, ComponentError, JacobianBuilder, ResidualVector, StateSlice,
};
use entropyk_solver::solver::{
    ConvergenceStatus, FallbackConfig, FallbackSolver, NewtonConfig, PicardConfig, Solver,
    SolverError, TimeoutConfig,
};
use entropyk_solver::system::System;
use std::time::Duration;

// ─────────────────────────────────────────────────────────────────────────────
// Mock Components for Testing
// ─────────────────────────────────────────────────────────────────────────────

/// A 2x2 linear system: r = A * x - b
struct LinearSystem2x2 {
    a: [[f64; 2]; 2],
    b: [f64; 2],
}

impl LinearSystem2x2 {
    fn well_conditioned() -> Self {
        Self {
            a: [[2.0, 1.0], [1.0, 2.0]],
            b: [3.0, 3.0],
        }
    }
}

impl Component for LinearSystem2x2 {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        residuals[0] = self.a[0][0] * state[0] + self.a[0][1] * state[1] - self.b[0];
        residuals[1] = self.a[1][0] * state[0] + self.a[1][1] * state[1] - self.b[1];
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        jacobian.add_entry(0, 0, self.a[0][0]);
        jacobian.add_entry(0, 1, self.a[0][1]);
        jacobian.add_entry(1, 0, self.a[1][0]);
        jacobian.add_entry(1, 1, self.a[1][1]);
        Ok(())
    }

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

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

fn create_test_system(component: Box<dyn Component>) -> System {
    let mut system = System::new();
    let n0 = system.add_component(component);
    system.add_edge(n0, n0).unwrap();
    system.finalize().unwrap();
    system
}

// ─────────────────────────────────────────────────────────────────────────────
// TimeoutConfig Tests (AC: #6)
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_timeout_config_defaults() {
    let config = TimeoutConfig::default();
    assert!(config.return_best_state_on_timeout);
    assert!(!config.zoh_fallback);
}

#[test]
fn test_timeout_config_zoh_enabled() {
    let config = TimeoutConfig {
        return_best_state_on_timeout: true,
        zoh_fallback: true,
    };
    assert!(config.zoh_fallback);
}

#[test]
fn test_timeout_config_return_error_on_timeout() {
    let config = TimeoutConfig {
        return_best_state_on_timeout: false,
        zoh_fallback: false,
    };
    assert!(!config.return_best_state_on_timeout);
}

// ─────────────────────────────────────────────────────────────────────────────
// AC: #1, #2 - Timeout Returns Best State
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_timeout_returns_best_state_not_error() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let timeout = Duration::from_nanos(1);
    let mut solver = NewtonConfig {
        timeout: Some(timeout),
        max_iterations: 10000,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: true,
            zoh_fallback: false,
        },
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    match result {
        Ok(state) => {
            assert!(
                state.status == ConvergenceStatus::Converged
                    || state.status == ConvergenceStatus::TimedOutWithBestState
            );
        }
        Err(SolverError::Timeout { .. }) => {}
        Err(other) => panic!("Unexpected error: {:?}", other),
    }
}

#[test]
fn test_best_state_is_lowest_residual() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let timeout = Duration::from_micros(100);
    let mut solver = NewtonConfig {
        timeout: Some(timeout),
        max_iterations: 10000,
        timeout_config: TimeoutConfig::default(),
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    if let Ok(state) = result {
        assert!(state.final_residual.is_finite());
        assert!(state.final_residual >= 0.0);
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// AC: #3 - ZOH Fallback
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_zoh_fallback_returns_previous_state() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let previous_state = vec![1.0, 2.0];
    let timeout = Duration::from_nanos(1);

    let mut solver = NewtonConfig {
        timeout: Some(timeout),
        max_iterations: 10000,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: true,
            zoh_fallback: true,
        },
        previous_state: Some(previous_state.clone()),
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    if let Ok(state) = result {
        if state.status == ConvergenceStatus::TimedOutWithBestState {
            assert_eq!(state.state, previous_state);
        }
    }
}

#[test]
fn test_zoh_fallback_ignored_without_previous_state() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let timeout = Duration::from_nanos(1);

    let mut solver = NewtonConfig {
        timeout: Some(timeout),
        max_iterations: 10000,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: true,
            zoh_fallback: true,
        },
        previous_state: None,
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    if let Ok(state) = result {
        if state.status == ConvergenceStatus::TimedOutWithBestState {
            assert_eq!(state.state.len(), 2);
        }
    }
}

#[test]
fn test_zoh_fallback_picard() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let previous_state = vec![5.0, 10.0];
    let timeout = Duration::from_nanos(1);

    let mut solver = PicardConfig {
        timeout: Some(timeout),
        max_iterations: 10000,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: true,
            zoh_fallback: true,
        },
        previous_state: Some(previous_state.clone()),
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    if let Ok(state) = result {
        if state.status == ConvergenceStatus::TimedOutWithBestState {
            assert_eq!(state.state, previous_state);
        }
    }
}

#[test]
fn test_zoh_fallback_uses_previous_residual() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let previous_state = vec![1.0, 2.0];
    let previous_residual = 1e-4;
    let timeout = Duration::from_nanos(1);

    let mut solver = NewtonConfig {
        timeout: Some(timeout),
        max_iterations: 10000,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: true,
            zoh_fallback: true,
        },
        previous_state: Some(previous_state.clone()),
        previous_residual: Some(previous_residual),
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    if let Ok(state) = result {
        if state.status == ConvergenceStatus::TimedOutWithBestState {
            assert_eq!(state.state, previous_state);
            assert!((state.final_residual - previous_residual).abs() < 1e-10);
        }
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// AC: #6 - return_best_state_on_timeout = false
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_timeout_returns_error_when_configured() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let timeout = Duration::from_millis(1);

    let mut solver = NewtonConfig {
        timeout: Some(timeout),
        max_iterations: 10000,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: false,
            zoh_fallback: false,
        },
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    match result {
        Err(SolverError::Timeout { .. }) | Ok(_) => {}
        Err(other) => panic!("Expected Timeout or Ok, got {:?}", other),
    }
}

#[test]
fn test_picard_timeout_returns_error_when_configured() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let timeout = Duration::from_millis(1);

    let mut solver = PicardConfig {
        timeout: Some(timeout),
        max_iterations: 10000,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: false,
            zoh_fallback: false,
        },
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    match result {
        Err(SolverError::Timeout { .. }) | Ok(_) => {}
        Err(other) => panic!("Expected Timeout or Ok, got {:?}", other),
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// AC: #4 - Timeout Across Fallback Switches
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_timeout_across_fallback_switches_preserves_best_state() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let timeout = Duration::from_millis(10);

    let mut solver = FallbackSolver::new(FallbackConfig {
        fallback_enabled: true,
        max_fallback_switches: 2,
        ..Default::default()
    })
    .with_timeout(timeout)
    .with_newton_config(NewtonConfig {
        max_iterations: 500,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: true,
            zoh_fallback: false,
        },
        ..Default::default()
    })
    .with_picard_config(PicardConfig {
        max_iterations: 500,
        timeout_config: TimeoutConfig {
            return_best_state_on_timeout: true,
            zoh_fallback: false,
        },
        ..Default::default()
    });

    let result = solver.solve(&mut system);
    match result {
        Ok(state) => {
            assert!(
                state.status == ConvergenceStatus::Converged
                    || state.status == ConvergenceStatus::TimedOutWithBestState
            );
            assert!(state.final_residual.is_finite());
        }
        Err(SolverError::Timeout { .. }) => {}
        Err(other) => panic!("Unexpected error: {:?}", other),
    }
}

#[test]
fn test_fallback_solver_total_timeout() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let timeout = Duration::from_millis(5);

    let mut solver = FallbackSolver::default_solver()
        .with_timeout(timeout)
        .with_newton_config(NewtonConfig {
            max_iterations: 10000,
            ..Default::default()
        })
        .with_picard_config(PicardConfig {
            max_iterations: 10000,
            ..Default::default()
        });

    let start = std::time::Instant::now();
    let result = solver.solve(&mut system);
    let elapsed = start.elapsed();

    if result.is_err()
        || matches!(result, Ok(ref s) if s.status == ConvergenceStatus::TimedOutWithBestState)
    {
        assert!(
            elapsed < timeout + Duration::from_millis(100),
            "Total solve time should respect timeout budget. Elapsed: {:?}, Timeout: {:?}",
            elapsed,
            timeout
        );
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Pre-allocation Tests (AC: #5)
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_newton_config_best_state_preallocated() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let mut solver = NewtonConfig {
        timeout: Some(Duration::from_millis(100)),
        max_iterations: 10,
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    assert!(result.is_ok() || matches!(result, Err(SolverError::Timeout { .. })));
}

#[test]
fn test_picard_config_best_state_preallocated() {
    let mut system = create_test_system(Box::new(LinearSystem2x2::well_conditioned()));
    let mut solver = PicardConfig {
        timeout: Some(Duration::from_millis(100)),
        max_iterations: 10,
        ..Default::default()
    };

    let result = solver.solve(&mut system);
    match result {
        Ok(_) | Err(SolverError::Timeout { .. }) | Err(SolverError::NonConvergence { .. }) => {}
        Err(other) => panic!("Unexpected error: {:?}", other),
    }
}