#![allow(unsafe_code)] // Required for CUDA DtoD memcpy via cudarc.

//! GPU-resident replay buffer — the only replay path.
//!
//! All experience data lives as contiguous GPU tensors. Sampling and
//! priority updates happen entirely on GPU — zero CPU round-trips.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};

/// Zero-alloc GPU batch: raw device pointers to pre-allocated replay buffer memory.
/// Valid until the next `sample()` call (buffers are reused each sample).
/// No cuMemAlloc, no DtoD clone, no CudaSlice ownership.
#[derive(Debug, Clone, Copy)]
pub struct GpuBatch {
    /// [batch_size, state_dim] f32 on GPU (pre-allocated in GpuReplayBuffer)
    pub states_ptr: u64,
    /// [batch_size] i32 on GPU (reinterpreted from u32 actions)
    pub actions_ptr: u64,
    /// [batch_size] f32 on GPU
    pub rewards_ptr: u64,
    /// [batch_size, state_dim] f32 on GPU
    pub next_states_ptr: u64,
    /// [batch_size] f32 on GPU (0.0/1.0)
    pub dones_ptr: u64,
    /// [batch_size] f32 on GPU (IS weights)
    pub weights_ptr: u64,
    /// [batch_size] u32 on GPU (buffer indices for PER update)
    pub indices_ptr: u64,
    /// [batch_size] i32 on GPU (episode IDs for HER)
    pub episode_ids_ptr: u64,
    /// Actual number of samples in this batch (may differ from trainer's max batch_size)
    pub batch_size: usize,
}

/// GPU-resident prioritized replay buffer.
/// Experiences are inserted directly via DtoD f32 copies.
/// No CPU staging — all data stays on GPU.
pub struct StagedGpuBuffer {
    pub gpu: GpuReplayBuffer,
}

impl std::fmt::Debug for StagedGpuBuffer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "StagedGpuBuffer(capacity={})", self.gpu.capacity())
    }
}

impl StagedGpuBuffer {
    /// Allocate GPU PER buffer with adaptive capacity halving on OOM.
    pub fn try_with_halving(
        capacity: usize,
        alpha: f64,
        beta: f64,
        beta_max: f64,
        beta_annealing_steps: usize,
        max_memory_bytes: usize,
        batch_size: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let min_capacity = batch_size.max(1);
        let mut try_cap = capacity;
        let state_dim = ml_core::state_layout::STATE_DIM;

        while try_cap >= min_capacity {
            let config = GpuReplayBufferConfig {
                capacity: try_cap,
                alpha: alpha as f32,
                beta_start: beta as f32,
                beta_max: beta_max as f32,
                beta_annealing_steps,
                epsilon: 1e-6,
                max_memory_bytes,
                max_batch_size: batch_size,
            };
            match GpuReplayBuffer::new(config, stream) {
                Ok(buffer) => {
                    if try_cap < capacity {
                        tracing::warn!(
                            "GPU PER replay buffer allocated at reduced capacity ({} -> {}, {} state_dim)",
                            capacity, try_cap, state_dim
                        );
                    } else {
                        tracing::info!(
                            "GPU PER replay buffer allocated ({} capacity, {} state_dim)",
                            try_cap, state_dim
                        );
                    }
                    return Ok(Self { gpu: buffer });
                }
                Err(e) => {
                    tracing::warn!(
                        "GPU PER at capacity {} failed ({}), retrying at {}",
                        try_cap, e, try_cap / 2
                    );
                    try_cap /= 2;
                    if try_cap < min_capacity {
                        return Err(e);
                    }
                }
            }
        }

        Err(MLError::ModelError(format!(
            "GPU PER allocation failed: capacity {} below minimum {} (batch_size)",
            try_cap, min_capacity
        )))
    }

    /// Sample a batch from the GPU PER buffer.
    pub fn sample(&mut self, batch_size: usize) -> Result<GpuBatch, MLError> {
        let ptrs = self.gpu.sample_proportional(batch_size)?;
        Ok(GpuBatch {
            states_ptr: ptrs.states_ptr,
            actions_ptr: ptrs.actions_ptr,
            rewards_ptr: ptrs.rewards_ptr,
            next_states_ptr: ptrs.next_states_ptr,
            dones_ptr: ptrs.dones_ptr,
            weights_ptr: ptrs.weights_ptr,
            indices_ptr: ptrs.indices_ptr,
            episode_ids_ptr: ptrs.episode_ids_ptr,
            batch_size,
        })
    }

    /// Step training counter for beta annealing.
    pub fn step(&mut self) {
        self.gpu.step();
    }

    /// Update PER priorities from GPU-resident TD error tensors.
    pub fn update_priorities_gpu(
        &mut self,
        indices: &cudarc::driver::CudaSlice<u32>,
        td_errors: &cudarc::driver::CudaSlice<f32>,
        bs: usize,
    ) -> Result<(), MLError> {
        if bs == 0 { return Ok(()); }
        self.gpu.update_priorities_gpu(indices, td_errors, bs, None)
    }

    /// Update PER priorities using td_errors from the last sample.
    ///
    /// `ext_stream`: the stream on which `td_errors` was computed. All GPU work
    /// runs on this stream to avoid cross-stream synchronization.
    pub fn update_priorities_from_td(
        &mut self,
        td_errors: &cudarc::driver::CudaSlice<f32>,
        batch_size: usize,
        ext_stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        if batch_size == 0 { return Ok(()); }
        self.gpu.update_priorities_gpu_raw(0, td_errors, batch_size, Some(ext_stream))
    }

    /// Flush GPU-accumulated max priority to CPU (single scalar readback).
    /// Call once per epoch after all priority updates.
    pub fn flush_max_priority(&mut self) -> Result<(), MLError> {
        self.gpu.flush_max_priority()
    }

    /// PER alpha/epsilon for GPU-native priority update kernel.
    pub fn per_alpha_epsilon(&self) -> (f32, f32) {
        (self.gpu.alpha(), self.gpu.epsilon())
    }

    /// Set `max_priority` from a CPU scalar (epoch-boundary flush from CUDA kernel).
    pub fn apply_max_priority_scalar(&mut self, max_prio: f32) -> Result<(), MLError> {
        self.gpu.apply_max_priority_scalar(max_prio)
    }

    /// Get current beta value for importance sampling.
    pub fn current_beta(&self) -> f32 {
        self.gpu.current_beta()
    }

    /// Get buffer size (number of stored transitions).
    pub fn len(&self) -> usize {
        self.gpu.len()
    }

    /// Check if buffer is empty.
    pub fn is_empty(&self) -> bool {
        self.gpu.len() == 0
    }

    /// Clear all experiences from the buffer.
    pub fn clear(&mut self) -> Result<(), MLError> {
        self.gpu.clear()
    }

    /// Check if buffer can sample given batch size.
    pub fn can_sample(&self, batch_size: usize) -> bool {
        self.gpu.len() >= batch_size
    }

    /// Get current buffer capacity.
    pub fn capacity(&self) -> usize {
        self.gpu.capacity()
    }

    /// Get current write cursor position (for staleness computation).
    pub fn write_cursor(&self) -> usize {
        self.gpu.write_cursor()
    }
}
