#![allow(unsafe_code)] // Required for CUDA kernel launch

//! GPU-fused epsilon-greedy action selection -- pure cudarc.
//!
//! All inputs are `CudaSlice<f32>` Q-values, all outputs are `CudaSlice<u32>` action
//! indices. The selector stores an `Arc<CudaStream>`.
//!
//! Pure cudarc implementation with no Candle dependency.
//! Kernel loaded from build-time precompiled cubin (zero runtime nvcc).
//!
//! NoisyNet exploration: Q-values received by the action selector already
//! include factorized Gaussian noise from NoisyLinear layers in the network
//! forward pass. Epsilon-greedy acts as a secondary fallback, not the primary
//! exploration mechanism when NoisyNet is enabled. Explicit noise injection
//! here would double-count exploration.

use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use std::sync::Arc;
use tracing::info;

use crate::MLError;
use super::mapped_pinned::{MappedF32Buffer, MappedU32Buffer};

/// Precompiled epsilon_greedy_kernel cubin, embedded at compile time by build.rs.
// SP22 H6 (2026-05-12): visibility lifted to `pub(crate)` so the collector
// can load the `fill_f32` kernel for `prev_aux_dir_prob` 0.5 sentinel init
// (cold-start + FoldReset). Same cubin, multiple in-crate consumers per the
// existing pattern shared by SP13/SP14 CUBIN statics in `gpu_dqn_trainer.rs`.
pub(crate) static EPSILON_GREEDY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/epsilon_greedy_kernel.cubin"));

/// GPU-fused epsilon-greedy action selector -- pure cudarc API.
pub struct GpuActionSelector {
    kernel_func: CudaFunction,
    routed_kernel_func: CudaFunction,
    branching_kernel_func: CudaFunction,
    route_func: CudaFunction,
    /// RNG state, mapped pinned. Kernel reads & writes via `rng_states_ptr`;
    /// the seed array is initialised via direct host_ptr writes (no HtoD).
    rng_states: MappedU32Buffer,
    rng_states_ptr: u64,
    actions_buf: CudaSlice<u32>,
    fill_mask_buf: CudaSlice<i32>,
    max_batch_size: usize,
    stream: Arc<CudaStream>,
    q_gap_threshold: f32,
    /// UCB count bonus for branching heads: [4], [3], [3] in mapped pinned memory.
    /// Device pointer = 0 means disabled (kernel treats NULL as no bonus).
    /// Mapped pinned per `feedback_no_htod_htoh_only_mapped_pinned.md` —
    /// `set_count_bonuses` is hot (per action-selection call). Direct host_ptr
    /// writes via `write_from_slice` eliminate the per-call HtoD memcpy.
    bonus_exposure_ptr: u64,
    bonus_order_ptr: u64,
    bonus_urgency_ptr: u64,
    bonus_exposure_buf: Option<MappedF32Buffer>,
    bonus_order_buf: Option<MappedF32Buffer>,
    bonus_urgency_buf: Option<MappedF32Buffer>,
    /// Per-sample epsilon buffer for branching kernel: [max_batch_size] on GPU.
    /// Filled via GPU-side fill_f32 kernel before each branching launch.
    epsilon_buf: CudaSlice<f32>,
    /// GPU-side fill kernel — eliminates per-step HtoD memcpy for epsilon buffer.
    fill_f32_kernel: CudaFunction,
}

impl GpuActionSelector {
    pub fn new(stream: Arc<CudaStream>, max_batch_size: usize, seed: u64) -> Result<Self, MLError> {
        let context = stream.context();
        let module = context.load_cubin(EPSILON_GREEDY_CUBIN.to_vec()).map_err(|e| MLError::ModelError(format!("load epsilon_greedy cubin: {e}")))?;
        let kernel_func = module.load_function("epsilon_greedy_select").map_err(|e| MLError::ModelError(format!("epsilon_greedy function load: {e}")))?;
        let routed_kernel_func = module.load_function("epsilon_greedy_routed").map_err(|e| MLError::ModelError(format!("epsilon_greedy_routed function load: {e}")))?;
        let branching_kernel_func = module.load_function("branching_action_select").map_err(|e| MLError::ModelError(format!("branching_action_select function load: {e}")))?;
        let route_func = module.load_function("batch_route_exposure_to_factored").map_err(|e| MLError::ModelError(format!("batch_route_exposure_to_factored function load: {e}")))?;
        let fill_f32_kernel = module.load_function("fill_f32").map_err(|e| MLError::ModelError(format!("fill_f32 function load: {e}")))?;
        let actions_buf = stream.alloc_zeros::<u32>(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc actions_buf: {e}")))?;
        let fill_mask_buf = stream.alloc_zeros::<i32>(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc fill_mask_buf: {e}")))?;
        // RNG state — mapped pinned per `feedback_no_htod_htoh_only_mapped_pinned.md`.
        // Seeds written via direct host_ptr writes; kernel reads & writes
        // back via rng_states_ptr (CUdeviceptr).
        // Safety: a CUDA context is active (we just resolved kernels through it).
        let rng_states = unsafe { MappedU32Buffer::new(max_batch_size) }
            .map_err(|e| MLError::ModelError(format!("alloc rng_states mapped pinned: {e}")))?;
        let rng_states_ptr = rng_states.dev_ptr;
        {
            let mut rng_seeds = Vec::with_capacity(max_batch_size);
            for i in 0..max_batch_size {
                let s = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(i as u64);
                rng_seeds.push((s >> 32) as u32 | 1);
            }
            rng_states.write_from_slice(&rng_seeds);
        }
        let epsilon_buf = stream.alloc_zeros::<f32>(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc epsilon_buf: {e}")))?;
        info!("GpuActionSelector initialized: max_batch_size={max_batch_size}, precompiled cubin loaded");
        Ok(Self { kernel_func, routed_kernel_func, branching_kernel_func, route_func, rng_states, rng_states_ptr, actions_buf, fill_mask_buf, max_batch_size, stream, q_gap_threshold: 0.0,
            bonus_exposure_ptr: 0, bonus_order_ptr: 0, bonus_urgency_ptr: 0,
            bonus_exposure_buf: None, bonus_order_buf: None, bonus_urgency_buf: None,
            epsilon_buf,
            fill_f32_kernel,
        })
    }

    pub fn stream(&self) -> &Arc<CudaStream> { &self.stream }

    /// Set the Q-gap conviction threshold (0.0 = disabled).
    pub fn set_q_gap_threshold(&mut self, threshold: f32) { self.q_gap_threshold = threshold; }

    /// Upload per-branch UCB count bonuses for branching action selection.
    /// Bonuses are added to Q-values before argmax in the greedy path.
    /// Pass zeros to disable (kernel sees zero contribution; same effect as NULL).
    ///
    /// HOT path: called per action-selection. Mapped pinned buffers (allocated
    /// lazily on first call) make subsequent updates pure host_ptr writes —
    /// no HtoD memcpy. Per `feedback_no_htod_htoh_only_mapped_pinned.md`.
    pub fn set_count_bonuses(&mut self, exposure: &[f32; 4], order: &[f32; 3], urgency: &[f32; 3]) -> Result<(), MLError> {
        // Lazy-allocate mapped pinned buffers on first call. Safety: a CUDA
        // context is active because `self.stream` was constructed against it.
        if self.bonus_exposure_buf.is_none() {
            let buf = unsafe { MappedF32Buffer::new(4) }
                .map_err(|e| MLError::ModelError(format!("alloc bonus_exposure mapped pinned: {e}")))?;
            self.bonus_exposure_ptr = buf.dev_ptr;
            self.bonus_exposure_buf = Some(buf);
        }
        if self.bonus_order_buf.is_none() {
            let buf = unsafe { MappedF32Buffer::new(3) }
                .map_err(|e| MLError::ModelError(format!("alloc bonus_order mapped pinned: {e}")))?;
            self.bonus_order_ptr = buf.dev_ptr;
            self.bonus_order_buf = Some(buf);
        }
        if self.bonus_urgency_buf.is_none() {
            let buf = unsafe { MappedF32Buffer::new(3) }
                .map_err(|e| MLError::ModelError(format!("alloc bonus_urgency mapped pinned: {e}")))?;
            self.bonus_urgency_ptr = buf.dev_ptr;
            self.bonus_urgency_buf = Some(buf);
        }
        // Direct host_ptr writes — no memcpy. Kernel reads via dev_ptr after
        // stream sync barrier (mapped pinned coherence).
        self.bonus_exposure_buf.as_ref().unwrap().write_from_slice(exposure);
        self.bonus_order_buf.as_ref().unwrap().write_from_slice(order);
        self.bonus_urgency_buf.as_ref().unwrap().write_from_slice(urgency);
        Ok(())
    }

    pub fn select_actions(&mut self, q_values: &CudaSlice<f32>, epsilon: f32, batch_size: usize, num_actions: usize) -> Result<CudaSlice<u32>, MLError> {
        if batch_size == 0 { return self.stream.alloc_zeros::<u32>(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); }
        if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {batch_size} exceeds max_batch_size {}", self.max_batch_size))); }
        let config = launch_config_1d(batch_size);
        let bs_i32 = batch_size as i32;
        let na_i32 = num_actions as i32;
        let q_gap = self.q_gap_threshold;
        let rng_ptr = self.rng_states_ptr;
        unsafe {
            self.stream.launch_builder(&self.kernel_func)
                .arg(q_values).arg(&rng_ptr).arg(&mut self.actions_buf)
                .arg(&epsilon).arg(&bs_i32).arg(&na_i32).arg(&q_gap)
                .launch(config).map_err(|e| MLError::ModelError(format!("epsilon_greedy kernel launch: {e}")))?;
        }
        self.copy_actions_out(batch_size)
    }

    #[allow(clippy::too_many_arguments)]
    pub fn select_actions_routed(&mut self, q_values: &CudaSlice<f32>, epsilon: f32, batch_size: usize, num_actions: usize,
        step_offset: i32, spread: f32, median_spread: f32, volatility: f32, median_vol: f32,
        spread_bps: f32, ioc_fill_prob: f32, limit_fill_min: f32, limit_fill_max: f32,
        spread_cost_frac: f32, spread_capture_frac: f32,
    ) -> Result<CudaSlice<u32>, MLError> {
        if batch_size == 0 { return self.stream.alloc_zeros::<u32>(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); }
        if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {batch_size} exceeds max_batch_size {}", self.max_batch_size))); }
        let config = launch_config_1d(batch_size);
        let bs_i32 = batch_size as i32;
        let na_i32 = num_actions as i32;
        let rng_ptr = self.rng_states_ptr;
        unsafe {
            self.stream.launch_builder(&self.routed_kernel_func)
                .arg(q_values).arg(&rng_ptr).arg(&mut self.actions_buf).arg(&mut self.fill_mask_buf)
                .arg(&epsilon).arg(&bs_i32).arg(&na_i32).arg(&step_offset)
                .arg(&spread).arg(&median_spread).arg(&volatility).arg(&median_vol)
                .arg(&spread_bps).arg(&ioc_fill_prob).arg(&limit_fill_min).arg(&limit_fill_max)
                .arg(&spread_cost_frac).arg(&spread_capture_frac).arg(&self.q_gap_threshold)
                .launch(config).map_err(|e| MLError::ModelError(format!("epsilon_greedy_routed kernel launch: {e}")))?;
        }
        self.copy_actions_out(batch_size)
    }

    pub fn select_actions_branching(&mut self, q_exposure: &CudaSlice<f32>, q_order: &CudaSlice<f32>,
        q_urgency: &CudaSlice<f32>, epsilon: f32, batch_size: usize,
    ) -> Result<CudaSlice<u32>, MLError> {
        if batch_size == 0 { return self.stream.alloc_zeros::<u32>(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); }
        if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {} exceeds max {}", batch_size, self.max_batch_size))); }
        // GPU-side fill: write uniform epsilon value directly — no HtoD memcpy.
        let eps_ptr = self.epsilon_buf.raw_ptr();
        let bs_fill = batch_size as i32;
        unsafe {
            self.stream.launch_builder(&self.fill_f32_kernel)
                .arg(&eps_ptr)
                .arg(&epsilon)
                .arg(&bs_fill)
                .launch(launch_config_1d(batch_size))
                .map_err(|e| MLError::ModelError(format!("fill_f32 epsilon: {e}")))?;
        }
        let config = launch_config_1d(batch_size);
        let bs_i32 = batch_size as i32;
        let be_ptr = self.bonus_exposure_ptr;
        let bo_ptr = self.bonus_order_ptr;
        let bu_ptr = self.bonus_urgency_ptr;
        let rng_ptr = self.rng_states_ptr;
        unsafe {
            self.stream.launch_builder(&self.branching_kernel_func)
                .arg(q_exposure).arg(q_order).arg(q_urgency)
                .arg(&rng_ptr).arg(&mut self.actions_buf)
                .arg(&self.epsilon_buf)
                .arg(&bs_i32).arg(&self.q_gap_threshold)
                .arg(&be_ptr).arg(&bo_ptr).arg(&bu_ptr) // UCB count bonuses (0 = NULL = disabled)
                .launch(config).map_err(|e| MLError::ModelError(format!("branching kernel launch: {e}")))?;
        }
        self.copy_actions_out(batch_size)
    }

    #[allow(clippy::too_many_arguments)]
    pub fn route_exposure_to_factored(&mut self, exposure_actions: &CudaSlice<u32>, batch_size: usize,
        spread: f32, median_spread: f32, volatility: f32, median_volatility: f32,
    ) -> Result<CudaSlice<u32>, MLError> {
        if batch_size == 0 { return self.stream.alloc_zeros::<u32>(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); }
        if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {batch_size} exceeds max_batch_size {}", self.max_batch_size))); }
        let config = launch_config_1d(batch_size);
        let n_i32 = batch_size as i32;
        unsafe {
            self.stream.launch_builder(&self.route_func)
                .arg(exposure_actions).arg(&mut self.actions_buf)
                .arg(&spread).arg(&median_spread).arg(&volatility).arg(&median_volatility).arg(&n_i32)
                .launch(config).map_err(|e| MLError::ModelError(format!("batch_route kernel launch: {e}")))?;
        }
        self.copy_actions_out(batch_size)
    }

    pub fn readback_actions(stream: &Arc<CudaStream>, actions: &CudaSlice<u32>, count: usize) -> Result<Vec<u32>, MLError> {
        let view = actions.slice(..count);
        let mut host = vec![0_u32; count];
        stream.memcpy_dtoh(&view, &mut host).map_err(|e| MLError::ModelError(format!("DtoH readback actions: {e}")))?; // gpu-exit: action indices
        Ok(host)
    }

    fn copy_actions_out(&self, batch_size: usize) -> Result<CudaSlice<u32>, MLError> {
        let out = self.stream.alloc_zeros::<u32>(batch_size).map_err(|e| MLError::ModelError(format!("alloc output slice: {e}")))?;
        let src_view = self.actions_buf.slice(..batch_size);
        let num_bytes = batch_size * std::mem::size_of::<u32>();
        {
            let (dst_ptr, _dst_sync) = out.device_ptr(&self.stream);
            let (src_ptr, _src_sync) = src_view.device_ptr(&self.stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(dst_ptr, src_ptr, num_bytes, self.stream.cu_stream())
                    .map_err(|e| MLError::ModelError(format!("DtoD copy actions out: {e}")))?;
            }
        }
        Ok(out)
    }
}

impl std::fmt::Debug for GpuActionSelector {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GpuActionSelector").field("max_batch_size", &self.max_batch_size).finish()
    }
}

fn launch_config_1d(n: usize) -> LaunchConfig {
    let threads_per_block = 256_u32;
    let blocks = (n as u32).div_ceil(threads_per_block);
    LaunchConfig { grid_dim: (blocks.max(1), 1, 1), block_dim: (threads_per_block, 1, 1), shared_mem_bytes: 0 }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn cuda_stream() -> Arc<CudaStream> {
        let context = cudarc::driver::CudaContext::new(0).expect("CUDA context required");
        context.default_stream()
    }

    #[test]
    fn test_cubin_loads() {
        let stream = cuda_stream();
        let context = stream.context();
        let module = context.load_cubin(EPSILON_GREEDY_CUBIN.to_vec()).expect("load precompiled cubin");
        module.load_function("epsilon_greedy_select").expect("load epsilon_greedy_select");
        module.load_function("epsilon_greedy_routed").expect("load epsilon_greedy_routed");
        module.load_function("branching_action_select").expect("load branching_action_select");
        module.load_function("batch_route_exposure_to_factored").expect("load batch_route_exposure_to_factored");
        module.load_function("fill_f32").expect("load fill_f32");
    }

    #[test]
    fn test_fused_epsilon_greedy_gpu() {
        let stream = cuda_stream();
        let batch_size = 32;
        let num_actions = 5;
        let mut selector = GpuActionSelector::new(stream.clone(), batch_size, 12345).expect("init");
        // Allocate Q-values on GPU via host upload
        let q_host: Vec<f32> = (0..batch_size * num_actions).map(|i| (i as f32) * 0.1 - 8.0).collect();
        let mut q_buf = stream.alloc_zeros::<f32>(batch_size * num_actions).expect("alloc q_values");
        stream.memcpy_htod(&q_host, &mut q_buf).expect("upload q_values");
        let greedy_actions = selector.select_actions(&q_buf, 0.0, batch_size, num_actions).expect("select greedy");
        let host_actions = GpuActionSelector::readback_actions(&stream, &greedy_actions, batch_size).expect("readback");
        for (i, &a) in host_actions.iter().enumerate() {
            assert!((a as usize) < num_actions, "action [{i}]={a} out of range");
        }
    }
}
