//! ComputeBackend trait for abstracting GPU operations (Phase 41)

//!

//! Enables mocking GPU host code for coverage testing.

use std::error::Error;

/// Result type for backend operations

pub type BackendResult<T> = Result<T, Box<dyn Error + Send + Sync>>;

/// Abstraction over GPU compute backends (CUDA, Mock, future: Metal, HIP)

///

/// This trait allows swapping out the actual GPU implementation with a mock

/// for testing purposes, enabling coverage of GPU host code without requiring

/// actual GPU hardware.

///

/// # Example

///

/// ```rust,ignore

/// use realizar::gpu::backend::{ComputeBackend, BackendResult};

///

/// struct MockBackend {

///     weights: std::collections::HashMap<String, Vec<f32>>,

/// }

///

/// impl ComputeBackend for MockBackend {

///     fn is_available() -> bool { true }

///     fn new(_device_id: u32) -> BackendResult<Self> {

///         Ok(Self { weights: Default::default() })

///     }

///     // ... implement other methods

/// }

/// ```

pub trait ComputeBackend: Send {

    /// Check if this backend type is available on the system

    fn is_available() -> bool

    where

        Self: Sized;

    /// Create a new backend instance for the given device

    ///

    /// # Errors

    ///

    /// Returns error if device initialization fails.

    fn new(device_id: u32) -> BackendResult<Self>

    where

        Self: Sized;

    /// Get device name

    fn device_name(&self) -> String;

    // =========================================================================

    // Weight Management

    // =========================================================================

    /// Load f32 weights to device memory

    ///

    /// Returns a handle/index for the loaded weights.

    ///

    /// # Errors

    ///

    /// Returns error if memory allocation or transfer fails.

    fn load_weights(&mut self, name: &str, data: &[f32]) -> BackendResult<usize>;

    /// Load quantized weights to device memory

    ///

    /// # Arguments

    ///

    /// * `name` - Unique identifier for the weights

    /// * `data` - Raw quantized bytes

    /// * `qtype` - Quantization type (e.g., 2 for Q4_0, 3 for Q4_K)

    ///

    /// # Errors

    ///

    /// Returns error if memory allocation or transfer fails.

    fn load_quantized_weights(&mut self, name: &str, data: &[u8], qtype: u32)

        -> BackendResult<usize>;

    /// Check if weights are loaded

    fn has_weights(&self, name: &str) -> bool;

    /// Clear all cached weights from device memory

    fn clear_weights(&mut self);

    /// Get count of cached weights

    fn cached_weight_count(&self) -> usize;

    // =========================================================================

    // Core Compute

    // =========================================================================

    /// Matrix multiplication: C = A @ B

    ///

    /// # Arguments

    ///

    /// * `a` - Left matrix, row-major [m, k]

    /// * `b` - Right matrix, row-major [k, n]

    /// * `m` - Rows in A

    /// * `k` - Cols in A / Rows in B

    /// * `n` - Cols in B

    ///

    /// # Errors

    ///

    /// Returns error if dimensions are invalid or compute fails.

    fn matmul(&mut self, a: &[f32], b: &[f32], m: u32, k: u32, n: u32) -> BackendResult<Vec<f32>>;

    /// Matrix multiplication using cached weights

    ///

    /// # Arguments

    ///

    /// * `weight_name` - Name of previously loaded weights

    /// * `x` - Input vector/matrix

    /// * `m` - Batch size (rows in input)

    /// * `k` - Input dimension

    /// * `n` - Output dimension

    ///

    /// # Errors

    ///

    /// Returns error if weights not found or compute fails.

    fn matmul_cached(

        &mut self,

        weight_name: &str,

        x: &[f32],

        m: u32,

        k: u32,

        n: u32,

    ) -> BackendResult<Vec<f32>>;

    // =========================================================================

    // Quantized Operations

    // =========================================================================

    /// Q4_K quantized GEMV using cached weights

    ///

    /// Computes y = dequant(W_q4k) @ x where W is Q4_K quantized.

    ///

    /// # Arguments

    ///

    /// * `weight_name` - Name of previously loaded Q4_K weights

    /// * `input` - Input vector [k]

    /// * `n` - Output dimension (rows in weight matrix)

    /// * `k` - Input dimension (cols in weight matrix)

    ///

    /// # Errors

    ///

    /// Returns error if weights not found, wrong quantization type, or compute fails.

    fn q4k_gemv_cached(

        &mut self,

        weight_name: &str,

        input: &[f32],

        n: u32,

        k: u32,

    ) -> BackendResult<Vec<f32>>;

    // =========================================================================

    // Synchronization

    // =========================================================================

    /// Synchronize device execution

    ///

    /// Blocks until all queued operations complete.

    ///

    /// # Errors

    ///

    /// Returns error if synchronization fails.

    fn synchronize(&self) -> BackendResult<()>;

}

#[cfg(test)]

mod tests {

    use super::*;

    use std::collections::HashMap;

    /// Mock backend for testing trait compliance

    struct MockBackend {

        device_name: String,

        weights: HashMap<String, Vec<f32>>,

        quantized_weights: HashMap<String, (Vec<u8>, u32)>,

    }

    impl ComputeBackend for MockBackend {

        fn is_available() -> bool {

            true

        }

        fn new(device_id: u32) -> BackendResult<Self> {

            Ok(Self {

                device_name: format!("MockDevice:{}", device_id),

                weights: HashMap::new(),

                quantized_weights: HashMap::new(),

            })

        }

        fn device_name(&self) -> String {

            self.device_name.clone()

        }

        fn load_weights(&mut self, name: &str, data: &[f32]) -> BackendResult<usize> {

            let len = data.len();

            self.weights.insert(name.to_string(), data.to_vec());

            Ok(len)

        }

        fn load_quantized_weights(

            &mut self,

            name: &str,

            data: &[u8],

            qtype: u32,

        ) -> BackendResult<usize> {

            let len = data.len();

            self.quantized_weights

                .insert(name.to_string(), (data.to_vec(), qtype));

            Ok(len)

        }

        fn has_weights(&self, name: &str) -> bool {

            self.weights.contains_key(name) || 
self.quantized_weights3
.
contains_key3
(
name3
)

        }

        fn clear_weights(&mut self) {

            self.weights.clear();

            self.quantized_weights.clear();

        }

        fn cached_weight_count(&self) -> usize {

            self.weights.len() + self.quantized_weights.len()

        }

        fn matmul(

            &mut self,

            a: &[f32],

            b: &[f32],

            m: u32,

            k: u32,

            n: u32,

        ) -> BackendResult<Vec<f32>> {

            let m = m as usize;

            let k = k as usize;

            let n = n as usize;

            if a.len() != m * k {

                return Err(format!("A size {} != m*k {}", a.len(), m * k).into());

            }

            if b.len() != k * n {

                return Err(format!("B size {} != k*n {}", b.len(), k * n).into());

            }

            let mut c = vec![0.0f32; m * n];

            for i in 0..
m2
 {

                for j in 0..
n4
 {

                    let mut sum = 0.0f32;

                    for p in 0..
k8
 {

                        sum += a[i * k + p] * b[p * n + j];

                    }

                    c[i * n + j] = sum;

                }

            }

            Ok(c)

        }

        fn matmul_cached(

            &mut self,

            weight_name: &str,

            x: &[f32],

            m: u32,

            k: u32,

            n: u32,

        ) -> BackendResult<Vec<f32>> {

            let weights = self

                .weights

                .get(weight_name)

                .ok_or_else(|| 
format!0
(
"Weights '{}' not found"0
, weight_name))
?0

                .clone();

            self.matmul(x, &weights, m, k, n)

        }

        fn q4k_gemv_cached(

            &mut self,

            weight_name: &str,

            input: &[f32],

            n: u32,

            k: u32,

        ) -> BackendResult<Vec<f32>> {

            // Mock: just return zeros of correct size

            if !self.quantized_weights.contains_key(weight_name) {

                return Err(format!("Quantized weights '{}' not found", weight_name).into());

            }

            let _ = input;

            let _ = k;

            Ok(vec![0.0f32; n as usize])

        }

        fn synchronize(&self) -> BackendResult<()> {

            Ok(())

        }

    }

    #[test]

    fn test_mock_backend_creation() {

        assert!(MockBackend::is_available());

        let backend = MockBackend::new(0).unwrap();

        assert_eq!(backend.device_name(), "MockDevice:0");

    }

    #[test]

    fn test_mock_backend_weight_management() {

        let mut backend = MockBackend::new(0).unwrap();

        // Initially empty

        assert_eq!(backend.cached_weight_count(), 0);

        assert!(!backend.has_weights("test"));

        // Load weights

        let data = vec![1.0, 2.0, 3.0, 4.0];

        let handle = backend.load_weights("test", &data).unwrap();

        assert_eq!(handle, 4);

        assert!(backend.has_weights("test"));

        assert_eq!(backend.cached_weight_count(), 1);

        // Load quantized weights

        let qdata = vec![0u8; 18]; // Q4_0 block

        let qhandle = backend.load_quantized_weights("test_q4", &qdata, 2).unwrap();

        assert_eq!(qhandle, 18);

        assert!(backend.has_weights("test_q4"));

        assert_eq!(backend.cached_weight_count(), 2);

        // Clear

        backend.clear_weights();

        assert_eq!(backend.cached_weight_count(), 0);

        assert!(!backend.has_weights("test"));

    }

    #[test]

    fn test_mock_backend_matmul() {

        let mut backend = MockBackend::new(0).unwrap();

        // 2x3 @ 3x2 = 2x2

        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];

        let b = vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0];

        let c = backend.matmul(&a, &b, 2, 3, 2).unwrap();

        // Expected: [[1*7+2*9+3*11, 1*8+2*10+3*12], [4*7+5*9+6*11, 4*8+5*10+6*12]]

        //         = [[7+18+33, 8+20+36], [28+45+66, 32+50+72]]

        //         = [[58, 64], [139, 154]]

        assert_eq!(c.len(), 4);

        assert!((c[0] - 58.0).abs() < 1e-5);

        assert!((c[1] - 64.0).abs() < 1e-5);

        assert!((c[2] - 139.0).abs() < 1e-5);

        assert!((c[3] - 154.0).abs() < 1e-5);

    }

    #[test]

    fn test_mock_backend_matmul_cached() {

        let mut backend = MockBackend::new(0).unwrap();

        // Load weights: 3x2 matrix

        let weights = vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0];

        backend.load_weights("W", &weights).unwrap();

        // Input: 2x3

        let x = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];

        let c = backend.matmul_cached("W", &x, 2, 3, 2).unwrap();

        assert_eq!(c.len(), 4);

    }

    #[test]

    fn test_mock_backend_sync() {

        let backend = MockBackend::new(0).unwrap();

        assert!(backend.synchronize().is_ok());

    }

}