use super::helpers::*;
use ml_core::cuda_autograd::{GpuTensor, GpuVarStore, GpuAdamW, AdamWConfig};

use std::sync::Arc;
use tracing::info;

// GPU replay buffer (ml-dqn crate)
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};

// ---------------------------------------------------------------------------
// GpuReplayBuffer tests (compiled only with cuda feature, GPU-only)
// ---------------------------------------------------------------------------

fn test_buffer_config(capacity: usize, state_dim: usize) -> GpuReplayBufferConfig {
    GpuReplayBufferConfig {
        capacity,
        state_dim,
        alpha: 0.6,
        beta_start: 0.4,
        beta_max: 1.0,
        beta_annealing_steps: 1000,
        epsilon: 1e-6,
        max_memory_bytes: 4 * 1024 * 1024 * 1024,
        max_batch_size: 1024,
    }
}

/// Get the CUDA stream from the shared smoke device.
fn smoke_stream() -> Arc<cudarc::driver::CudaStream> {
    let dev = cuda_device();
    Arc::clone(dev.cuda_stream().expect("cuda stream"))
}

/// Create random f32 CudaSlice on GPU (host-generated, uploaded).
fn random_cuda_f32(n: usize, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<cudarc::driver::CudaSlice<f32>> {
    use rand::Rng;
    let mut rng = rand::thread_rng();
    let host: Vec<f32> = (0..n).map(|_| rng.gen_range(-1.0_f32..1.0)).collect();
    let slice = stream.clone_htod(&host).map_err(|e| anyhow::anyhow!("htod f32: {e}"))?;
    Ok(slice)
}

/// Create zeroed u32 CudaSlice on GPU.
fn zeros_cuda_u32(n: usize, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<cudarc::driver::CudaSlice<u32>> {
    let slice = stream.alloc_zeros::<u32>(n).map_err(|e| anyhow::anyhow!("alloc u32: {e}"))?;
    Ok(slice)
}

/// Create zeroed f32 CudaSlice on GPU.
fn zeros_cuda_f32(n: usize, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<cudarc::driver::CudaSlice<f32>> {
    let slice = stream.alloc_zeros::<f32>(n).map_err(|e| anyhow::anyhow!("alloc f32: {e}"))?;
    Ok(slice)
}

/// Insert random experiences into a GPU replay buffer.
fn insert_random_batch(
    buf: &mut GpuReplayBuffer,
    n: usize,
    state_dim: usize,
) -> anyhow::Result<()> {
    let stream = buf.stream();
    let states = random_cuda_f32(n * state_dim, stream)?;
    let next_states = random_cuda_f32(n * state_dim, stream)?;
    let actions = zeros_cuda_u32(n, stream)?;
    let rewards = random_cuda_f32(n, stream)?;
    let dones = zeros_cuda_f32(n, stream)?;
    buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, n)?;
    Ok(())
}

/// Download a CudaSlice<f32> to host for assertion.
fn dtoh_f32(slice: &cudarc::driver::CudaSlice<f32>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<f32>> {
    let n = slice.len();
    let mut host = vec![0.0_f32; n];
    stream.memcpy_dtoh(slice, &mut host).map_err(|e| anyhow::anyhow!("dtoh f32: {e}"))?; // test-only readback
    Ok(host)
}

/// Download a CudaSlice<u32> to host for assertion.
fn dtoh_u32(slice: &cudarc::driver::CudaSlice<u32>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<u32>> {
    let n = slice.len();
    let mut host = vec![0_u32; n];
    stream.memcpy_dtoh(slice, &mut host).map_err(|e| anyhow::anyhow!("dtoh u32: {e}"))?; // test-only readback
    Ok(host)
}

#[test]
fn test_gpu_replay_buffer_insert_and_len() -> anyhow::Result<()> {
    let stream = smoke_stream();
    let config = test_buffer_config(100, 48);
    let mut buf = GpuReplayBuffer::new(config, &stream)?;

    insert_random_batch(&mut buf, 20, 48)?;

    assert_eq!(buf.len(), 20);
    Ok(())
}

#[test]
fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<()> {
    let stream = smoke_stream();
    let config = test_buffer_config(100, 48);
    let mut buf = GpuReplayBuffer::new(config, &stream)?;

    insert_random_batch(&mut buf, 50, 48)?;

    let batch = buf.sample_proportional(16)?;

    assert_eq!(batch.batch_size, 16);

    // Download weights and indices to host for assertions
    let weights_host = dtoh_f32(&batch.weights, &stream)?;
    let indices_host = dtoh_u32(&batch.indices, &stream)?;

    // All weights must be positive
    for (i, &w) in weights_host.iter().enumerate() {
        assert!(w > 0.0, "weight[{i}] is not positive: {w}");
    }

    // Weights must be finite
    let w_sum: f32 = weights_host.iter().sum();
    assert!(w_sum.is_finite(), "weight sum is not finite: {w_sum}");

    // All indices must be in range [0, 50)
    for (i, &idx) in indices_host.iter().enumerate() {
        assert!(idx < 50, "index[{i}] out of range: {idx}");
    }

    // Shape sanity
    assert_eq!(indices_host.len(), 16);

    Ok(())
}

#[test]
fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> {
    let stream = smoke_stream();
    let config = test_buffer_config(100, 48);
    let mut buf = GpuReplayBuffer::new(config, &stream)?;

    insert_random_batch(&mut buf, 50, 48)?;

    let batch = buf.sample_proportional(16)?;

    assert_eq!(batch.batch_size, 16);

    // Download weights and indices to host for assertions
    let weights_host = dtoh_f32(&batch.weights, &stream)?;
    let indices_host = dtoh_u32(&batch.indices, &stream)?;

    // All weights must be positive
    for (i, &w) in weights_host.iter().enumerate() {
        assert!(w > 0.0, "weight[{i}] is not positive: {w}");
    }

    // Weights must be finite
    let w_sum: f32 = weights_host.iter().sum();
    assert!(w_sum.is_finite(), "weight sum is not finite: {w_sum}");

    // All indices must be in range [0, 50)
    for (i, &idx) in indices_host.iter().enumerate() {
        assert!(idx < 50, "index[{i}] out of range: {idx}");
    }

    // Shape sanity
    assert_eq!(indices_host.len(), 16);

    Ok(())
}

#[test]
fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> {
    let stream = smoke_stream();
    let config = test_buffer_config(100, 48);
    let epsilon = config.epsilon;
    let mut buf = GpuReplayBuffer::new(config, &stream)?;

    insert_random_batch(&mut buf, 50, 48)?;

    // Sample and get indices
    let batch = buf.sample_proportional(16)?;
    let td_errors = random_cuda_f32(16, &stream)?;
    buf.update_priorities_gpu(&batch.indices, &td_errors, 16)?;

    // Verify sampling still works with updated priorities (proves update succeeded)
    let batch2 = buf.sample_proportional(16)?;
    let weights_host = dtoh_f32(&batch2.weights, &stream)?;

    // Weights should still be positive after priority update
    for (i, &w) in weights_host.iter().enumerate() {
        assert!(w > 0.0, "weight[{i}] not positive after update: {w}");
    }
    let w_sum: f32 = weights_host.iter().sum();
    assert!(w_sum.is_finite(), "weight sum not finite after update: {w_sum}");

    // Epsilon is the priority floor -- verify it's reasonable
    assert!(epsilon > 0.0, "epsilon should be positive");

    Ok(())
}

#[test]
fn test_searchsorted_gpu_correctness() -> anyhow::Result<()> {
    let stream = smoke_stream();
    let config = test_buffer_config(50, 4);
    let mut buf = GpuReplayBuffer::new(config, &stream)?;

    // Fill with exactly 50 experiences
    insert_random_batch(&mut buf, 50, 4)?;
    assert_eq!(buf.len(), 50);

    // Sample multiple times to exercise the searchsorted codepath (GPU-resident)
    for trial in 0..5 {
        let batch = buf.sample_proportional(10)?;
        let indices_host = dtoh_u32(&batch.indices, &stream)?;
        for (i, &idx) in indices_host.iter().enumerate() {
            assert!(
                idx < 50,
                "trial {trial}: index[{i}] = {idx} out of range [0, 50)"
            );
        }
        assert_eq!(indices_host.len(), 10);
    }

    Ok(())
}

#[test]
#[ignore] // Loads real training data — run via nightly CI or manual trigger
fn test_train_step_produces_finite_metrics() -> anyhow::Result<()> {
    let data_dir = test_data_dir()
        .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback");
    let mut trainer = smoke_trainer()?;
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;
    let metrics = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _is_best| {
        Ok(String::new())
    }))?;
    assert_finite(metrics.loss, "final_loss");
    drop(trainer);
    drop(rt);
    Ok(())
}

/// Verify training dtype is BF16 on CUDA.
#[test]
fn test_gpu_training_dtype_bf16() -> anyhow::Result<()> {
    let dtype = ml_core::native_types::NativeDType::BF16;
    assert_eq!(dtype, ml_core::native_types::NativeDType::BF16, "CUDA should use BF16 training dtype");
    Ok(())
}

/// Diagnose: GPU training dtype is BF16.
#[test]
fn test_gpu_training_dtype_diagnosis() -> anyhow::Result<()> {
    let dev = cuda_device();
    let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
    info!(device = ?dev, "GPU training dtype");

    // Verify a linear layer forward pass works
    let mut vs = GpuVarStore::new(Arc::clone(&stream));
    let layer = vs.linear("test", 48, 32)?;
    let input = GpuTensor::zeros(&[2, 48], &stream)?;
    let cublas = cudarc::cublas::CudaBlas::new(Arc::clone(&stream)).map_err(|e| anyhow::anyhow!("cuBLAS: {e}"))?;
    let (out, _activations) = layer.forward(&input, &vs, &cublas, &stream)?;
    info!(dims = ?out.dims(), "Linear forward OK");
    Ok(())
}

// test_training_rejects_missing_gpu_collector: REMOVED.
// The GPU experience collector is mandatory — there is no disable path.
// Testing "what if we disable a mandatory component" is testing dead code.

/// Diagnose: can we create + step an AdamW optimizer on GPU?
#[test]
fn test_gpu_adamw_creation() -> anyhow::Result<()> {
    let dev = cuda_device();
    let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
    info!(device = ?dev, "GPU device");

    let mut vs = GpuVarStore::new(Arc::clone(&stream));
    let layer = vs.linear("test", 48, 32)?;

    // Forward
    let input = GpuTensor::zeros(&[2, 48], &stream)?;
    let cublas = cudarc::cublas::CudaBlas::new(Arc::clone(&stream)).map_err(|e| anyhow::anyhow!("cuBLAS: {e}"))?;
    let (out, _activations) = layer.forward(&input, &vs, &cublas, &stream)?;
    let loss_val = out.mean_all(&stream)?;
    info!(loss = loss_val, "Forward OK");

    // Create optimizer
    let adamw_config = AdamWConfig {
        lr: 1e-4,
        ..Default::default()
    };
    let mut opt = GpuAdamW::new(adamw_config, Arc::clone(&stream))?;

    // Compute gradients: use a simple synthetic gradient (same shape as params)
    let mut grads = std::collections::BTreeMap::new();
    for name in vs.param_names() {
        if let Some(param) = vs.get(name) {
            let grad = GpuTensor::from_host(
                &vec![0.01_f32; param.data.len()],
                param.shape.clone(),
                &stream,
            )?;
            grads.insert(name.to_owned(), grad);
        }
    }
    let _grad_norm = opt.step(&mut vs, &grads)?;
    info!("AdamW step OK");
    Ok(())
}
