//! Stencil kernel support for CUDA.
//!
//! This module provides functionality to load and execute stencil kernels
//! that were generated by the `#[stencil_kernel]` macro.

use std::sync::Arc;

use ringkernel_core::__private::{find_stencil_kernel, StencilKernelRegistration};
use ringkernel_core::error::{Result, RingKernelError};

use crate::device::CudaDevice;

/// A compiled stencil kernel ready for execution.
pub struct CompiledStencilKernel {
    /// Kernel ID.
    pub id: String,
    /// Grid dimensionality.
    pub grid: String,
    /// Tile dimensions.
    pub tile_size: (u32, u32),
    /// Halo width.
    pub halo: u32,
    /// The CUDA device.
    device: Arc<cudarc::driver::CudaDevice>,
    /// Module name for function lookup.
    module_name: String,
}

impl CompiledStencilKernel {
    /// Compile a stencil kernel from its registration.
    pub fn from_registration(reg: &StencilKernelRegistration, device: &CudaDevice) -> Result<Self> {
        let cuda_device = device.inner();

        // Compile CUDA source to PTX using NVRTC
        let ptx = cudarc::nvrtc::compile_ptx(reg.cuda_source).map_err(|e| {
            RingKernelError::CompilationError(format!(
                "NVRTC compilation failed for '{}': {}",
                reg.id, e
            ))
        })?;

        // Load the PTX module
        let module_name = format!("stencil_{}", reg.id);
        cuda_device
            .load_ptx(ptx, &module_name, &[reg.id])
            .map_err(|e| {
                RingKernelError::CompilationError(format!(
                    "Failed to load PTX for '{}': {}",
                    reg.id, e
                ))
            })?;

        tracing::info!(
            kernel_id = %reg.id,
            grid = %reg.grid,
            tile = ?(&reg.tile_width, &reg.tile_height),
            halo = %reg.halo,
            "Compiled stencil kernel"
        );

        Ok(Self {
            id: reg.id.to_string(),
            grid: reg.grid.to_string(),
            tile_size: (reg.tile_width, reg.tile_height),
            halo: reg.halo,
            device: cuda_device.clone(),
            module_name,
        })
    }

    /// Get the CUDA function for this kernel.
    pub fn get_function(&self) -> Result<cudarc::driver::CudaFunction> {
        self.device
            .get_func(&self.module_name, &self.id)
            .ok_or_else(|| {
                RingKernelError::KernelNotFound(format!(
                    "Function '{}' not found in module '{}'",
                    self.id, self.module_name
                ))
            })
    }

    /// Get the buffer width (tile + 2*halo).
    pub fn buffer_width(&self) -> u32 {
        self.tile_size.0 + 2 * self.halo
    }

    /// Get the buffer height (tile + 2*halo).
    pub fn buffer_height(&self) -> u32 {
        self.tile_size.1 + 2 * self.halo
    }

    /// Calculate launch configuration for a grid of tiles.
    pub fn launch_config(&self, num_tiles: u32) -> LaunchConfig {
        LaunchConfig {
            grid_dim: (num_tiles, 1, 1),
            block_dim: (self.tile_size.0, self.tile_size.1, 1),
            shared_mem_bytes: 0,
        }
    }
}

/// Launch configuration for a stencil kernel.
#[derive(Debug, Clone, Copy)]
pub struct LaunchConfig {
    /// Grid dimensions (blocks).
    pub grid_dim: (u32, u32, u32),
    /// Block dimensions (threads per block).
    pub block_dim: (u32, u32, u32),
    /// Shared memory in bytes.
    pub shared_mem_bytes: u32,
}

/// Loader for stencil kernels registered via `#[stencil_kernel]`.
pub struct StencilKernelLoader {
    device: CudaDevice,
}

impl StencilKernelLoader {
    /// Create a new loader for the given device.
    pub fn new(device: CudaDevice) -> Self {
        Self { device }
    }

    /// Load a stencil kernel by ID.
    ///
    /// Looks up the kernel in the static registry (populated by `#[stencil_kernel]` macro),
    /// compiles it to PTX, and loads it onto the GPU.
    pub fn load(&self, kernel_id: &str) -> Result<CompiledStencilKernel> {
        let registration = find_stencil_kernel(kernel_id).ok_or_else(|| {
            RingKernelError::KernelNotFound(format!(
                "Stencil kernel '{}' not registered. Did you use #[stencil_kernel]?",
                kernel_id
            ))
        })?;

        CompiledStencilKernel::from_registration(registration, &self.device)
    }

    /// Load a stencil kernel from raw CUDA source.
    ///
    /// This is useful for testing or dynamically generated kernels.
    pub fn load_from_source(
        &self,
        kernel_id: &str,
        cuda_source: &str,
        tile_size: (u32, u32),
        halo: u32,
    ) -> Result<CompiledStencilKernel> {
        let cuda_device = self.device.inner();

        // Compile CUDA source to PTX
        let ptx = cudarc::nvrtc::compile_ptx(cuda_source).map_err(|e| {
            RingKernelError::CompilationError(format!("NVRTC compilation failed: {}", e))
        })?;

        // Load the PTX module - cudarc requires &[&str] with 'static lifetime for function names
        // We use the module name as both the module and function name
        let module_name = format!("stencil_{}", kernel_id);
        let kernel_id_owned = kernel_id.to_string();

        // Use Box::leak to create 'static strings (acceptable for kernel registration)
        let module_name_static: &'static str = Box::leak(module_name.clone().into_boxed_str());
        let kernel_id_static: &'static str = Box::leak(kernel_id_owned.into_boxed_str());

        cuda_device
            .load_ptx(ptx, module_name_static, &[kernel_id_static])
            .map_err(|e| RingKernelError::CompilationError(format!("Failed to load PTX: {}", e)))?;

        Ok(CompiledStencilKernel {
            id: kernel_id.to_string(),
            grid: "2d".to_string(),
            tile_size,
            halo,
            device: cuda_device.clone(),
            module_name,
        })
    }

    /// List all registered stencil kernels.
    pub fn list_registered() -> Vec<&'static str> {
        ringkernel_core::__private::registered_stencil_kernels()
            .map(|r| r.id)
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_launch_config() {
        let config = LaunchConfig {
            grid_dim: (16, 1, 1),
            block_dim: (16, 16, 1),
            shared_mem_bytes: 0,
        };

        assert_eq!(config.grid_dim.0, 16);
        assert_eq!(config.block_dim.0, 16);
        assert_eq!(config.block_dim.1, 16);
    }
}
