//! WGPU compute backend — wraps CubeCL's WGPU runtime.
//!
//! This module is only compiled when the `wgpu` feature is enabled.
//! It uses the same CubeCL kernels as the CUDA backend, but dispatches
//! through WGPU to support ROCm, Metal, Vulkan, and DX12.

use cubecl::prelude::*;
use cubecl::wgpu::{WgpuDevice, WgpuRuntime};

use crate::compute::ComputeBackend;
use crate::kernels::{matmul, rmsnorm, rope_gpu, silu, softmax};
use crate::tensor::Tensor;

/// WGPU backend using CubeCL for cross-platform GPU kernel dispatch.
pub struct WgpuBackend {
    /// CubeCL compute client bound to the WGPU adapter.
    client: ComputeClient<<WgpuRuntime as Runtime>::Server, <WgpuRuntime as Runtime>::Channel>,
    /// Human-readable adapter name (e.g. "Metal", "Vulkan").
    adapter: String,
}

impl WgpuBackend {
    /// Attempt to initialize a WGPU adapter.
    ///
    /// Returns `None` if no suitable GPU adapter is found.
    pub fn try_new() -> Option<Self> {
        let device = WgpuDevice::default();
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            WgpuRuntime::client(&device)
        }));
        match result {
            Ok(client) => {
                tracing::info!("WGPU backend initialized");
                Some(Self {
                    client,
                    adapter: "wgpu".to_string(),
                })
            }
            Err(_) => {
                tracing::debug!("WGPU adapter not available");
                None
            }
        }
    }

    /// Human-readable adapter name (e.g. "Metal", "Vulkan", "DX12").
    pub fn adapter_name(&self) -> &str {
        &self.adapter
    }
}

impl ComputeBackend for WgpuBackend {
    fn name(&self) -> &str {
        "wgpu"
    }

    fn matmul(&self, a: &Tensor, b: &Tensor) -> Tensor {
        assert_eq!(a.shape.len(), 2, "matmul requires 2D tensors");
        assert_eq!(b.shape.len(), 2, "matmul requires 2D tensors");

        let m = a.shape[0] as u32;
        let k = a.shape[1] as u32;
        let k2 = b.shape[0] as u32;
        let n = b.shape[1] as u32;
        assert_eq!(k, k2, "inner dimensions must match: {} vs {}", k, k2);

        let a_handle = self.client.create(f32::as_bytes(&a.data));
        let b_handle = self.client.create(f32::as_bytes(&b.data));
        let out_handle = self.client.empty((m as usize) * (n as usize) * std::mem::size_of::<f32>());

        let a_strides = [k as usize, 1];
        let a_shape = [m as usize, k as usize];
        let b_strides = [n as usize, 1];
        let b_shape = [k as usize, n as usize];
        let out_strides = [n as usize, 1];
        let out_shape = [m as usize, n as usize];

        let (cube_count, cube_dim) = matmul::matmul_launch_config(m, n);

        unsafe {
            matmul::matmul_kernel::launch::<f32, WgpuRuntime>(
                &self.client,
                cube_count,
                cube_dim,
                TensorArg::from_raw_parts::<f32>(&a_handle, &a_strides, &a_shape, 1),
                TensorArg::from_raw_parts::<f32>(&b_handle, &b_strides, &b_shape, 1),
                TensorArg::from_raw_parts::<f32>(&out_handle, &out_strides, &out_shape, 1),
                m,
                n,
                k,
            );
        }

        let result_bytes = self.client.read_one(out_handle.binding());
        let result_data = f32::from_bytes(&result_bytes).to_vec();

        Tensor {
            data: result_data,
            shape: vec![m as usize, n as usize],
        }
    }

    fn softmax(&self, x: &mut Tensor, _dim: usize) {
        let last_dim = *x.shape.last().expect("softmax on empty tensor");
        let num_rows = (x.data.len() / last_dim) as u32;
        let dim_size = last_dim as u32;

        let input_handle = self.client.create(f32::as_bytes(&x.data));
        let out_handle = self.client.empty(x.data.len() * std::mem::size_of::<f32>());

        let strides = [last_dim, 1];
        let shape = [num_rows as usize, last_dim];

        let (cube_count, cube_dim) = softmax::softmax_launch_config(num_rows);

        unsafe {
            softmax::softmax_kernel::launch::<f32, WgpuRuntime>(
                &self.client,
                cube_count,
                cube_dim,
                TensorArg::from_raw_parts::<f32>(&input_handle, &strides, &shape, 1),
                TensorArg::from_raw_parts::<f32>(&out_handle, &strides, &shape, 1),
                dim_size,
            );
        }

        let result_bytes = self.client.read_one(out_handle.binding());
        x.data = f32::from_bytes(&result_bytes).to_vec();
    }

    fn rope(
        &self,
        q: &mut Tensor,
        k: &mut Tensor,
        freqs_cos: &[f32],
        freqs_sin: &[f32],
        num_heads: usize,
        num_kv_heads: usize,
        head_dim: usize,
        start_pos: usize,
    ) {
        let seq_len = q.shape[0];
        let half_dim = head_dim / 2;

        self.apply_rope_single(
            &mut q.data,
            freqs_cos,
            freqs_sin,
            seq_len,
            num_heads,
            head_dim,
            half_dim,
            start_pos,
        );

        self.apply_rope_single(
            &mut k.data,
            freqs_cos,
            freqs_sin,
            seq_len,
            num_kv_heads,
            head_dim,
            half_dim,
            start_pos,
        );
    }

    fn rmsnorm(&self, x: &Tensor, weight: &Tensor, eps: f32) -> Tensor {
        let hidden_size = weight.data.len();
        let num_rows = (x.data.len() / hidden_size) as u32;

        let x_handle = self.client.create(f32::as_bytes(&x.data));
        let w_handle = self.client.create(f32::as_bytes(&weight.data));
        let out_handle = self.client.empty(x.data.len() * std::mem::size_of::<f32>());

        let x_strides = [hidden_size, 1];
        let x_shape = [num_rows as usize, hidden_size];
        let w_strides = [1];
        let w_shape = [hidden_size];
        let out_strides = [hidden_size, 1];
        let out_shape = [num_rows as usize, hidden_size];

        let (cube_count, cube_dim) = rmsnorm::rmsnorm_launch_config(num_rows);

        unsafe {
            rmsnorm::rmsnorm_kernel::launch::<f32, WgpuRuntime>(
                &self.client,
                cube_count,
                cube_dim,
                TensorArg::from_raw_parts::<f32>(&x_handle, &x_strides, &x_shape, 1),
                TensorArg::from_raw_parts::<f32>(&w_handle, &w_strides, &w_shape, 1),
                TensorArg::from_raw_parts::<f32>(&out_handle, &out_strides, &out_shape, 1),
                hidden_size as u32,
            );
        }

        let result_bytes = self.client.read_one(out_handle.binding());
        let result_data = f32::from_bytes(&result_bytes).to_vec();

        Tensor {
            data: result_data,
            shape: x.shape.clone(),
        }
    }

    fn silu_elementwise_mul(&self, gate: &mut Tensor, up: &Tensor) {
        let len = gate.data.len() as u32;

        let gate_handle = self.client.create(f32::as_bytes(&gate.data));
        let up_handle = self.client.create(f32::as_bytes(&up.data));
        let out_handle = self.client.empty(gate.data.len() * std::mem::size_of::<f32>());

        let strides = [1];
        let shape = [gate.data.len()];

        let (cube_count, cube_dim) = silu::silu_mul_launch_config(len);

        unsafe {
            silu::silu_mul_kernel::launch::<f32, WgpuRuntime>(
                &self.client,
                cube_count,
                cube_dim,
                TensorArg::from_raw_parts::<f32>(&gate_handle, &strides, &shape, 1),
                TensorArg::from_raw_parts::<f32>(&up_handle, &strides, &shape, 1),
                TensorArg::from_raw_parts::<f32>(&out_handle, &strides, &shape, 1),
                len,
            );
        }

        let result_bytes = self.client.read_one(out_handle.binding());
        gate.data = f32::from_bytes(&result_bytes).to_vec();
    }
}

impl WgpuBackend {
    /// Helper: apply RoPE to a single tensor (Q or K) on GPU.
    fn apply_rope_single(
        &self,
        data: &mut Vec<f32>,
        freqs_cos: &[f32],
        freqs_sin: &[f32],
        seq_len: usize,
        num_heads: usize,
        head_dim: usize,
        half_dim: usize,
        start_pos: usize,
    ) {
        let total_len = data.len();

        let x_handle = self.client.create(f32::as_bytes(data));
        let cos_handle = self.client.create(f32::as_bytes(freqs_cos));
        let sin_handle = self.client.create(f32::as_bytes(freqs_sin));
        let out_handle = self.client.empty(total_len * std::mem::size_of::<f32>());

        let x_strides = [num_heads * head_dim, 1];
        let x_shape = [seq_len, num_heads * head_dim];
        let table_strides = [1];
        let table_shape = [freqs_cos.len()];
        let out_strides = [num_heads * head_dim, 1];
        let out_shape = [seq_len, num_heads * head_dim];

        let (cube_count, cube_dim) =
            rope_gpu::rope_launch_config(seq_len as u32, num_heads as u32, half_dim as u32);

        unsafe {
            rope_gpu::rope_kernel::launch::<f32, WgpuRuntime>(
                &self.client,
                cube_count,
                cube_dim,
                TensorArg::from_raw_parts::<f32>(&x_handle, &x_strides, &x_shape, 1),
                TensorArg::from_raw_parts::<f32>(&cos_handle, &table_strides, &table_shape, 1),
                TensorArg::from_raw_parts::<f32>(&sin_handle, &table_strides, &table_shape, 1),
                TensorArg::from_raw_parts::<f32>(&out_handle, &out_strides, &out_shape, 1),
                num_heads as u32,
                head_dim as u32,
                half_dim as u32,
                start_pos as u32,
            );
        }

        let result_bytes = self.client.read_one(out_handle.binding());
        *data = f32::from_bytes(&result_bytes).to_vec();
    }
}
