use gpu_ffi::{bc_mem_pool, bc_stream};

use super::{
    CudaResult, DChunks, DChunksMut, DIter, DIterMut, DVec, GlobalDeviceStatic, HostAllocator,
};

pub struct DSlice<T>([T]);

impl<T> DSlice<T> {
    pub const fn len(&self) -> usize {
        std::ptr::metadata(self)
    }

    pub const fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub unsafe fn from_raw_parts<'a>(ptr: *const T, len: usize) -> &'a Self {
        std::mem::transmute(std::slice::from_raw_parts(ptr, len))
    }

    pub unsafe fn from_raw_parts_mut<'a>(ptr: *mut T, len: usize) -> &'a mut Self {
        std::mem::transmute(std::slice::from_raw_parts_mut(ptr, len))
    }

    pub fn as_ptr(&self) -> *const T {
        &self.0 as *const [T] as *mut T
    }

    pub fn as_mut_ptr(&mut self) -> *mut T {
        &mut self.0 as *mut [T] as *mut T
    }

    pub fn get(&self, index: usize) -> &T {
        unsafe { &*self.as_ptr().add(index) }
    }

    pub fn get_mut(&mut self, index: usize) -> &mut T {
        unsafe { &mut *self.as_mut_ptr().add(index) }
    }

    pub fn split_at(&self, mid: usize) -> (&Self, &Self) {
        let ptr = self.as_ptr();
        let len = self.len();
        unsafe {
            (
                DSlice::from_raw_parts(ptr, mid),
                DSlice::from_raw_parts(ptr.add(mid), len - mid),
            )
        }
    }

    pub fn split_at_mut(&mut self, mid: usize) -> (&mut Self, &mut Self) {
        let ptr = self.as_mut_ptr();
        let len = self.len();
        unsafe {
            (
                DSlice::from_raw_parts_mut(ptr, mid),
                DSlice::from_raw_parts_mut(ptr.add(mid), len - mid),
            )
        }
    }

    pub fn chunks(&self, chunk_size: usize) -> DChunks<T> {
        DChunks::new(self, chunk_size)
    }

    pub fn chunks_mut(&mut self, chunk_size: usize) -> DChunksMut<T> {
        DChunksMut::new(self, chunk_size)
    }

    pub fn iter<'a>(&'a self) -> DIter<'a, T> {
        DIter::new(self)
    }

    pub fn iter_mut<'a>(&'a mut self) -> DIterMut<'a, T> {
        DIterMut::new(self)
    }

    pub fn copy_from_slice_on(&mut self, other: &DSlice<T>, stream: bc_stream) -> CudaResult<()> {
        super::mem::d2d_on(other, self, stream)
    }

    pub fn to_vec(&self, stream: bc_stream) -> CudaResult<Vec<T>> {
        self.to_vec_in(stream, std::alloc::Global)
    }

    pub fn to_vec_in<A: HostAllocator>(
        &self,
        stream: bc_stream,
        alloc: A,
    ) -> CudaResult<Vec<T, A>> {
        let mut dst: Vec<_, A> = Vec::with_capacity_in(self.len(), alloc);
        unsafe { dst.set_len(self.len()) };
        super::mem::d2h_on(self, &mut dst, stream)?;

        Ok(dst)
    }

    pub fn to_dvec(&self, stream: bc_stream) -> CudaResult<DVec<T, GlobalDeviceStatic>> {
        let mut dst = DVec::allocate_zeroed(self.len());
        super::mem::d2d_on(self, &mut dst, stream)?;

        Ok(dst)
    }
}

impl<T> std::ops::Index<usize> for DSlice<T> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        self.get(index)
    }
}
impl<T> std::ops::IndexMut<usize> for DSlice<T> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        self.get_mut(index)
    }
}
impl<T> std::ops::Index<std::ops::Range<usize>> for DSlice<T> {
    type Output = DSlice<T>;

    fn index(&self, index: std::ops::Range<usize>) -> &Self::Output {
        let start = index.start;
        let len = index.len();

        unsafe {
            let ptr = self.as_ptr().add(start);
            Self::from_raw_parts(ptr, len)
        }
    }
}
impl<T> std::ops::IndexMut<std::ops::Range<usize>> for DSlice<T> {
    fn index_mut(&mut self, index: std::ops::Range<usize>) -> &mut Self::Output {
        let start = index.start;
        let len = index.len();

        unsafe {
            let ptr = self.as_mut_ptr().add(start);
            Self::from_raw_parts_mut(ptr, len)
        }
    }
}

impl<T> std::ops::Index<std::ops::RangeFrom<usize>> for DSlice<T> {
    type Output = DSlice<T>;

    fn index(&self, index: std::ops::RangeFrom<usize>) -> &Self::Output {
        let end = self.len();
        &self[index.start..end]
    }
}
impl<T> std::ops::IndexMut<std::ops::RangeFrom<usize>> for DSlice<T> {
    fn index_mut(&mut self, index: std::ops::RangeFrom<usize>) -> &mut Self::Output {
        let end = self.len();
        &mut self[index.start..end]
    }
}

impl<T> std::ops::Index<std::ops::RangeTo<usize>> for DSlice<T> {
    type Output = DSlice<T>;

    fn index(&self, index: std::ops::RangeTo<usize>) -> &Self::Output {
        &self[0..index.end]
    }
}
impl<T> std::ops::IndexMut<std::ops::RangeTo<usize>> for DSlice<T> {
    fn index_mut(&mut self, index: std::ops::RangeTo<usize>) -> &mut Self::Output {
        &mut self[0..index.end]
    }
}

impl<T> std::ops::Index<std::ops::RangeFull> for DSlice<T> {
    type Output = DSlice<T>;

    fn index(&self, _index: std::ops::RangeFull) -> &Self::Output {
        &self[0..self.len()]
    }
}

impl<T> std::ops::IndexMut<std::ops::RangeFull> for DSlice<T> {
    fn index_mut(&mut self, _index: std::ops::RangeFull) -> &mut Self::Output {
        let end = self.len();
        &mut self[0..end]
    }
}
