//! Rollout buffer for on-policy PPO with GAE.
//!
//! Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.1
//! Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md
//!
//! Stores `[B × T_rollout]` trajectories of (rewards, dones, v_t, actions,
//! log_pi_old, h_t encoder snapshot) collected by the rollout-collection
//! loop. The trainer then invokes `compute_gae` (wraps the
//! `gae_backward_sweep` kernel) to compute advantages + returns, and runs
//! multi-epoch minibatch updates over the rollout.
//!
//! Phase 1B-A scope: standalone infrastructure. The struct is allocated,
//! the GAE kernel is loaded, and GPU-oracle invariant tests cover the
//! math. **No trainer integration yet** — wiring into `IntegratedTrainer`
//! lands in Phase 1B-B+.
//!
//! ## Determinism contract
//!
//! Per `pearl_determinism_achieved`: the GAE backward-sweep kernel is
//! single-thread-per-batch with a sequential reverse-T loop, so its
//! outputs are bit-equal across runs with identical inputs. No atomic
//! ops, no parallel reductions inside a trajectory, no inter-block
//! synchronisation needed.
//!
//! ## Constraints honoured
//!
//!   * `feedback_no_atomicadd` — kernel has zero atomic operations.
//!   * `feedback_cpu_is_read_only` — buffers live on the device; CPU
//!     reads only via mapped-pinned `read_slice_d_into<T>` from outside
//!     this module.
//!   * `feedback_no_htod_htoh_only_mapped_pinned` — no raw `memcpy_*` in
//!     this module. All CPU↔GPU paths go through `read_slice_d_into<T>`
//!     in the test harness.
//!   * `feedback_no_nvrtc` — pre-compiled cubin via `build.rs`.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{
    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream,
    DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg,
};

const GAE_BACKWARD_SWEEP_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/gae_backward_sweep.cubin"));

/// PPO rollout buffer + GAE backward-sweep kernel handle.
///
/// All device buffers are sized `[B × T_rollout]` (or `[B]` for the
/// bootstrap V(s_T) vector, or `[B × T × H]` for the encoder snapshot).
/// The collection loop fills the per-step buffers; `compute_gae` reads
/// rewards/dones/v_t/v_T_bootstrap and writes advantages/returns.
pub struct RolloutBuffer {
    /// Batch dimension (parallel envs).
    pub b_size: usize,
    /// Trajectory length T_rollout.
    pub t_rollout: usize,
    /// Encoder hidden dim (size of `h_t` snapshot per env-step).
    pub hidden_dim: usize,

    /// Per-step reward `r_t` ([B × T]).
    pub rewards_d: CudaSlice<f32>,
    /// Per-step done flag `1.0` if terminal at step t else `0.0`, encoded
    /// as `u8` to keep the GAE kernel's `non_terminal = 1.0f - (float)d`
    /// branch-free and 4× smaller than `f32`. ([B × T])
    pub dones_d: CudaSlice<u8>,
    /// Per-step value prediction `V(s_t)` snapshot at collection time.
    /// Frozen during the multi-epoch PPO update so the advantage targets
    /// stay consistent across epochs. ([B × T])
    pub v_t_d: CudaSlice<f32>,
    /// Per-step action index (output of action selection). ([B × T])
    pub actions_d: CudaSlice<i32>,
    /// Per-step log π_old(a_t | s_t) snapshot. PPO ratio = exp(log π_new
    /// − log π_old). ([B × T])
    pub log_pi_old_d: CudaSlice<f32>,
    /// Per-step encoder hidden state `h_t`. Multi-epoch PPO re-evaluates
    /// the heads on this snapshot rather than re-running the encoder.
    /// ([B × T × HIDDEN_DIM])
    pub h_t_d: CudaSlice<f32>,

    /// Trajectory-end bootstrap V(s_T) — the value prediction at the
    /// step immediately after the last collected step. ([B])
    pub v_t_bootstrap_d: CudaSlice<f32>,

    /// GAE advantages A_t (output of compute_gae). ([B × T])
    pub advantages_d: CudaSlice<f32>,
    /// GAE returns R_t = A_t + V_t (output of compute_gae). ([B × T])
    pub returns_d: CudaSlice<f32>,

    /// GAE backward-sweep kernel handle.
    gae_kernel: CudaFunction,
    /// Module owning the kernel — kept alive for the kernel's lifetime.
    _gae_module: Arc<CudaModule>,

    /// Write cursor for the collection loop (Phase 1B-B+). Reset at the
    /// start of each rollout and advanced by `push_step()` once that
    /// path is wired.
    pub current_t: usize,
}

impl RolloutBuffer {
    /// Allocate all device buffers and load the GAE kernel.
    ///
    /// Buffer sizes are fixed at construction; reuse across rollouts is
    /// the expected pattern (no per-rollout cudaMalloc on the hot path).
    pub fn new(
        ctx: &Arc<CudaContext>,
        stream: &Arc<CudaStream>,
        b_size: usize,
        t_rollout: usize,
        hidden_dim: usize,
    ) -> Result<Self> {
        let bt = b_size.checked_mul(t_rollout)
            .context("rollout buffer B × T overflow")?;
        let bth = bt.checked_mul(hidden_dim)
            .context("rollout buffer B × T × HIDDEN_DIM overflow")?;

        let rewards_d = stream.alloc_zeros::<f32>(bt).context("rewards_d alloc")?;
        let dones_d = stream.alloc_zeros::<u8>(bt).context("dones_d alloc")?;
        let v_t_d = stream.alloc_zeros::<f32>(bt).context("v_t_d alloc")?;
        let actions_d = stream.alloc_zeros::<i32>(bt).context("actions_d alloc")?;
        let log_pi_old_d = stream.alloc_zeros::<f32>(bt).context("log_pi_old_d alloc")?;
        let h_t_d = stream.alloc_zeros::<f32>(bth).context("h_t_d alloc")?;
        let v_t_bootstrap_d = stream.alloc_zeros::<f32>(b_size).context("v_t_bootstrap_d alloc")?;
        let advantages_d = stream.alloc_zeros::<f32>(bt).context("advantages_d alloc")?;
        let returns_d = stream.alloc_zeros::<f32>(bt).context("returns_d alloc")?;

        let gae_module = ctx
            .load_cubin(GAE_BACKWARD_SWEEP_CUBIN.to_vec())
            .context("load gae_backward_sweep cubin")?;
        let gae_kernel = gae_module
            .load_function("gae_backward_sweep")
            .context("load gae_backward_sweep function")?;

        Ok(Self {
            b_size,
            t_rollout,
            hidden_dim,
            rewards_d,
            dones_d,
            v_t_d,
            actions_d,
            log_pi_old_d,
            h_t_d,
            v_t_bootstrap_d,
            advantages_d,
            returns_d,
            gae_kernel,
            _gae_module: gae_module,
            current_t: 0,
        })
    }

    /// Reset the write cursor. Called at the start of each rollout once
    /// the collection path is wired (Phase 1B-B+).
    pub fn reset(&mut self) {
        self.current_t = 0;
    }

    /// Invoke the GAE backward-sweep kernel. Reads `rewards_d`,
    /// `dones_d`, `v_t_d`, `v_t_bootstrap_d`; writes `advantages_d`,
    /// `returns_d`.
    ///
    /// Determinism: single thread per batch with a sequential reverse-T
    /// loop, so the output is bit-equal across runs for identical inputs
    /// per `pearl_determinism_achieved`.
    pub fn compute_gae(
        &mut self,
        stream: &Arc<CudaStream>,
        gamma: f32,
        lambda: f32,
    ) -> Result<()> {
        let b = self.b_size as i32;
        let t = self.t_rollout as i32;
        let threads: u32 = 256;
        let blocks: u32 = ((self.b_size as u32) + threads - 1) / threads;
        let cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: 0,
        };

        // Resolve device pointers up-front so the launch_builder borrows
        // are scoped narrowly (the borrow guards drop after `.launch()`).
        let rewards_ptr = self.rewards_d.device_ptr(stream).0;
        let dones_ptr = self.dones_d.device_ptr(stream).0;
        let v_t_ptr = self.v_t_d.device_ptr(stream).0;
        let v_t_bootstrap_ptr = self.v_t_bootstrap_d.device_ptr(stream).0;
        let advantages_ptr = self.advantages_d.device_ptr_mut(stream).0;
        let returns_ptr = self.returns_d.device_ptr_mut(stream).0;

        unsafe {
            stream
                .launch_builder(&self.gae_kernel)
                .arg(&rewards_ptr)
                .arg(&dones_ptr)
                .arg(&v_t_ptr)
                .arg(&v_t_bootstrap_ptr)
                .arg(&advantages_ptr)
                .arg(&returns_ptr)
                .arg(&b)
                .arg(&t)
                .arg(&gamma)
                .arg(&lambda)
                .launch(cfg)
                .context("gae_backward_sweep launch")?;
        }
        Ok(())
    }
}
