//! Plan 4 Task 4 (E.4): explicit per-branch value-decoder boundary.
//!
//! `ValueDecoder` is a thin Rust API wrapper around one of the four DQN
//! advantage branches (Direction, Magnitude, Order, Urgency). Each
//! instance forwards an encoder's `h_s2` activation through its branch's
//! FC + adv-logits heads via
//! `BatchedForward::decoder_forward_only`. It owns no parameters of its
//! own — references the trainer's per-branch weights via the
//! `weight_ptrs` slice the caller passes in.
//!
//! Plan 2 D.3 (V_short + V_long horizon-decomposed value head) is
//! already landed on the value-head side; the decoder API surfaces the
//! `BranchDecoderOutput::v_short_dev_ptr` / `v_long_dev_ptr` fields so
//! Plan 4 Task 3 (multi-Q IQN) and Task 6 (auxiliary heads) can attach
//! cleanly. For decoders that do not own a value head, callers may set
//! the `v_short` / `v_long` pointers to 0 — they are not consumed by
//! `decoder_forward_only` itself today (they pass through into the
//! struct so downstream code sees a uniform shape).
//!
//! Design intent (E.4 spec):
//! - **Additive** to `BatchedForward::forward_online_raw` — that
//!   monolithic dispatch keeps the existing multi-stream branch
//!   fork/join. `ValueDecoder::forward` is for callers that need to
//!   drive a single branch on a single stream (validation paths,
//!   per-branch experiments, future migrations). No call-site
//!   migration required for production training.
//! - 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;

/// Identifies one of the four DQN advantage branches. The `repr(usize)` is
/// chosen so `Branch::Dir as usize == 0` etc., matching the
/// `branch_idx` convention used throughout `BatchedForward`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(usize)]
pub enum Branch {
    Dir = 0,
    Mag = 1,
    Ord = 2,
    Urg = 3,
}

impl Branch {
    /// Branch index in `[0, 4)`, matching `BatchedForward`'s internal layout.
    pub fn idx(self) -> usize { self as usize }

    /// Short label, useful for diagnostics.
    pub fn label(self) -> &'static str {
        match self {
            Branch::Dir => "dir",
            Branch::Mag => "mag",
            Branch::Ord => "ord",
            Branch::Urg => "urg",
        }
    }
}

/// Output of one `ValueDecoder::forward` call.
///
/// Carries the device pointer for the per-action Q-distribution
/// (`q_per_action_dev_ptr`, shape `[B, branch_q_dim * num_atoms]` in
/// f32) plus the Plan 2 D.3 horizon-decomposed value pointers. The
/// value pointers are pass-through fields — the decoder does not own
/// them, the caller supplies them so downstream consumers see a uniform
/// type. Pass `0` for either pointer if the caller does not maintain a
/// horizon-decomposed value head for this branch.
#[derive(Debug, Clone, Copy)]
pub struct BranchDecoderOutput {
    pub q_per_action_dev_ptr: u64,
    pub q_per_action_dim: usize,
    /// Number of distributional atoms per action (C51 support size).
    pub num_atoms: usize,
    pub v_short_dev_ptr: u64,
    pub v_long_dev_ptr: u64,
    pub batch_size: usize,
}

/// Plan 4 Task 4 (E.4) per-branch value-decoder façade.
///
/// One instance per branch. Borrows a `CublasGemmSet` (`BatchedForward`)
/// and dispatches `decoder_forward_only` for the configured branch.
#[allow(missing_debug_implementations)] // CublasGemmSet borrow is not Debug
pub struct ValueDecoder<'a> {
    batched: &'a CublasGemmSet,
    branch: Branch,
}

impl<'a> ValueDecoder<'a> {
    /// Construct a decoder bound to a specific branch.
    pub fn new(batched: &'a CublasGemmSet, branch: Branch) -> Self {
        Self { batched, branch }
    }

    /// Branch this decoder targets.
    pub fn branch(&self) -> Branch { self.branch }

    /// Number of actions per branch (pre-num_atoms multiplication).
    pub fn q_per_action_dim(&self) -> usize {
        self.batched.branch_q_dim_internal(self.branch.idx())
    }

    /// Forward `h_s2` through this branch's decoder heads.
    ///
    /// Reuses `BatchedForward::decoder_forward_only`, which mirrors the
    /// per-branch loop body in `forward_online_raw` (sequential
    /// single-stream variant). Behaviorally a no-op for callers that
    /// keep using `forward_online_raw` directly.
    ///
    /// Pointer ownership (caller-supplied):
    /// - `h_s2_dev_ptr`        : `[B, shared_h2]` encoder output (read-only)
    /// - `mag_concat_dev_ptr`  : `[B, shared_h2 + branch_0_size]` magnitude
    ///   branch wider input (direction-conditioned: SH2 + b0 = trunk activation
    ///   concatenated with per-direction Q values); pass `0` for legacy
    ///   `[B, shared_h2]` input
    /// - `branch_h_dev_ptr`    : `[B, adv_h]` per-branch hidden scratch
    /// - `q_per_action_dev_ptr`: `[B, branch_q_dim * num_atoms]` output
    /// - `v_short_dev_ptr`,
    ///   `v_long_dev_ptr`      : optional horizon-decomposed value
    ///   pointers, pass `0` if not used.
    #[allow(clippy::too_many_arguments)]
    pub fn forward(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_dev_ptr: u64,
        mag_concat_dev_ptr: u64,
        weight_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        branch_h_dev_ptr: u64,
        q_per_action_dev_ptr: u64,
        v_short_dev_ptr: u64,
        v_long_dev_ptr: u64,
        batch_size: usize,
    ) -> Result<BranchDecoderOutput, MLError> {
        if batch_size != self.batched.batch_size() {
            return Err(MLError::ModelError(format!(
                "ValueDecoder::forward: batch_size {} != BatchedForward batch_size {}",
                batch_size,
                self.batched.batch_size(),
            )));
        }

        self.batched.decoder_forward_only(
            stream,
            self.branch.idx(),
            h_s2_dev_ptr,
            mag_concat_dev_ptr,
            0u64, // dir_qaux_concat_ptr: SP14 EGF wire not yet exposed via this Rust API.
            weight_ptrs,
            branch_h_dev_ptr,
            q_per_action_dev_ptr,
        )?;

        Ok(BranchDecoderOutput {
            q_per_action_dev_ptr,
            q_per_action_dim: self.batched.branch_q_dim_internal(self.branch.idx()),
            num_atoms: self.batched.num_atoms(),
            v_short_dev_ptr,
            v_long_dev_ptr,
            batch_size,
        })
    }
}
