use crate::core::image::Image;
use crate::core::texture::{FloatTexture, SpectrumTexture};
use crate::shapes::{BilinearPatchMesh, TriangleMesh};
use crate::spectra::DenselySampledSpectrumBuffer;
use crate::utils::mipmap::MIPMap;
use crate::utils::sampling::{PiecewiseConstant2D, WindowedPiecewiseConstant2D};
use shared::core::color::RGBToSpectrumTable;
use shared::core::image::DeviceImage;
use shared::core::shape::Shape;
use shared::core::spectrum::Spectrum;
use shared::core::texture::{GPUFloatTexture, GPUSpectrumTexture};
use shared::spectra::{DenselySampledSpectrum, DeviceStandardColorSpaces, RGBColorSpace};
use shared::textures::*;
use shared::utils::Ptr;
use shared::utils::mesh::{DeviceBilinearPatchMesh, DeviceTriangleMesh};
use shared::utils::sampling::{
    DevicePiecewiseConstant1D, DevicePiecewiseConstant2D, DeviceWindowedPiecewiseConstant2D,
};
use std::alloc::Layout;
use std::collections::HashMap;
use std::slice::from_raw_parts;
use std::sync::Arc;

pub struct Arena {
    buffer: Vec<(*mut u8, Layout)>,
    texture_cache: HashMap<usize, u64>,
}

impl Arena {
    pub fn new() -> Self {
        Self {
            buffer: Vec::new(),
            texture_cache: HashMap::new(),
        }
    }

    pub fn alloc<T>(&mut self, value: T) -> Ptr<T> {
        let layout = Layout::new::<T>();

        let ptr = unsafe { self.alloc_unified(layout) } as *mut T;

        unsafe {
            ptr.write(value);
        }

        Ptr::from_raw(ptr)
    }

    pub fn alloc_opt<T>(&mut self, value: Option<T>) -> Ptr<T> {
        match value {
            Some(v) => self.alloc(v),
            None => Ptr::null(),
        }
    }

    pub fn alloc_slice<T: Copy>(&mut self, values: &[T]) -> (Ptr<T>, usize) {
        if values.is_empty() {
            return (Ptr::null(), 0);
        }

        let layout = Layout::array::<T>(values.len()).unwrap();
        let ptr = unsafe { self.alloc_unified(layout) } as *mut T;

        unsafe {
            std::ptr::copy_nonoverlapping(values.as_ptr(), ptr, values.len());
        }

        (Ptr::from_raw(ptr), values.len())
    }

    #[cfg(feature = "cuda")]
    unsafe fn alloc_unified(&mut self, layout: Layout) -> *mut u8 {
        use cuda_runtime_sys::*;

        let mut ptr: *mut std::ffi::c_void = std::ptr::null_mut();
        let size = layout.size().max(layout.align());

        let result = unsafe { cudaMallocManaged(&mut ptr, size, cudaMemAttachGlobal) };

        if result != cudaError::cudaSuccess {
            panic!("cudaMallocManaged failed: {:?}", result);
        }

        self.buffer.push((ptr as *mut u8, layout));
        ptr as *mut u8
    }

    pub fn get_texture_object(&mut self, mipmap: &Arc<MIPMap>) -> u64 {
        let key = Arc::as_ptr(mipmap) as usize;

        if let Some(&tex_obj) = self.texture_cache.get(&key) {
            return tex_obj;
        }

        #[cfg(feature = "cuda")]
        let tex_obj = self.create_cuda_texture(mipmap);

        #[cfg(not(feature = "cuda"))]
        let tex_obj = 0u64;

        self.texture_cache.insert(key, tex_obj);
        tex_obj
    }

    #[cfg(feature = "cuda")]
    fn create_cuda_texture(&self, mipmap: &MIPMap) -> u64 {
        // TODO: Implement with cudarc
        // 1. Get image data from mipmap.base_image()
        // 2. cudaMallocArray
        // 3. cudaMemcpy2DToArray
        // 4. cudaCreateTextureObject
        // 5. Return handle
        0
    }

    #[cfg(not(feature = "cuda"))]
    unsafe fn alloc_unified(&mut self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { std::alloc::alloc(layout) };
        self.buffer.push((ptr, layout));
        ptr
    }

    // pub fn raw_data(&self) -> &[u8] {
    //     &self.buffer
    // }
}

impl Drop for Arena {
    fn drop(&mut self) {
        for (ptr, layout) in self.buffer.drain(..) {
            unsafe {
                #[cfg(feature = "cuda")]
                {
                    cuda_runtime_sys::cudaFree(ptr as *mut _);
                }
                #[cfg(not(feature = "cuda"))]
                {
                    std::alloc::dealloc(ptr, layout);
                }
            }
        }
    }
}

pub trait Upload {
    type Target: Copy;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target>;
}

impl Upload for Shape {
    type Target = Shape;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        arena.alloc(self.clone())
    }
}

impl Upload for Image {
    type Target = DeviceImage;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        arena.alloc(*self.device_image())
    }
}

impl Upload for Spectrum {
    type Target = Spectrum;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        arena.alloc(self.clone())
    }
}

impl Upload for DenselySampledSpectrumBuffer {
    type Target = DenselySampledSpectrum;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        arena.alloc(self.device())
    }
}

impl Upload for SpectrumTexture {
    type Target = GPUSpectrumTexture;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let gpu_variant = match self {
            SpectrumTexture::Constant(tex) => GPUSpectrumTexture::Constant(tex.clone()),
            SpectrumTexture::Checkerboard(tex) => GPUSpectrumTexture::Checkerboard(tex.clone()),
            SpectrumTexture::Dots(tex) => GPUSpectrumTexture::Dots(tex.clone()),
            SpectrumTexture::Image(tex) => {
                let tex_obj = arena.get_texture_object(&tex.base.mipmap);
                let gpu_img = GPUSpectrumImageTexture {
                    mapping: tex.base.mapping,
                    tex_obj,
                    scale: tex.base.scale,
                    invert: tex.base.invert,
                    is_single_channel: tex.base.mipmap.is_single_channel(),
                    color_space: tex.base.mipmap.color_space.clone().unwrap(),
                    spectrum_type: tex.spectrum_type,
                };
                GPUSpectrumTexture::Image(gpu_img)
            }
            SpectrumTexture::Bilerp(tex) => GPUSpectrumTexture::Bilerp(tex.clone()),
            SpectrumTexture::Scaled(tex) => {
                let child_ptr = tex.tex.upload(arena);

                let gpu_scaled = GPUSpectrumScaledTexture {
                    tex: child_ptr,
                    scale: tex.scale.upload(arena),
                };
                GPUSpectrumTexture::Scaled(gpu_scaled)
            }
            // SpectrumTexture::Ptex(tex) => {
            //     let gpu_ptex = GPUSpectrumPtexTexture::new(
            //         tex.base.filename,
            //         tex.base.encoding,
            //         tex.base.scale,
            //         tex.spectrum_type,
            //     );
            //     GPUSpectrumTexture::Ptex(gpu_ptex)
            // }
            SpectrumTexture::Marble(tex) => GPUSpectrumTexture::Marble(tex.clone()),
            SpectrumTexture::Mix(tex) => {
                let tex1_ptr = tex.tex1.upload(arena);
                let tex2_ptr = tex.tex2.upload(arena);
                let amount_ptr = tex.amount.upload(arena);

                let gpu_mix = GPUSpectrumMixTexture {
                    tex1: tex1_ptr,
                    tex2: tex2_ptr,
                    amount: amount_ptr,
                };

                GPUSpectrumTexture::Mix(gpu_mix)
            }
            SpectrumTexture::DirectionMix(tex) => {
                let tex1_ptr = tex.tex1.upload(arena);
                let tex2_ptr = tex.tex2.upload(arena);

                let gpu_mix = GPUSpectrumDirectionMixTexture {
                    tex1: tex1_ptr,
                    tex2: tex2_ptr,
                    dir: tex.dir,
                };

                GPUSpectrumTexture::DirectionMix(gpu_mix)
            }
        };

        arena.alloc(gpu_variant)
    }
}

impl Upload for FloatTexture {
    type Target = GPUFloatTexture;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let gpu_variant = match self {
            FloatTexture::Constant(tex) => GPUFloatTexture::Constant(tex.clone()),
            FloatTexture::Checkerboard(tex) => GPUFloatTexture::Checkerboard(tex.clone()),
            FloatTexture::Dots(tex) => GPUFloatTexture::Dots(tex.clone()),
            FloatTexture::FBm(tex) => GPUFloatTexture::FBm(tex.clone()),
            FloatTexture::Windy(tex) => GPUFloatTexture::Windy(tex.clone()),
            FloatTexture::Wrinkled(tex) => GPUFloatTexture::Wrinkled(tex.clone()),
            FloatTexture::Scaled(tex) => {
                let child_ptr = tex.tex.upload(arena);

                let gpu_scaled = GPUFloatScaledTexture {
                    tex: child_ptr,
                    scale: tex.scale.upload(arena),
                };
                GPUFloatTexture::Scaled(gpu_scaled)
            }

            FloatTexture::Mix(tex) => {
                let tex1_ptr = tex.tex1.upload(arena);
                let tex2_ptr = tex.tex2.upload(arena);
                let amount_ptr = tex.amount.upload(arena);

                let gpu_mix = GPUFloatMixTexture {
                    tex1: tex1_ptr,
                    tex2: tex2_ptr,
                    amount: amount_ptr,
                };
                GPUFloatTexture::Mix(gpu_mix)
            }

            FloatTexture::DirectionMix(tex) => {
                let tex1_ptr = tex.tex1.upload(arena);
                let tex2_ptr = tex.tex2.upload(arena);
                let gpu_dmix = GPUFloatDirectionMixTexture {
                    tex1: tex1_ptr,
                    tex2: tex2_ptr,
                    dir: tex.dir,
                };
                GPUFloatTexture::DirectionMix(gpu_dmix)
            }

            FloatTexture::Image(tex) => {
                let gpu_image_tex = GPUFloatImageTexture {
                    mapping: tex.base.mapping,
                    tex_obj: tex.base.mipmap.texture_object(),
                    scale: tex.base.scale,
                    invert: tex.base.invert,
                };
                GPUFloatTexture::Image(gpu_image_tex)
            }

            FloatTexture::Bilerp(tex) => GPUFloatTexture::Bilerp(tex.clone()),
        };

        arena.alloc(gpu_variant)
    }
}

impl Upload for RGBToSpectrumTable {
    type Target = RGBToSpectrumTable;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let n_nodes = self.n_nodes as usize;
        let z_slice = unsafe { from_raw_parts(self.z_nodes.as_raw(), n_nodes) };
        let coeffs_slice = unsafe { from_raw_parts(self.coeffs.as_raw(), n_nodes) };
        let (z_ptr, _) = arena.alloc_slice(z_slice);
        let (c_ptr, _) = arena.alloc_slice(coeffs_slice);

        let shared_table = RGBToSpectrumTable {
            z_nodes: z_ptr,
            coeffs: c_ptr,
            n_nodes: self.n_nodes,
        };

        arena.alloc(shared_table)
    }
}

impl Upload for RGBColorSpace {
    type Target = RGBColorSpace;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let table_ptr = self.rgb_to_spectrum_table.upload(arena);

        let shared_space = RGBColorSpace {
            r: self.r,
            g: self.g,
            b: self.b,
            w: self.w,
            illuminant: self.illuminant.clone(),
            rgb_to_spectrum_table: table_ptr,
            xyz_from_rgb: self.xyz_from_rgb,
            rgb_from_xyz: self.rgb_from_xyz,
        };

        arena.alloc(shared_space)
    }
}

impl Upload for DeviceStandardColorSpaces {
    type Target = DeviceStandardColorSpaces;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let srgb_ptr = self.srgb.upload(arena);
        let dci_ptr = self.dci_p3.upload(arena);
        let rec_ptr = self.rec2020.upload(arena);
        let aces_ptr = self.aces2065_1.upload(arena);

        let registry = DeviceStandardColorSpaces {
            srgb: srgb_ptr,
            dci_p3: dci_ptr,
            rec2020: rec_ptr,
            aces2065_1: aces_ptr,
        };

        arena.alloc(registry)
    }
}

impl Upload for PiecewiseConstant2D {
    type Target = DevicePiecewiseConstant2D;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let marginal_shared = self.marginal.to_shared(arena);

        let conditionals_shared: Vec<DevicePiecewiseConstant1D> = self
            .conditionals
            .iter()
            .map(|c| c.to_shared(arena))
            .collect();

        let (conditionals_ptr, _) = arena.alloc_slice(&conditionals_shared);

        let shared_2d = DevicePiecewiseConstant2D {
            conditionals: conditionals_ptr,
            marginal: marginal_shared,
            ..self.device
        };

        arena.alloc(shared_2d)
    }
}

impl Upload for WindowedPiecewiseConstant2D {
    type Target = DeviceWindowedPiecewiseConstant2D;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let specific = DeviceWindowedPiecewiseConstant2D {
            sat: self.sat,
            func: self.func,
        };
        arena.alloc(specific)
    }
}

impl Upload for TriangleMesh {
    type Target = DeviceTriangleMesh;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let storage = &self.storage;

        // Upload all arrays to arena
        let (vertex_indices_ptr, _) = arena.alloc_slice(&storage.vertex_indices);
        let (p_ptr, _) = arena.alloc_slice(&storage.p);

        let (n_ptr, _) = if storage.n.is_empty() {
            (Ptr::null(), 0)
        } else {
            arena.alloc_slice(&storage.n)
        };

        let (s_ptr, _) = if storage.s.is_empty() {
            (Ptr::null(), 0)
        } else {
            arena.alloc_slice(&storage.s)
        };

        let (uv_ptr, _) = if storage.uv.is_empty() {
            (Ptr::null(), 0)
        } else {
            arena.alloc_slice(&storage.uv)
        };

        let (face_indices_ptr, _) = if storage.face_indices.is_empty() {
            (Ptr::null(), 0)
        } else {
            arena.alloc_slice(&storage.face_indices)
        };

        let device = DeviceTriangleMesh {
            vertex_indices: vertex_indices_ptr,
            p: p_ptr,
            n: n_ptr,
            s: s_ptr,
            uv: uv_ptr,
            face_indices: face_indices_ptr,
            ..self.device
        };

        arena.alloc(device)
    }
}

impl Upload for BilinearPatchMesh {
    type Target = DeviceBilinearPatchMesh;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let storage = &self.storage;

        let (vertex_indices_ptr, _) = arena.alloc_slice(&storage.vertex_indices);
        let (p_ptr, _) = arena.alloc_slice(&storage.p);

        let (n_ptr, _) = if storage.n.is_empty() {
            (Ptr::null(), 0)
        } else {
            arena.alloc_slice(&storage.n)
        };

        let (uv_ptr, _) = if storage.uv.is_empty() {
            (Ptr::null(), 0)
        } else {
            arena.alloc_slice(&storage.uv)
        };

        let image_dist_ptr = storage.image_distribution.upload(arena);

        let device = DeviceBilinearPatchMesh {
            vertex_indices: vertex_indices_ptr,
            p: p_ptr,
            n: n_ptr,
            uv: uv_ptr,
            image_distribution: image_dist_ptr,
            ..self.device
        };

        arena.alloc(device)
    }
}

impl<T: Upload> Upload for Option<T> {
    type Target = T::Target;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        match self {
            Some(val) => val.upload(arena),
            None => Ptr::null(),
        }
    }
}

impl<T: Upload> Upload for Arc<T> {
    type Target = T::Target;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        (**self).upload(arena)
    }
}
