//! Batch size resolution using GPU memory profiling.
//!
//! Combines `ModelMemoryEstimate` from `gpu::memory_profile` with
//! `AutoBatchSizer` from `memory_optimization::auto_batch_size` to
//! compute hardware-appropriate batch sizes.

use tracing::{info, warn};

use crate::gpu::capabilities::GpuCapabilities;
use crate::gpu::memory_profile::ModelMemoryEstimate;

use crate::memory_optimization::auto_batch_size::{
    AutoBatchSizer, BatchSizeConfig, ModelPrecision, OptimizerType,
};

/// Resolve the optimal batch size for a model given GPU capabilities.
///
/// Uses the existing `AutoBatchSizer` math from `memory_optimization`.
/// If the GPU has no VRAM (CPU mode), returns the requested batch size unchanged.
///
/// # Arguments
/// * `capabilities` - Detected GPU hardware (from `GpuCapabilities::detect()`)
/// * `estimate` - Memory footprint of the model to train
/// * `requested_batch_size` - The batch size the user/config requested (may be shrunk)
///
/// # Returns
/// The effective batch size (may be smaller than requested if VRAM is insufficient).
pub fn resolve_batch_size(
    capabilities: &GpuCapabilities,
    estimate: &ModelMemoryEstimate,
    requested_batch_size: usize,
) -> usize {
    if !capabilities.is_cuda {
        info!(
            "CPU mode: using requested batch_size {} for {}",
            requested_batch_size, estimate.name
        );
        return requested_batch_size;
    }

    // Scale model memory by activation multiplier to account for attention layers etc.
    // Models with attention (TFT=2.5x, TLOB=2.0x) need more memory for intermediate activations.
    let model_size_mb = estimate.estimated_size_mb() * estimate.activation_multiplier;

    let batch_config = BatchSizeConfig {
        model_memory_mb: model_size_mb,
        model_precision: ModelPrecision::FP32,
        sequence_length: estimate.default_seq_len,
        feature_dim: estimate.default_feature_dim,
        gradient_checkpointing: estimate.supports_checkpointing,
        optimizer_type: OptimizerType::AdamW,
        safety_margin: 0.20,
        min_batch_size: 1,
        max_batch_size: requested_batch_size,
    };

    let sizer = AutoBatchSizer::with_manual_memory(
        capabilities.total_vram_mb,
        capabilities.free_vram_mb,
        capabilities.device_name.clone(),
    );

    match sizer.calculate_optimal_batch_size(&batch_config) {
        Ok(optimal) => {
            if optimal < requested_batch_size {
                warn!(
                    "{}: batch_size {} exceeds GPU capacity ({} {:.0} MB free), shrinking to {}",
                    estimate.name,
                    requested_batch_size,
                    capabilities.device_name,
                    capabilities.free_vram_mb,
                    optimal
                );
            } else {
                info!(
                    "{}: batch_size {} fits on {} ({:.0} MB free)",
                    estimate.name,
                    requested_batch_size,
                    capabilities.device_name,
                    capabilities.free_vram_mb
                );
            }
            optimal
        }
        Err(e) => {
            warn!(
                "{}: GPU memory insufficient ({}), falling back to batch_size 1",
                estimate.name, e
            );
            1
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::gpu::memory_profile::estimates;

    #[test]
    fn test_resolve_batch_size_cpu_passthrough() {
        let caps = GpuCapabilities::cpu_fallback();
        let batch = resolve_batch_size(&caps, &estimates::DQN, 256);
        assert_eq!(batch, 256, "CPU mode should pass through requested batch size");
    }

    #[test]
    fn test_resolve_batch_size_small_gpu() {
        let caps = GpuCapabilities::with_vram("RTX 3050 Ti", 4096.0, 3700.0);
        let batch = resolve_batch_size(&caps, &estimates::DQN, 512);
        assert!(batch >= 32, "DQN on 4GB GPU should handle at least 32, got {}", batch);
    }

    #[test]
    fn test_resolve_batch_size_tft_constrained() {
        let caps = GpuCapabilities::with_vram("RTX 3050 Ti", 4096.0, 3700.0);
        let batch_tft = resolve_batch_size(&caps, &estimates::TFT, 512);
        let batch_dqn = resolve_batch_size(&caps, &estimates::DQN, 512);
        assert!(
            batch_tft <= batch_dqn,
            "TFT ({}) should get <= batch_size than DQN ({}) on 4GB",
            batch_tft, batch_dqn
        );
    }

    #[test]
    fn test_resolve_batch_size_large_gpu() {
        let caps = GpuCapabilities::with_vram("A100-SXM4-80GB", 81920.0, 80000.0);
        let batch = resolve_batch_size(&caps, &estimates::TFT, 256);
        // A100 should handle TFT at a reasonable batch size (AutoBatchSizer rounds to power-of-2)
        // The key invariant: batch should be > 0 and not shrunk to 1
        assert!(batch >= 64, "A100 should handle large TFT batches, got {}", batch);
    }
}
