//! Plan 4 Task 4 (E.4): explicit state-encoder boundary for the DQN trunk.
//!
//! `StateEncoder` is a thin Rust API wrapper around `BatchedForward`'s trunk
//! encoder portion (the two shared layers `h_s1` + `h_s2`). It owns no
//! parameters and no buffers — it borrows a `BatchedForward` and forwards
//! caller-owned device pointers through the existing
//! `BatchedForward::encoder_forward_only` dispatch.
//!
//! Design intent (E.4 spec):
//! - Provide a clean type boundary so Plan 4 Task 1 (Full VSN, attaches
//!   pre-encoder), Task 3 (multi-Q IQN, attaches at decoder), and Task 6
//!   (auxiliary heads, attaches at decoder peer) have a stable attachment
//!   point.
//! - **Additive** to the existing monolithic `forward_online_raw` — that
//!   function still routes through `encoder_forward_only` internally, so
//!   call sites remain unchanged.
//! - Pure Rust API refactor: no kernel changes, no parameter changes, no
//!   checkpoint break.

use std::sync::Arc;

use cudarc::driver::CudaStream;

use crate::MLError;
use crate::cuda_pipeline::batched_forward::CublasGemmSet;

/// Output of one `StateEncoder::forward` call.
///
/// Carries the device pointer the encoder wrote to (`h_s2_dev_ptr`), the
/// shared-trunk output dimension, and the batch size. `h_s2_dev_ptr` is
/// always equal to the pointer the caller passed in — the type exists to
/// document the contract for downstream `ValueDecoder` consumers.
#[derive(Debug, Clone, Copy)]
pub struct EncoderOutput {
    pub h_s2_dev_ptr: u64,
    pub h_s2_dim: usize,
    pub batch_size: usize,
}

/// Plan 4 Task 4 (E.4) state-encoder façade.
///
/// Borrows a `CublasGemmSet` (`BatchedForward`) and exposes the trunk
/// encoder (`h_s1` + `h_s2` GEMMs) as a separately-callable forward step.
/// The encoder runs on the caller's stream; callers own all device
/// buffers (`states_ptr`, `h_s1_ptr`, `h_s2_ptr`).
///
/// Owns no parameters of its own — references the trainer's params via
/// the `weight_ptrs` slice the caller passes in (positions 0..=3 are the
/// trunk encoder weights `W_s1, b_s1, W_s2, b_s2`).
#[allow(missing_debug_implementations)] // CublasGemmSet borrow is not Debug
pub struct StateEncoder<'a> {
    batched: &'a mut CublasGemmSet,
}

impl<'a> StateEncoder<'a> {
    /// Wrap an existing `BatchedForward` (`CublasGemmSet`) without taking
    /// ownership of any GPU resources. The returned encoder is a mutable
    /// borrow because Plan 4 Task 2c.3b's GRN encoder writes save-for-
    /// backward state into the wrapper's GrnBlock instances.
    pub fn new(batched: &'a mut CublasGemmSet) -> Self {
        Self { batched }
    }

    /// Forward-pass through the trunk encoder (`h_s1` + `h_s2` GRN blocks).
    ///
    /// Reuses `BatchedForward::encoder_forward_only`, which is the same
    /// GRN composition the existing `forward_online_raw` runs for the
    /// encoder portion (Plan 4 Task 2c.3b).
    ///
    /// `weight_ptrs` is the full DQN weight pointer array; positions
    /// `[0..=12]` are the trunk encoder GRN tensors.
    pub fn forward(
        &mut self,
        stream: &Arc<CudaStream>,
        state_dev_ptr: u64,
        weight_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        h_s1_dev_ptr: u64,
        h_s2_dev_ptr: u64,
        batch_size: usize,
    ) -> Result<EncoderOutput, MLError> {
        if batch_size != self.batched.batch_size() {
            return Err(MLError::ModelError(format!(
                "StateEncoder::forward: batch_size {} != BatchedForward batch_size {}",
                batch_size,
                self.batched.batch_size(),
            )));
        }
        self.batched.encoder_forward_only(
            stream,
            state_dev_ptr,
            weight_ptrs,
            h_s1_dev_ptr,
            h_s2_dev_ptr,
        )?;
        Ok(EncoderOutput {
            h_s2_dev_ptr,
            h_s2_dim: self.batched.shared_h2(),
            batch_size,
        })
    }

    /// Output activation dimension (`shared_h2`).
    pub fn h_s2_dim(&self) -> usize {
        self.batched.shared_h2()
    }
}
