use bytemuck::{Pod, Zeroable};
use cust::context::{CacheConfig, CurrentContext, ResourceLimit};
use cust::device::DeviceAttribute;
use cust::memory::DeviceCopy;
use cust::prelude::*;
use lazy_static::lazy_static;
use parking_lot::Mutex;
use std::error::Error;
use std::sync::Arc;

use shared::Float;
use shared::core::geometry::{Normal, Point, Vector};
use shared::core::options::get_options;
use shared::spectra::{SampledSpectrum, SampledWavelengths};
use shared::utils::interval::Interval;

#[macro_export]
macro_rules! impl_gpu_traits {
    ($name:ty) => {
        unsafe impl DeviceCopy for $name {}
        unsafe impl Zeroable for $name {}
        unsafe impl Pod for $name {}
    };
}

#[macro_export]
macro_rules! impl_math_gpu_traits {
    ($Struct:ident) => {
        #[cfg(feature = "use_gpu")]
        unsafe impl<T, const N: usize> DeviceCopy for $Struct<T, N> where T: DeviceCopy + Copy {}

        unsafe impl<T, const N: usize> Zeroable for $Struct<T, N> where T: Zeroable {}

        unsafe impl<T, const N: usize> Pod for $Struct<T, N> where T: Pod {}
    };
}

impl_math_gpu_traits!(Vector);
impl_math_gpu_traits!(Normal);
impl_math_gpu_traits!(Point);
impl_gpu_traits!(Interval);
impl_gpu_traits!(Float4);
impl_gpu_traits!(SampledSpectrum);
impl_gpu_traits!(SampledWavelengths);

#[repr(C, align(16))]
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct Float4 {
    pub v: [f32; 4],
}
pub type Vec4 = Vector<Float, 4>;

impl From<Vec4> for Float4 {
    #[inline]
    fn from(vec: Vector<f32, 4>) -> Self {
        Self { v: vec.0 }
    }
}

impl From<Float4> for Vec4 {
    #[inline]
    fn from(storage: Float4) -> Self {
        Vector(storage.v)
    }
}

struct KernelStats {
    description: String,
    num_launches: usize,
    sum_ms: f32,
    min_ms: f32,
    max_ms: f32,
}

impl KernelStats {
    fn new(description: &str) -> Self {
        Self {
            description: description.to_string(),
            num_launches: 0,
            sum_ms: 0.0,
            min_ms: 0.0,
            max_ms: 0.0,
        }
    }
}

struct ProfilerEvent {
    start: Event,
    stop: Event,
    active: bool,
    stats: Option<Arc<Mutex<KernelStats>>>,
}

impl ProfilerEvent {
    fn new() -> Result<Self, cust::error::CudaError> {
        let start = Event::new(EventFlags::DEFAULT)?;
        let stop = Event::new(EventFlags::DEFAULT)?;
        Ok(Self {
            start,
            stop,
            active: false,
            stats: None,
        })
    }

    fn sync(&mut self) {
        if !self.active {
            return;
        }

        if self.stop.synchronize().is_ok() {
            match self.stop.elapsed_time_f32(&self.start) {
                Ok(ms) => {
                    if let Some(stats_arc) = &self.stats {
                        let mut stats = stats_arc.lock();
                        stats.num_launches += 1;
                        if stats.num_launches == 1 {
                            stats.sum_ms = ms;
                            stats.min_ms = ms;
                            stats.max_ms = ms;
                        } else {
                            stats.sum_ms += ms;
                            stats.min_ms = stats.min_ms.min(ms);
                            stats.max_ms = stats.max_ms.max(ms);
                        }
                    }
                }
                Err(e) => log::error!("Failed to get elapsed time: {:?}", e),
            }
        }
        self.active = false;
    }
}

struct Profiler {
    kernel_stats: Vec<Arc<Mutex<KernelStats>>>,
    event_pool: Vec<ProfilerEvent>,
    pool_offset: usize,
}

impl Profiler {
    fn new() -> Self {
        Self {
            kernel_stats: Vec::new(),
            event_pool: Vec::new(),
            pool_offset: 0,
        }
    }

    fn prepare<'a>(&'a mut self, description: &str) -> &'a mut ProfilerEvent {
        if self.event_pool.is_empty() {
            for _ in 0..128 {
                if let Ok(e) = ProfilerEvent::new() {
                    self.event_pool.push(e);
                }
            }
        }

        if self.pool_offset >= self.event_pool.len() {
            self.pool_offset = 0;
        }

        let idx = self.pool_offset;
        self.pool_offset += 1;

        let pe = &mut self.event_pool[idx];

        if pe.active {
            pe.sync();
        }

        pe.active = true;
        pe.stats = None;

        let mut found = None;
        for s in &self.kernel_stats {
            if s.lock().description == description {
                found = Some(s.clone());
                break;
            }
        }

        if found.is_none() {
            let new_stats = Arc::new(Mutex::new(KernelStats::new(description)));
            self.kernel_stats.push(new_stats.clone());
            found = Some(new_stats);
        }

        pe.stats = found;
        pe
    }
}

pub struct GpuState {
    pub context: Context,
    pub stream: Stream,
    profiler: Profiler,
}

impl GpuState {
    fn init(device_index: u32) -> Result<Self, Box<dyn Error>> {
        match cust::init(CudaFlags::empty()) {
            Ok(_) => {}
            Err(cust::error::CudaError::SharedObjectInitFailed) => {}
            Err(e) => return Err(Box::new(e)),
        }

        let device = Device::get_device(device_index)?;
        let name = device.name().unwrap_or_else(|_| "Unknown".into());

        log::info!("Selected GPU: {}", name);

        let has_unified = device
            .get_attribute(DeviceAttribute::UnifiedAddressing)
            .unwrap_or(0);
        if has_unified == 0 {
            panic!("Selected GPU does not support unified addressing.");
        }

        let context = Context::new(device)?;

        CurrentContext::set_resource_limit(ResourceLimit::StackSize, 8192)?;
        CurrentContext::set_resource_limit(ResourceLimit::PrintfFifoSize, 32 * 1024 * 1024)?;
        CurrentContext::set_cache_config(CacheConfig::PreferL1)?;

        let stream = Stream::new(StreamFlags::DEFAULT, None)?;

        Ok(Self {
            context,
            stream,
            profiler: Profiler::new(),
        })
    }
}

lazy_static! {
    pub static ref GPU_STATE: Mutex<Option<GpuState>> = Mutex::new(None);
}

pub fn gpu_init() {
    if !get_options().use_gpu {
        return;
    }

    let device_id = get_options().gpu_device.unwrap_or(0);
    log::info!("Initializing GPU Device {}", device_id);

    match GpuState::init(device_id) {
        Ok(state) => {
            #[cfg(feature = "use_nvtx")]
            nvtx::name_thread("MAIN_THREAD");
            *GPU_STATE.lock() = Some(state);
        }
        Err(e) => {
            panic!("Failed to initialize GPU: {:?}", e);
        }
    }
}

pub fn gpu_thread_init() {
    if let Some(state) = GPU_STATE.lock().as_ref() {
        if let Err(e) = CurrentContext::set_current(&state.context) {
            log::error!("Failed to set CUDA context for thread: {:?}", e);
        }
    }
}

const SANITY_PTX: &str = r#"
.version 6.5
.target sm_30
.address_size 64
.visible .entry scale_array_kernel(.param .u64 ptr, .param .u32 n, .param .f32 scale) {
    .reg .u64 %ptr; .reg .u32 %n; .reg .f32 %scale;
    .reg .u32 %tid; .reg .u32 %ntid; .reg .u32 %ctid; .reg .u32 %idx; 
    .reg .pred %p; .reg .u64 %addr; .reg .f32 %val;

    ld.param.u64 %ptr, [ptr];
    ld.param.u32 %n, [n];
    ld.param.f32 %scale, [scale];

    mov.u32 %tid, %tid.x;
    mov.u32 %ntid, %ntid.x;
    mov.u32 %ctid, %ctaid.x;
    mad.lo.s32 %idx, %ctid, %ntid, %tid;

    setp.ge.u32 %p, %idx, %n;
    @%p ret;

    mul.wide.u32 %addr, %idx, 4;
    add.u64 %addr, %ptr, %addr;
    ld.global.f32 %val, [%addr];
    mul.f32 %val, %val, %scale;
    st.global.f32 [%addr], %val;
}
"#;

pub fn launch_scale_kernel(ptr: *mut f32, len: usize, scale: f32) -> Result<(), Box<dyn Error>> {
    // Note: quick_init works best for isolated tests.
    // If the main engine runs, this might conflict, but for the sanity test it is correct.
    let _ctx = cust::quick_init()?;

    let module = Module::from_str(SANITY_PTX)?;
    let stream = Stream::new(StreamFlags::NON_BLOCKING, None)?;
    let mut kernel = module.get_function("scale_array_kernel")?;

    let block_size = 256;
    let grid_size = (len as u32 + block_size - 1) / block_size;

    unsafe {
        launch!(
            kernel<<<grid_size, block_size, 0, stream>>>(
                ptr as u64,
                len as u32,
                scale
            )
        )?;
    }
    stream.synchronize()?;
    Ok(())
}
