#![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.

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

use crate::MLError;

static EPSILON_GREEDY_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();

fn compile_kernel_ptx(context: &CudaContext) -> Result<Ptx, String> {
    let defines = "\
        #define STATE_DIM 72\n\
        #define MARKET_DIM 42\n\
        #define PORTFOLIO_DIM 8\n";
    let common_src = include_str!("common_device_functions.cuh");
    let kernel_src = include_str!("epsilon_greedy_kernel.cu");
    let full_source = format!("{defines}{common_src}\n{kernel_src}");
    crate::cuda_pipeline::compile_ptx_for_device(&full_source, context)
}

/// 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_states: CudaSlice<u32>,
    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: [5], [3], [3] on GPU.
    /// Device pointer = 0 means disabled (kernel treats NULL as no bonus).
    bonus_exposure_ptr: u64,
    bonus_order_ptr: u64,
    bonus_urgency_ptr: u64,
    bonus_exposure_buf: Option<CudaSlice<f32>>,
    bonus_order_buf: Option<CudaSlice<f32>>,
    bonus_urgency_buf: Option<CudaSlice<f32>>,
}

impl GpuActionSelector {
    pub fn new(stream: Arc<CudaStream>, max_batch_size: usize, seed: u64) -> Result<Self, MLError> {
        let context = stream.context();
        let ptx_result = EPSILON_GREEDY_PTX.get_or_init(|| compile_kernel_ptx(&context));
        let ptx = ptx_result.as_ref().map_err(|e| MLError::ModelError(format!("epsilon_greedy PTX: {e}")))?;
        let module = context.load_module(ptx.clone()).map_err(|e| MLError::ModelError(format!("epsilon_greedy module load: {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 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}")))?;
        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);
        }
        let mut rng_states = stream.alloc_zeros::<u32>(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?;
        stream.memcpy_htod(&rng_seeds, &mut rng_states).map_err(|e| MLError::ModelError(format!("upload rng_states: {e}")))?;
        info!("GpuActionSelector initialized: max_batch_size={max_batch_size}, kernel compiled");
        Ok(Self { kernel_func, routed_kernel_func, branching_kernel_func, route_func, rng_states, 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,
        })
    }

    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 empty slices or zeros to disable.
    pub fn set_count_bonuses(&mut self, exposure: &[f32; 5], order: &[f32; 3], urgency: &[f32; 3]) -> Result<(), MLError> {
        // Allocate or reuse GPU buffers
        let be = match self.bonus_exposure_buf.take() {
            Some(mut buf) => { self.stream.memcpy_htod(exposure, &mut buf).map_err(|e| MLError::ModelError(format!("bonus_exposure upload: {e}")))?; buf }
            None => self.stream.clone_htod(exposure).map_err(|e| MLError::ModelError(format!("bonus_exposure alloc: {e}")))?,
        };
        let bo = match self.bonus_order_buf.take() {
            Some(mut buf) => { self.stream.memcpy_htod(order, &mut buf).map_err(|e| MLError::ModelError(format!("bonus_order upload: {e}")))?; buf }
            None => self.stream.clone_htod(order).map_err(|e| MLError::ModelError(format!("bonus_order alloc: {e}")))?,
        };
        let bu = match self.bonus_urgency_buf.take() {
            Some(mut buf) => { self.stream.memcpy_htod(urgency, &mut buf).map_err(|e| MLError::ModelError(format!("bonus_urgency upload: {e}")))?; buf }
            None => self.stream.clone_htod(urgency).map_err(|e| MLError::ModelError(format!("bonus_urgency alloc: {e}")))?,
        };
        self.bonus_exposure_ptr = be.device_ptr(&self.stream).0;
        self.bonus_order_ptr = bo.device_ptr(&self.stream).0;
        self.bonus_urgency_ptr = bu.device_ptr(&self.stream).0;
        self.bonus_exposure_buf = Some(be);
        self.bonus_order_buf = Some(bo);
        self.bonus_urgency_buf = Some(bu);
        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;
        unsafe {
            self.stream.launch_builder(&self.kernel_func)
                .arg(q_values).arg(&mut self.rng_states).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;
        unsafe {
            self.stream.launch_builder(&self.routed_kernel_func)
                .arg(q_values).arg(&mut self.rng_states).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))); }
        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;
        unsafe {
            self.stream.launch_builder(&self.branching_kernel_func)
                .arg(q_exposure).arg(q_order).arg(q_urgency)
                .arg(&mut self.rng_states).arg(&mut self.actions_buf).arg(&epsilon).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_ptx_compilation() {
        let stream = cuda_stream();
        let context = stream.context();
        let result = compile_kernel_ptx(&context);
        if let Err(ref e) = result { panic!("PTX compilation failed: {e}"); }
    }

    #[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 random 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");
        }
    }
}
