//! Comprehensive health check tests for ML Training Service
//!
//! Tests cover:
//! - HTTP /health and /ready endpoints
//! - Service startup state transitions
//! - GPU availability checks
//! - Model checkpoint accessibility
//! - Database connectivity
//! - Training job status
//! - Dependency cascade failures
//! - Health check latency
//! - Concurrent health checks
//! - Health during model training
//! - Resource exhaustion (GPU memory, disk space)

use axum::{
    body::Body,
    http::{Request, StatusCode},
};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tower::ServiceExt;

/// Mock health state for ML training service
#[derive(Clone)]
struct MockMlHealthState {
    healthy: Arc<RwLock<bool>>,
    ready: Arc<RwLock<bool>>,
    gpu_available: Arc<RwLock<bool>>,
    checkpoints_accessible: Arc<RwLock<bool>>,
    database_connected: Arc<RwLock<bool>>,
    training_active: Arc<RwLock<bool>>,
    gpu_memory_available: Arc<RwLock<bool>>,
}

impl MockMlHealthState {
    fn new() -> Self {
        Self {
            healthy: Arc::new(RwLock::new(true)),
            ready: Arc::new(RwLock::new(true)),
            gpu_available: Arc::new(RwLock::new(true)),
            checkpoints_accessible: Arc::new(RwLock::new(true)),
            database_connected: Arc::new(RwLock::new(true)),
            training_active: Arc::new(RwLock::new(false)),
            gpu_memory_available: Arc::new(RwLock::new(true)),
        }
    }

    async fn set_healthy(&self, healthy: bool) {
        *self.healthy.write().await = healthy;
    }

    async fn set_ready(&self, ready: bool) {
        *self.ready.write().await = ready;
    }

    async fn set_gpu_available(&self, available: bool) {
        *self.gpu_available.write().await = available;
    }

    async fn set_checkpoints_accessible(&self, accessible: bool) {
        *self.checkpoints_accessible.write().await = accessible;
    }

    async fn set_database_connected(&self, connected: bool) {
        *self.database_connected.write().await = connected;
    }

    async fn set_training_active(&self, active: bool) {
        *self.training_active.write().await = active;
    }

    async fn set_gpu_memory_available(&self, available: bool) {
        *self.gpu_memory_available.write().await = available;
    }

    async fn is_healthy(&self) -> bool {
        *self.healthy.read().await
    }

    async fn is_ready(&self) -> bool {
        *self.ready.read().await
    }

    async fn is_gpu_available(&self) -> bool {
        *self.gpu_available.read().await
    }

    async fn is_checkpoints_accessible(&self) -> bool {
        *self.checkpoints_accessible.read().await
    }

    async fn is_database_connected(&self) -> bool {
        *self.database_connected.read().await
    }

    async fn is_training_active(&self) -> bool {
        *self.training_active.read().await
    }

    async fn is_gpu_memory_available(&self) -> bool {
        *self.gpu_memory_available.read().await
    }
}

/// Create health router for ML training service
fn create_ml_health_router(state: MockMlHealthState) -> axum::Router {
    use axum::{extract::State, routing::get, Json, Router};
    use serde_json::json;

    async fn health_handler(
        State(state): State<MockMlHealthState>,
    ) -> Result<Json<serde_json::Value>, StatusCode> {
        if state.is_healthy().await {
            Ok(Json(json!({
                "status": "healthy",
                "service": "ml_training",
                "version": env!("CARGO_PKG_VERSION")
            })))
        } else {
            Err(StatusCode::SERVICE_UNAVAILABLE)
        }
    }

    async fn ready_handler(
        State(state): State<MockMlHealthState>,
    ) -> Result<Json<serde_json::Value>, StatusCode> {
        if state.is_ready().await {
            Ok(Json(json!({
                "status": "ready",
                "service": "ml_training",
                "version": env!("CARGO_PKG_VERSION")
            })))
        } else {
            Err(StatusCode::SERVICE_UNAVAILABLE)
        }
    }

    async fn deep_health_handler(
        State(state): State<MockMlHealthState>,
    ) -> Result<Json<serde_json::Value>, StatusCode> {
        let gpu_ok = state.is_gpu_available().await;
        let checkpoints_ok = state.is_checkpoints_accessible().await;
        let db_ok = state.is_database_connected().await;
        let gpu_mem_ok = state.is_gpu_memory_available().await;
        let healthy = state.is_healthy().await;
        let training_active = state.is_training_active().await;

        if healthy && gpu_ok && checkpoints_ok && db_ok && gpu_mem_ok {
            Ok(Json(json!({
                "status": "healthy",
                "service": "ml_training",
                "version": env!("CARGO_PKG_VERSION"),
                "dependencies": {
                    "gpu": "available",
                    "checkpoints": "accessible",
                    "database": "connected",
                    "gpu_memory": "available"
                },
                "training_active": training_active
            })))
        } else {
            Err(StatusCode::SERVICE_UNAVAILABLE)
        }
    }

    Router::new()
        .route("/health", get(health_handler))
        .route("/ready", get(ready_handler))
        .route("/health/deep", get(deep_health_handler))
        .with_state(state)
}

#[tokio::test]
async fn test_ml_health_basic_healthy() {
    let state = MockMlHealthState::new();
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_ml_health_unhealthy() {
    let state = MockMlHealthState::new();
    state.set_healthy(false).await;
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[tokio::test]
async fn test_ml_readiness_check() {
    let state = MockMlHealthState::new();
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/ready")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_ml_service_startup() {
    let state = MockMlHealthState::new();
    state.set_ready(false).await;

    // Initially not ready
    assert!(!state.is_ready().await);
    assert!(state.is_healthy().await);

    // Simulate GPU initialization
    tokio::time::sleep(Duration::from_millis(100)).await;
    state.set_ready(true).await;

    assert!(state.is_ready().await);
}

#[tokio::test]
async fn test_ml_gpu_unavailable() {
    let state = MockMlHealthState::new();
    state.set_gpu_available(false).await;
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[tokio::test]
async fn test_ml_checkpoints_inaccessible() {
    let state = MockMlHealthState::new();
    state.set_checkpoints_accessible(false).await;
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[tokio::test]
async fn test_ml_database_disconnection() {
    let state = MockMlHealthState::new();
    state.set_database_connected(false).await;
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[tokio::test]
async fn test_ml_gpu_memory_exhausted() {
    let state = MockMlHealthState::new();
    state.set_gpu_memory_available(false).await;
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[tokio::test]
async fn test_ml_health_during_training() {
    let state = MockMlHealthState::new();
    state.set_training_active(true).await;

    let app = create_ml_health_router(state.clone());

    // Service should still be healthy during training
    let response = app
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_ml_dependency_cascade_failure() {
    let state = MockMlHealthState::new();

    // All dependencies fail
    state.set_gpu_available(false).await;
    state.set_checkpoints_accessible(false).await;
    state.set_database_connected(false).await;

    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[tokio::test]
async fn test_ml_health_check_latency() {
    let state = MockMlHealthState::new();
    let app = create_ml_health_router(state.clone());

    let start = Instant::now();
    let response = app
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    let latency = start.elapsed();

    assert_eq!(response.status(), StatusCode::OK);
    assert!(
        latency < Duration::from_millis(100),
        "Health check latency: {:?}",
        latency
    );
}

#[tokio::test]
async fn test_ml_concurrent_health_checks() {
    let state = MockMlHealthState::new();

    let mut handles = vec![];
    for _ in 0..100 {
        let state_clone = state.clone();
        let handle = tokio::spawn(async move {
            let app = create_ml_health_router(state_clone);
            let response = app
                .oneshot(
                    Request::builder()
                        .uri("/health")
                        .body(Body::empty())
                        .unwrap(),
                )
                .await
                .unwrap();
            assert_eq!(response.status(), StatusCode::OK);
        });
        handles.push(handle);
    }

    for handle in handles {
        handle.await.unwrap();
    }
}

#[tokio::test]
async fn test_ml_health_during_shutdown() {
    let state = MockMlHealthState::new();

    // Graceful shutdown
    state.set_ready(false).await;
    state.set_healthy(true).await;

    let app = create_ml_health_router(state.clone());

    // Ready check fails
    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/ready")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    // Health check passes
    let response = app
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_ml_rapid_health_checks() {
    let state = MockMlHealthState::new();

    let start = Instant::now();
    for _ in 0..500 {
        let app = create_ml_health_router(state.clone());
        let response = app
            .oneshot(
                Request::builder()
                    .uri("/health")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::OK);
    }
    let duration = start.elapsed();

    assert!(
        duration < Duration::from_secs(1),
        "500 health checks took: {:?}",
        duration
    );
}

#[tokio::test]
async fn test_ml_deep_vs_shallow_health() {
    let state = MockMlHealthState::new();
    let app = create_ml_health_router(state.clone());

    // Shallow health
    let start = Instant::now();
    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    let shallow_latency = start.elapsed();
    assert_eq!(response.status(), StatusCode::OK);

    // Deep health
    let start = Instant::now();
    let response = app
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    let deep_latency = start.elapsed();
    assert_eq!(response.status(), StatusCode::OK);

    assert!(shallow_latency < Duration::from_millis(50));
    assert!(deep_latency < Duration::from_millis(100));
}

#[tokio::test]
async fn test_ml_partial_availability() {
    let state = MockMlHealthState::new();

    // GPU down, checkpoints up
    state.set_gpu_available(false).await;
    state.set_checkpoints_accessible(true).await;
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[tokio::test]
async fn test_ml_recovery_after_failure() {
    let state = MockMlHealthState::new();

    // Service fails
    state.set_healthy(false).await;
    let app = create_ml_health_router(state.clone());

    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    // Service recovers
    state.set_healthy(true).await;
    let response = app
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_ml_health_json_format() {
    let state = MockMlHealthState::new();
    let app = create_ml_health_router(state.clone());

    let response = app
        .oneshot(
            Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    let content_type = response.headers().get("content-type");
    assert!(content_type.is_some());
    let content_type_str = content_type.unwrap().to_str().expect("INVARIANT: Path should be valid UTF-8");
    assert!(content_type_str.contains("application/json"));
}

#[tokio::test]
async fn test_ml_gpu_recovery_scenario() {
    let state = MockMlHealthState::new();

    // GPU fails
    state.set_gpu_available(false).await;
    let app = create_ml_health_router(state.clone());

    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    // GPU recovers
    state.set_gpu_available(true).await;
    let response = app
        .oneshot(
            Request::builder()
                .uri("/health/deep")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}