use super::*;
use crate::dqn::TradingState;
use crate::hyperopt::ParameterSpace;
use std::sync::OnceLock;

/// Shared CUDA device across all trainer tests.
///
/// Root cause fix for CUBLAS_STATUS_NOT_INITIALIZED cascades: each
/// `Device::cuda_if_available(0)` creates a new cuBLAS handle. With
/// 400+ tests doing this in rapid succession (even with --test-threads=1),
/// the driver's internal handle pool is exhausted. Sharing one device
/// eliminates the churn entirely.
static SHARED_DEVICE: OnceLock<MlDevice> = OnceLock::new();

fn shared_cuda_device() -> MlDevice {
    // Initialize tracing so kernel compilation/launch logs are visible.
    static TRACING_INIT: std::sync::Once = std::sync::Once::new();
    TRACING_INIT.call_once(|| {
        let filter = tracing_subscriber::EnvFilter::try_from_default_env()
            .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
        tracing_subscriber::fmt()
            .with_env_filter(filter)
            .with_test_writer()
            .try_init()
            .ok();
    });
    SHARED_DEVICE
        .get_or_init(|| {
            MlDevice::new_cuda(0).expect("CUDA required for tests")
        })
        .clone()
}

// Helper function to create test hyperparameters
// Uses conservative defaults suitable for testing
fn create_test_params() -> DQNHyperparameters {
    let mut params = DQNHyperparameters::conservative();
    // Production default: branching DQN (3 heads: exposure, order, urgency).
    // Always enabled — the warp-cooperative kernel on H100 requires it.
    params.hidden_dim_base = Some(32); // Small for fast test iterations
    params.buffer_size = 10_000;
    params.replay_buffer_vram_fraction = 0.0; // Disable AutoReplaySizer in tests
    params
}

fn create_test_trainer() -> Result<DQNTrainer> {
    DQNTrainer::new_with_device(create_test_params(), shared_cuda_device())
}

fn create_test_trainer_with(params: DQNHyperparameters) -> Result<DQNTrainer> {
    DQNTrainer::new_with_device(params, shared_cuda_device())
}

/// Pad a TradingState's regime_features so that `state.dimension()` matches the
/// trainer's aligned state_dim (e.g. 45→48 on CUDA due to tensor core alignment).
fn pad_state_to_aligned(state: &mut TradingState) {
    let aligned_dim = (state.dimension() + 7) & !7;
    let pad = aligned_dim.saturating_sub(state.dimension());
    if pad > 0 {
        state.regime_features.extend(vec![0.0_f32; pad]);
    }
}

#[tokio::test]
async fn test_dqn_trainer_creation() {
    let hyperparams = create_test_params();
    let trainer = create_test_trainer_with(hyperparams);

    assert!(
        trainer.is_ok(),
        "Failed to create DQN trainer: {:?}",
        trainer.err()
    );
}

#[tokio::test]
async fn test_batch_size_validation() {
    let mut hyperparams = create_test_params();
    hyperparams.batch_size = 500;

    // VRAM ceiling clamps if needed, never rejects
    let trainer = create_test_trainer_with(hyperparams);
    assert!(
        trainer.is_ok(),
        "Should clamp oversized batch, not reject: {:?}",
        trainer.err()
    );
}

#[tokio::test]
async fn test_feature_vector_to_state() {
    let hyperparams = create_test_params();
    let trainer = create_test_trainer_with(hyperparams).unwrap();

    // Create a synthetic 42-dim feature vector (42 market features)
    let mut feature_vec = [0.0; 42];
    feature_vec[0] = 4000.0; // open
    feature_vec[1] = 4010.0; // high
    feature_vec[2] = 3990.0; // low
    feature_vec[3] = 4005.0; // close
    feature_vec[4] = 1000.0; // volume
                             // Fill remaining features with synthetic data
    for i in 5..42 {
        feature_vec[i] = (i as f64) * 0.1;
    }

    let close_price =
        rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO);
    let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price));

    assert!(
        state.is_ok(),
        "Failed to convert feature vector: {:?}",
        state.err()
    );

    let state = state.unwrap();
    // State dimension is 45 (42 market + 3 portfolio)
    // - Market features: 0-41 (42 features)
    // - Portfolio features: 42-44 (3 features, populated by PortfolioTracker)
    assert_eq!(
        state.dimension(),
        45,
        "State dimension should be 45 (42 market + 3 portfolio features)"
    );
}

#[tokio::test]
async fn test_batched_action_selection() {
    let hyperparams = create_test_params();
    let mut trainer = create_test_trainer_with(hyperparams).unwrap();

    // Create multiple synthetic states for batched action selection
    let batch_size = 10;
    let mut states = Vec::with_capacity(batch_size);

    for i in 0..batch_size {
        let mut feature_vec = [0.0; 42]; // 42 market features
                                          // Create varied states for testing
        feature_vec[0] = 4000.0 + (i as f64 * 10.0); // open
        feature_vec[1] = 4010.0 + (i as f64 * 10.0); // high
        feature_vec[2] = 3990.0 + (i as f64 * 10.0); // low
        feature_vec[3] = 4005.0 + (i as f64 * 10.0); // close
        feature_vec[4] = 1000.0 + (i as f64 * 100.0); // volume

        // Fill remaining features
        for j in 5..42 {
            feature_vec[j] = (j as f64 + i as f64) * 0.1;
        }

        let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
            .unwrap_or(rust_decimal::Decimal::ZERO);
        let mut state = trainer
            .feature_vector_to_state(&feature_vec, Some(close_price))
            .unwrap();
        pad_state_to_aligned(&mut state);
        states.push(state);
    }

    // Test batched action selection
    let actions_result = trainer.select_actions_batch(&states).await;

    assert!(
        actions_result.is_ok(),
        "Batched action selection failed: {:?}",
        actions_result.err()
    );

    let actions = actions_result.unwrap();
    assert_eq!(
        actions.len(),
        batch_size,
        "Expected {} actions, got {}",
        batch_size,
        actions.len()
    );

    // Verify all actions have valid exposure indices (0-4)
    for (i, action) in actions.iter().enumerate() {
        let exp_idx = action.exposure as usize;
        assert!(
            exp_idx < 5,
            "Action {} has invalid exposure index {}: {:?}",
            i,
            exp_idx,
            action
        );
    }
}

#[tokio::test]
async fn test_batched_vs_sequential_action_selection_consistency() {
    let hyperparams = create_test_params();
    let mut trainer = create_test_trainer_with(hyperparams).unwrap();

    // Create test states
    let batch_size = 5;
    let mut states = Vec::with_capacity(batch_size);

    for i in 0..batch_size {
        let mut feature_vec = [0.0; 42]; // 42 market features
        feature_vec[0] = 4000.0 + (i as f64 * 50.0);
        feature_vec[1] = 4050.0 + (i as f64 * 50.0);
        feature_vec[2] = 3950.0 + (i as f64 * 50.0);
        feature_vec[3] = 4025.0 + (i as f64 * 50.0);
        feature_vec[4] = 5000.0 + (i as f64 * 500.0);

        for j in 5..42 {
            feature_vec[j] = (j as f64) * 0.5 + (i as f64);
        }

        let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
            .unwrap_or(rust_decimal::Decimal::ZERO);
        let mut state = trainer
            .feature_vector_to_state(&feature_vec, Some(close_price))
            .unwrap();
        pad_state_to_aligned(&mut state);
        states.push(state);
    }

    // Get batched actions (GPU-optimized)
    let batched_actions = trainer.select_actions_batch(&states).await.unwrap();

    // Both should return valid actions
    assert_eq!(
        batched_actions.len(),
        batch_size,
        "Batched action count mismatch"
    );

    // Verify all actions have valid exposure indices (0-4)
    for action in &batched_actions {
        let exp_idx = action.exposure as usize;
        assert!(exp_idx < 5, "Invalid exposure index {}: {:?}", exp_idx, action);
    }
}

#[tokio::test]
async fn test_empty_batch_handling() {
    let hyperparams = create_test_params();
    let mut trainer = create_test_trainer_with(hyperparams).unwrap();

    let empty_states: Vec<TradingState> = Vec::new();
    let result = trainer.select_actions_batch(&empty_states).await;

    assert!(result.is_ok(), "Empty batch should be handled gracefully");
    assert_eq!(
        result.unwrap().len(),
        0,
        "Empty batch should return empty actions"
    );
}

#[tokio::test]
async fn test_zero_batch_size_handling() {
    // Test DQN rejects zero batch size
    let mut hyperparams = create_test_params();
    hyperparams.batch_size = 0;

    let result = create_test_trainer_with(hyperparams);

    // Should fail with descriptive error
    assert!(
        result.is_err(),
        "DQN should reject zero batch size, but got: {:?}",
        result
    );

    // Error message should mention batch size
    let error_msg = result.unwrap_err().to_string();
    assert!(
        error_msg.to_lowercase().contains("batch"),
        "Error message should mention batch size, got: {}",
        error_msg
    );
}

// ===== Agent 23 Test #6: Batch Size Mismatch Validation Tests =====

/// Production-critical test: Verify trainer handles batch smaller than configured
#[tokio::test]
async fn test_batch_size_mismatch_smaller_than_configured() {
    let mut hyperparams = create_test_params();
    hyperparams.batch_size = 32;
    let mut trainer = create_test_trainer_with(hyperparams).unwrap();

    // Create batch with 16 states (half of configured 32)
    let mut feature_vec = [0.0; 42]; // 42 market features
    for i in 0..4 {
        feature_vec[i] = 4000.0 + (i as f64 * 10.0);
    }
    for i in 5..42 {
        feature_vec[i] = (i as f64) * 0.1;
    }

    let close_price =
        rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO);
    let mut state = trainer
        .feature_vector_to_state(&feature_vec, Some(close_price))
        .unwrap();
    pad_state_to_aligned(&mut state);
    let smaller_batch = vec![state.clone(); 16];

    let result = trainer.select_actions_batch(&smaller_batch).await;
    assert!(
        result.is_ok(),
        "DQN should handle smaller batches: {:?}",
        result.err()
    );
    assert_eq!(
        result.unwrap().len(),
        16,
        "Should return action for each state"
    );
}

/// Production-critical test: Verify trainer handles batch larger than configured
#[tokio::test]
async fn test_batch_size_mismatch_larger_than_configured() {
    let mut hyperparams = create_test_params();
    hyperparams.batch_size = 16;
    let mut trainer = create_test_trainer_with(hyperparams).unwrap();

    // Create batch with 64 states (4x configured 16)
    let mut feature_vec = [0.0; 42]; // 42 market features
    for i in 0..4 {
        feature_vec[i] = 4000.0 + (i as f64 * 10.0);
    }
    for i in 5..42 {
        feature_vec[i] = (i as f64) * 0.1;
    }

    let close_price =
        rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO);
    let mut state = trainer
        .feature_vector_to_state(&feature_vec, Some(close_price))
        .unwrap();
    pad_state_to_aligned(&mut state);
    let larger_batch = vec![state.clone(); 64];

    let result = trainer.select_actions_batch(&larger_batch).await;
    assert!(
        result.is_ok(),
        "DQN should handle larger batches: {:?}",
        result.err()
    );
    assert_eq!(
        result.unwrap().len(),
        64,
        "Should return action for each state"
    );
}

/// Production-critical test: Verify empty batch handling
#[tokio::test]
async fn test_empty_batch_returns_empty_actions() {
    let mut trainer = create_test_trainer().unwrap();
    let empty_batch: Vec<TradingState> = vec![];

    let result = trainer.select_actions_batch(&empty_batch).await;
    assert!(result.is_ok(), "Should handle empty batch gracefully");
    assert_eq!(
        result.unwrap().len(),
        0,
        "Empty batch should return empty actions"
    );
}

/// Production-critical test: Verify single-sample batch handling
#[tokio::test]
async fn test_single_sample_batch() {
    let mut hyperparams = create_test_params();
    hyperparams.batch_size = 32;
    let mut trainer = create_test_trainer_with(hyperparams).unwrap();

    let mut feature_vec = [0.0; 42]; // 42 market features
    for i in 0..4 {
        feature_vec[i] = 4000.0;
    }
    for i in 5..42 {
        feature_vec[i] = (i as f64) * 0.1;
    }

    let close_price =
        rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO);
    let mut state = trainer
        .feature_vector_to_state(&feature_vec, Some(close_price))
        .unwrap();
    pad_state_to_aligned(&mut state);
    let single_batch = vec![state];

    let result = trainer.select_actions_batch(&single_batch).await;
    assert!(
        result.is_ok(),
        "Should handle single-sample batch: {:?}",
        result.err()
    );
    assert_eq!(result.unwrap().len(), 1, "Should return exactly one action");
}

/// Large batch sizes are accepted (VRAM ceiling is the only cap)
#[test]
fn test_large_batch_size_accepted() {
    let mut hyperparams = create_test_params();
    hyperparams.batch_size = 2048;

    let result = create_test_trainer_with(hyperparams);
    assert!(
        result.is_ok(),
        "Should accept large batch sizes within VRAM ceiling: {:?}",
        result.err()
    );
}

/// Production-critical test: Non-power-of-2 batch sizes
#[tokio::test]
async fn test_non_power_of_two_batch_size() {
    let mut hyperparams = create_test_params();
    hyperparams.batch_size = 13; // Not a power of 2

    let result = create_test_trainer_with(hyperparams);
    assert!(
        result.is_ok(),
        "Should accept non-power-of-2 batch sizes: {:?}",
        result.err()
    );
}

/// Production-critical test: Train with empty dataset doesn't crash
#[tokio::test]
async fn test_train_with_empty_data_completes_gracefully() {
    let mut params = create_test_params();
    params.epochs = 5; // Short run — just checking it doesn't panic
    params.early_stopping_enabled = false;
    params.gradient_collapse_patience = 1000;
    params.buffer_size = 1024; // MIN_GPU_CAPACITY — GPU PER mandatory
    let device = MlDevice::new_cuda(0).expect("CUDA device required");
    let mut trainer = DQNTrainer::new_with_device(params, device).unwrap();
    let empty_data: Vec<(FeatureVector, Vec<f64>)> = vec![];
    let checkpoint_callback = |_, _, _| Ok(String::new());

    let result = trainer
        .train_with_data_full_loop(&empty_data, checkpoint_callback)
        .await;

    assert!(
        result.is_err(),
        "Training with empty data should return an error (no CPU fallback)"
    );
}

/// Test reward function calculates actual price changes correctly
#[test]
fn test_reward_function_price_changes() {
    let trainer = create_test_trainer().unwrap();

    // Test upward price move (+14.25 points, should clamp to +1.0)
    let reward_up = trainer.calculate_reward(5900.0, 5914.25);
    assert!(
        (reward_up - 1.0).abs() < 1e-6,
        "Upward move should return +1.0 (clamped), got: {}",
        reward_up
    );

    // Test downward price move (-14.25 points, should clamp to -1.0)
    let reward_down = trainer.calculate_reward(5914.25, 5900.0);
    assert!(
        (reward_down - (-1.0)).abs() < 1e-6,
        "Downward move should return -1.0 (clamped), got: {}",
        reward_down
    );

    // Test flat market (0 points, should return 0.0)
    let reward_flat = trainer.calculate_reward(5900.0, 5900.0);
    assert!(
        reward_flat.abs() < 1e-6,
        "Flat market should return 0.0, got: {}",
        reward_flat
    );

    // Test small upward move (+5 points, should return +0.5)
    let reward_small_up = trainer.calculate_reward(5900.0, 5905.0);
    assert!(
        (reward_small_up - 0.5).abs() < 1e-6,
        "Small upward move (+5) should return +0.5, got: {}",
        reward_small_up
    );

    // Test small downward move (-5 points, should return -0.5)
    let reward_small_down = trainer.calculate_reward(5905.0, 5900.0);
    assert!(
        (reward_small_down - (-0.5)).abs() < 1e-6,
        "Small downward move (-5) should return -0.5, got: {}",
        reward_small_down
    );

    // Test unclamped move (+3 points, should return +0.3)
    let reward_unclamped = trainer.calculate_reward(5900.0, 5903.0);
    assert!(
        (reward_unclamped - 0.3).abs() < 1e-6,
        "Move of +3 points should return +0.3, got: {}",
        reward_unclamped
    );
}

#[test]
fn test_dynamic_batch_size_l4() {
    // L4 has 24GB VRAM — HardwareBudget should allow batch_size >> 230
    let budget = crate::hyperopt::HardwareBudget {
        gpu_memory_mb: 24_000,
        gpu_name: "NVIDIA L4".to_string(),
    };
    let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0);
    assert!(batch.unwrap_or(0.0) > 230.0, "L4 should support DQN batch > 230, got {:?}", batch);
}

#[test]
fn test_dynamic_batch_size_h100() {
    // H100 has 80GB VRAM — should hit the 8192 ceiling
    let budget = crate::hyperopt::HardwareBudget {
        gpu_memory_mb: 81_920,
        gpu_name: "NVIDIA H100".to_string(),
    };
    let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0);
    assert!((batch.unwrap_or(0.0) - 8192.0).abs() < 1.0, "H100 should hit 8192 ceiling, got {:?}", batch);
}

// ── C2 Overhaul Smoke Tests ─────────────────────────────────────────

/// Verify DQN action space is 9 exposure levels (not 81 factored actions).
#[test]
fn test_c2_dqn_default_num_actions_is_9() {
    let config = crate::dqn::DQNConfig::default();
    assert_eq!(config.num_actions, 9, "DQN default must be 9 exposure-level actions (25% steps)");
}

/// Verify 5 exposure indices produce 5 distinct exposure levels.
#[test]
fn test_c2_five_actions_produce_distinct_exposures() {
    use crate::dqn::action_space::ExposureLevel;
    use crate::dqn::order_router::OrderRouter;

    let actions: Vec<_> = (0..5)
        .filter_map(|idx| ExposureLevel::from_index(idx).ok())
        .map(|e| OrderRouter::route_default(e))
        .collect();

    assert_eq!(actions.len(), 5);

    let unique: std::collections::HashSet<_> = actions.iter().map(|a| a.exposure).collect();
    assert_eq!(unique.len(), 5, "All 5 exposure levels must be distinct");
}

/// Verify hyperopt search space is 45D (Task 11: +4 new dims).
#[test]
fn test_c3_search_space_is_40d() {
    let bounds = crate::hyperopt::adapters::dqn::DQNParams::continuous_bounds();
    assert_eq!(bounds.len(), 45, "Search space must be 45D");

    let names = crate::hyperopt::adapters::dqn::DQNParams::param_names();
    assert_eq!(names.len(), 45);
    assert!(names.contains(&"count_bonus_coefficient"), "count_bonus_coefficient must be in search space (C3)");
    assert!(names.contains(&"sharpe_weight"), "sharpe_weight must be in search space (C4)");
    assert!(names.contains(&"branch_hidden_dim"), "branch_hidden_dim must be in search space (L2)");
    assert!(names.contains(&"c51_warmup_epochs"), "c51_warmup_epochs must be in search space");
    // Task 11: curiosity_weight is now in the search space
    assert!(names.contains(&"curiosity_weight"), "curiosity_weight must be in search space (Task 11)");
    assert!(names.contains(&"her_ratio"), "her_ratio must be in search space (Task 11)");
    assert!(names.contains(&"use_cvar_action_selection"), "use_cvar_action_selection must be in search space (Task 11)");
    assert!(names.contains(&"cvar_alpha"), "cvar_alpha must be in search space (Task 11)");
    assert!(!names.contains(&"noisy_epsilon_floor"), "noisy_epsilon_floor must not be in search space");
}

/// Verify noisy_epsilon_floor is fixed to 0.10 (prevents action collapse).
#[test]
fn test_noisy_epsilon_floor_fixed() {
    let params = crate::hyperopt::adapters::dqn::DQNParams::default();
    assert!(
        (params.noisy_epsilon_floor - 0.10).abs() < 1e-6,
        "noisy_epsilon_floor must default to 0.10 (prevents action collapse)"
    );
}

/// Verify exploration params are fixed after C2 cleanup.
#[test]
fn test_c2_exploration_params_fixed() {
    let params = crate::hyperopt::adapters::dqn::DQNParams::default();
    assert!(
        params.curiosity_weight.abs() < f64::EPSILON,
        "curiosity_weight must be fixed at 0.0"
    );
    assert!(
        (params.noisy_epsilon_floor - 0.10).abs() < 1e-6,
        "noisy_epsilon_floor must be fixed at 0.10"
    );
    assert!(
        params.count_bonus_coefficient.abs() < f64::EPSILON,
        "count_bonus_coefficient must be fixed at 0.0"
    );

    // Roundtrip through from_continuous should preserve fixed values
    let continuous = params.to_continuous();
    let recovered = crate::hyperopt::adapters::dqn::DQNParams::from_continuous(&continuous).unwrap();
    assert!(
        recovered.curiosity_weight.abs() < f64::EPSILON,
        "curiosity_weight must remain 0.0 after roundtrip"
    );
    assert!(
        recovered.count_bonus_coefficient.abs() < f64::EPSILON,
        "count_bonus_coefficient must remain 0.0 after roundtrip"
    );
}
