#![allow(clippy::too_many_arguments)]
use super::Float4;
use crate::Float;
use crate::core::geometry::{Normal3f, Point2f, Point2i, Point3f, Point3fi, Ray, Vector3f};
use crate::lights::LightSampleContext;
use crate::soa_struct;
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use cust::memory::{CopyDestination, DeviceMemory};
use cust::prelude::*;

#[macro_export]
macro_rules! soa_struct {
    (
    $(#[$outer:meta])*
    pub struct $name:ident {
    $(
        pub $field:ident : $type:ty
    ),* $(,)?
    }
) => {
    #[cfg(feature = "use_gpu")]
    $(#[$outer])*
    pub struct $name {
        capacity: u32,
        pub count: cust::memory::DeviceBuffer<u32>,
        $(
            pub $field: cust::memory::DeviceBuffer<$type>,
        )*
    }

    #[cfg(feature = "use_gpu")]
    impl $name {
        pub fn new(capacity: usize) -> cust::error::CudaResult<Self> {
            use cust::memory::DeviceBuffer;
            Ok(Self {
                capacity: capacity as u32,
                count: DeviceBuffer::zeroed(1)?,
                $(
                    $field: DeviceBuffer::zeroed(capacity)?,
                )*
            })
        }

        pub fn len(&self) -> cust::error::CudaResult<u32> {
            let mut host_count = [0u32; 1];
            self.count.copy_to(&mut host_count)?;
            Ok(host_count[0])
        }

        pub fn reset(&mut self) -> cust::error::CudaResult<()> {
            self.count.copy_from(&[0])
        }

        // Generate the View name
        pub fn as_view(&mut self) -> paste::paste! { [<$name View>] } {
            paste::paste! {
                [<$name View>] {
                    capacity: self.capacity,
                    count: self.count.as_device_ptr().as_mut_ptr(),
                    $(
                        $field: self.$field.as_device_ptr().as_raw() as *mut $type,
                    )*
                }
            }
        }
    }

    paste::paste! {
            #[repr(C)]
            #[derive(Clone, Copy)]
            pub struct [<$name View>] {
                pub capacity: u32,
                pub count: *mut u32,
                $(
                    pub $field: *mut $type,
                )*
            }

            unsafe impl cust::memory::DeviceCopy for [<$name View>] {}

            impl [<$name View>] {
                // The raw push that fills every field
                #[cfg(feature = "use_gpu")]
                pub unsafe fn push(&self, $( $field : $type ),* ) -> Option<u32> {
                    use core::sync::atomic::{AtomicU32, Ordering};

                    let index = unsafe {
                        let counter_ptr = self.count as *mut AtomicU32;
                        (*counter_ptr).fetch_add(1, Ordering::Relaxed)
                    };

                    if index >= self.capacity {
                        return None;
                    }

                    unsafe {
                        $(
                            *self.$field.add(index as usize) = $field;
                        )*
                    }

                    Some(index)
                }

                #[cfg(feature = "use_gpu")]
                pub unsafe fn size(&self) -> u32 {
                    use core::sync::atomic::{AtomicU32, Ordering};
                    unsafe {
                        (*(self.count as *const AtomicU32)).load(Ordering::Relaxed)
                    }
                }

                $(
                    #[cfg(feature = "use_gpu")]
                    pub fn [<$field _ptr>](&self) -> *mut $type {
                        self.$field
                    }
                )*
            }
        }
    };
}

#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct RaySamplesDirect {
    pub u: Point2f,
    pub uc: Float,
}

#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct RaySamplesIndirect {
    pub uc: Float,
    pub rr: Float,
    pub u: Point2f,
}

#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct RaySamplesSubsurface {
    pub uc: Float,
    pub u: Point2f,
}

#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct RaySamples {
    pub direct: RaySamplesDirect,
    pub indirect: RaySamplesIndirect,
    pub have_subsurface: bool,
    pub subsurface: RaySamplesSubsurface,
}

soa_struct! {
    pub struct RayQueue {
        pub ray_o: Point3f,
        pub ray_d: Vector3f,

        pub depth: i32,
        pub lambda: SampledWavelengths,
        pub pixel_index: u32,

        pub beta: SampledSpectrum,
        pub r_u: SampledSpectrum,
        pub r_l: SampledSpectrum,

        pub ctx_pi: Point3f,
        pub ctx_n: Normal3f,
        pub ctx_ns: Normal3f,

        pub eta_scale: Float,
        pub specular_bounce: u32,
        pub any_non_specular_bounces: u32,
    }
}

soa_struct! {
    pub struct PixelSampleStateStorage {
        pub p_pixel: Point2i,
        pub l: SampledSpectrum,
        pub lambda: SampledWavelengths,
        pub filter_weight: Float,
        pub visible_surface: u32,
        pub camera_ray_weight: SampledSpectrum,

        pub rs_direct_packed: Float4,
        pub rs_indirect_packed: Float4,
        pub rs_subsurface_packed: Float4,
    }
}

soa_struct! {
    pub struct EscapedRayQueue {
        pub ray_o: Point3f,
        pub ray_d: Vector3f,
        pub depth: i32,
        pub lambda: SampledWavelengths,
        pub pixel_index: u32,
        pub beta: SampledSpectrum,
        pub specular_bounce: u32,
        pub r_u: SampledSpectrum,
        pub r_l: SampledSpectrum,
        pub ctx_pi: Point3f,
        pub ctx_n: Normal3f,
        pub ctx_ns: Normal3f,
    }
}

soa_struct! {
    pub struct HitAreaLightQueue {
        pub area_light_id: u32, // Light ID
        pub p: Point3f,
        pub n: Normal3f,
        pub uv: Point2f,
        pub wo: Vector3f,
        pub lambda: SampledWavelengths,
        pub depth: i32,
        pub beta: SampledSpectrum,
        pub r_u: SampledSpectrum,
        pub r_l: SampledSpectrum,
        pub ctx_pi: Point3f,
        pub ctx_n: Normal3f,
        pub ctx_ns: Normal3f,
        pub specular_bounce: u32,
        pub pixel_index: u32,
    }
}

soa_struct! {
    pub struct ShadowRayQueue {
        pub ray_o: Point3f,
        pub ray_d: Vector3f,
        pub t_max: Float,
        pub lambda: SampledWavelengths,
        pub ld: SampledSpectrum,
        pub r_u: SampledSpectrum,
        pub r_l: SampledSpectrum,
        pub pixel_index: u32,
    }
}

soa_struct! {
    pub struct GetBSSRDFAndProbeRayQueue {
        pub material_id: u32,
        pub lambda: SampledWavelengths,
        pub beta: SampledSpectrum,
        pub r_u: SampledSpectrum,
        pub p: Point3f,
        pub wo: Vector3f,
        pub n: Normal3f,
        pub ns: Normal3f,
        pub dpdus: Vector3f,
        pub uv: Point2f,
        pub depth: i32,
        pub mi_inside: u32,
        pub mi_outside: u32,
        pub eta_scale: Float,
        pub pixel_index: u32,
    }
}

soa_struct! {
    pub struct SubsurfaceScatterQueue {
        pub p0: Point3f,
        pub p1: Point3f,
        pub depth: i32,
        pub material_id: u32,
        pub lambda: SampledWavelengths,
        pub beta: SampledSpectrum,
        pub r_u: SampledSpectrum,
        pub mi_inside: u32,
        pub mi_outside: u32,
        pub eta_scale: Float,
        pub pixel_index: u32,
    }
}

soa_struct! {
    pub struct MediumSampleQueue {
        pub ray_o: Point3f,
        pub ray_d: Vector3f,
        pub t_max: Float,
        pub lambda: SampledWavelengths,
        pub beta: SampledSpectrum,
        pub r_u: SampledSpectrum,
        pub r_l: SampledSpectrum,
        pub pixel_index: u32,

        pub ctx_pi: Point3f,
        pub ctx_n: Normal3f,
        pub ctx_ns: Normal3f,

        pub specular_bounce: u32,
        pub any_non_specular_bounces: u32,
        pub eta_scale: Float,

        pub area_light_id: u32,
        pub pi: Point3fi,
        pub n: Normal3f,
        pub dpdu: Vector3f,
        pub dpdv: Vector3f,
        pub wo: Vector3f,
        pub uv: Point2f,
        pub material_id: u32,
        pub ns: Normal3f,
        pub dpdus: Vector3f,
        pub dpdvs: Vector3f,
        pub dndus: Normal3f,
        pub dndvs: Normal3f,
        pub face_index: i32,
        pub mi_inside: u32,
        pub mi_outside: u32,
    }
}

soa_struct! {
    pub struct MaterialEvalQueue {
        pub material_id: u32,
        pub pi: Point3fi,
        pub n: Normal3f,
        pub dpdu: Vector3f,
        pub dpdv: Vector3f,
        pub time: Float,
        pub depth: i32,
        pub ns: Normal3f,
        pub dpdus: Vector3f,
        pub dpdvs: Vector3f,
        pub dndus: Normal3f,
        pub dndvs: Normal3f,
        pub uv: Point2f,
        pub face_index: i32,
        pub lambda: SampledWavelengths,
        pub pixel_index: u32,
        pub any_non_specular_bounces: u32,
        pub wo: Vector3f,
        pub beta: SampledSpectrum,
        pub r_u: SampledSpectrum,
        pub eta_scale: Float,
        pub mi_inside: u32,
        pub mi_outside: u32,
    }
}

soa_struct! {
    pub struct MediumScatterQueue {
        pub p: Point3f,
        pub depth: usize,
        pub lambda: SampledWavelengths,
        pub beta: SampledSpectrum,
        pub r_u: SampledSpectrum,
        pub wo: Vector3f,
        pub time: Float,
        pub eta_scale: Float,
        pub pixel_index: usize,

        // ID
        pub phase_function: u32,
        pub medium: u32,
    }
}

#[repr(C)]
#[derive(Clone, Copy)]
pub struct RayWorkItem {
    pub ray: Ray,
    pub depth: i32,
    pub lambda: SampledWavelengths,
    pub pixel_index: u32,
    pub beta: SampledSpectrum,
    pub r_u: SampledSpectrum,
    pub r_l: SampledSpectrum,
    pub prev_intr_ctx: LightSampleContext,
    pub eta_scale: Float,
    pub specular_bounce: bool,
    pub any_non_specular_bounces: bool,
}

#[repr(C)]
#[derive(Clone, Copy)]
pub struct EscapedRayWorkItem {
    pub ray_o: Point3f,
    pub ray_d: Vector3f,
    pub depth: i32,
    pub lambda: SampledWavelengths,
    pub pixel_index: u32,
    pub beta: SampledSpectrum,
    pub specular_bounce: bool,
    pub r_u: SampledSpectrum,
    pub r_l: SampledSpectrum,
    pub prev_intr_ctx: LightSampleContext,
}

#[repr(C)]
#[derive(Clone, Copy)]
pub struct ShadowRayWorkItem {
    pub ray: Ray,
    pub t_max: Float,
    pub lambda: SampledWavelengths,
    pub ld: SampledSpectrum,
    pub r_u: SampledSpectrum,
    pub r_l: SampledSpectrum,
    pub pixel_index: u32,
}

impl RayQueueView {
    #[cfg(feature = "use_gpu")]
    pub unsafe fn push_work_item(&self, item: RayWorkItem) -> Option<u32> {
        unsafe {
            self.push(
                item.ray.o,
                item.ray.d,
                item.depth,
                item.lambda,
                item.pixel_index,
                item.beta,
                item.r_u,
                item.r_l,
                item.prev_intr_ctx.pi.into(),
                item.prev_intr_ctx.n,
                item.prev_intr_ctx.ns,
                item.eta_scale,
                if item.specular_bounce { 1 } else { 0 },
                if item.any_non_specular_bounces { 1 } else { 0 },
            )
        }
    }
}

impl EscapedRayQueueView {
    #[cfg(feature = "use_gpu")]
    pub unsafe fn push_work_item(&self, r: &RayWorkItem) -> Option<u32> {
        unsafe {
            self.push(
                r.ray.o,
                r.ray.d,
                r.depth,
                r.lambda,
                r.pixel_index,
                r.beta,
                if r.specular_bounce { 1 } else { 0 },
                r.r_u,
                r.r_l,
                r.prev_intr_ctx.pi.into(),
                r.prev_intr_ctx.n,
                r.prev_intr_ctx.ns,
            )
        }
    }
}

impl PixelSampleStateStorageView {
    #[cfg(feature = "use_gpu")]
    pub unsafe fn get_samples(&self, index: u32) -> RaySamples {
        let i = index as usize;

        let (dir, ind, ss) = unsafe {
            (
                *self.rs_direct_packed.add(i),
                *self.rs_indirect_packed.add(i),
                *self.rs_subsurface_packed.add(i),
            )
        };

        let direct_u = Point2f::new(dir.v[0], dir.v[1]);
        let direct_uc = dir.v[2];
        let flags = dir.v[3] as i32;
        let have_subsurface = (flags & 1) != 0;

        let indirect_uc = ind.v[0];
        let indirect_rr = ind.v[1];
        let indirect_u = Point2f::new(ind.v[2], ind.v[3]);

        let subsurface_uc = ss.v[0];
        let subsurface_u = Point2f::new(ss.v[1], ss.v[2]);

        RaySamples {
            direct: RaySamplesDirect {
                u: direct_u,
                uc: direct_uc,
            },
            indirect: RaySamplesIndirect {
                uc: indirect_uc,
                rr: indirect_rr,
                u: indirect_u,
            },
            have_subsurface,
            subsurface: RaySamplesSubsurface {
                uc: subsurface_uc,
                u: subsurface_u,
            },
        }
    }

    #[cfg(feature = "use_gpu")]
    pub unsafe fn set_samples(&self, index: u32, rs: RaySamples) {
        if index >= self.capacity {
            return;
        }
        let i = index as usize;

        let flags = if rs.have_subsurface { 1.0 } else { 0.0 };
        let dir = Float4 {
            v: [rs.direct.u.0[0], rs.direct.u.0[1], rs.direct.uc, flags],
        };

        let ind = Float4 {
            v: [
                rs.indirect.uc,
                rs.indirect.rr,
                rs.indirect.u.0[0],
                rs.indirect.u.0[1],
            ],
        };

        unsafe {
            *self.rs_direct_packed.add(i) = dir;
            *self.rs_indirect_packed.add(i) = ind;
        }

        if rs.have_subsurface {
            let ss = Float4 {
                v: [
                    rs.subsurface.uc,
                    rs.subsurface.u.0[0],
                    rs.subsurface.u.0[1],
                    0.0,
                ],
            };
            unsafe {
                *self.rs_subsurface_packed.add(i) = ss;
            }
        }
    }
}