use super::*;

#[derive(derivative::Derivative)]
#[derivative(Debug)]
pub struct DVec<T, A: DeviceAllocator = GlobalDeviceStatic> {
    #[derivative(Debug = "ignore")]
    pub(crate) buf: RawDVec<T, A>,
}

impl<T, A: DeviceAllocator> DVec<T, A> {
    pub unsafe fn into_owned_chunks(self, chunk_size: usize) -> Vec<Self> {
        assert_eq!(self.len() % chunk_size, 0);
        let num_chunks = self.len() / chunk_size;
        let mut result = vec![];
        let (ptr, len, alloc, pool, stream) = self.into_raw_parts();
        assert_eq!(len, chunk_size * num_chunks);
        for chunk_idx in 0..num_chunks {
            let ptr = ptr.add(chunk_idx * chunk_size);
            let new = Self::from_raw_parts_in(ptr, chunk_size, A::default(), pool, stream);
            result.push(new);
        }
        result
    }

    pub unsafe fn split_into_owned_array(self, mid: usize) -> [Self; 2] {
        assert!(mid < self.len());
        let (ptr, len, _, pool, stream) = self.into_raw_parts();

        let left = Self::from_raw_parts_in(ptr, mid, A::default(), pool, stream);
        let right = Self::from_raw_parts_in(ptr.add(mid), len - mid, A::default(), pool, stream);

        [left, right]
    }

    pub fn len(&self) -> usize {
        self.buf.len()
    }

    pub fn as_ptr(&self) -> *const T {
        self.buf.as_ptr().cast()
    }

    pub fn as_mut_ptr(&mut self) -> *mut T {
        self.buf.as_ptr()
    }

    pub fn get(&self, index: usize) -> &T {
        self.as_ref().get(index)
    }

    pub fn get_mut(&mut self, index: usize) -> &mut T {
        self.as_mut().get_mut(index)
    }

    pub fn split_at(&self, mid: usize) -> (&DSlice<T>, &DSlice<T>) {
        self.as_ref().split_at(mid)
    }

    pub fn split_at_mut(&mut self, mid: usize) -> (&mut DSlice<T>, &mut DSlice<T>) {
        self.as_mut().split_at_mut(mid)
    }

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

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

    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 to_vec_on(&self, stream: bc_stream) -> CudaResult<Vec<T>> {
        self.as_ref().to_vec(stream)
    }

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

    pub unsafe fn into_raw_parts(
        self,
    ) -> (*mut T, usize, A, Option<bc_mem_pool>, Option<bc_stream>) {
        let mut me = std::mem::ManuallyDrop::new(self);
        let len = me.buf.len();
        let ptr = me.as_mut_ptr();
        (ptr, len, A::default(), me.buf.pool, me.buf.stream)
    }

    pub unsafe fn from_raw_parts_in(
        ptr: *mut T,
        len: usize,
        alloc: A,
        pool: Option<bc_mem_pool>,
        stream: Option<bc_stream>,
    ) -> Self {
        unsafe {
            DVec {
                buf: RawDVec::from_raw_parts_in(ptr, len, alloc, pool, stream),
            }
        }
    }
}

impl<T> DVec<T, PoolAllocator> {
    pub fn allocate_on(length: usize, pool: bc_mem_pool, stream: bc_stream) -> Self {
        Self {
            buf: RawDVec::allocate_on(length, PoolAllocator, pool, stream),
        }
    }

    pub fn allocate_zeroed_on(length: usize, pool: bc_mem_pool, stream: bc_stream) -> Self {
        Self {
            buf: RawDVec::allocate_zeroed_on(length, PoolAllocator, pool, stream),
        }
    }

    pub fn from_host_slice_on(src: &[T], pool: bc_mem_pool, stream: bc_stream) -> CudaResult<Self> {
        let mut dst = Self::allocate_on(src.len(), pool, stream);
        mem::h2d_on(src, &mut dst, stream)?;

        Ok(dst)
    }
}

impl<T> DVec<T, GlobalDeviceStatic> {
    pub fn dangling() -> Self {
        Self {
            buf: RawDVec::dangling(),
        }
    }

    pub fn allocate(length: usize) -> Self {
        Self {
            buf: RawDVec::allocate(length, _static_alloc()),
        }
    }

    pub fn allocate_zeroed(length: usize) -> Self {
        Self {
            buf: RawDVec::allocate_zeroed(length, _static_alloc()),
        }
    }

    pub fn from_host_slice(src: &[T], stream: bc_stream) -> CudaResult<Self> {
        let mut dst = Self::allocate(src.len());
        mem::h2d_on(src, &mut dst, stream)?;

        Ok(dst)
    }
}

impl<T> CloneStatic for DVec<T, GlobalDeviceStatic> {
    fn clone(&self, stream: bc_stream) -> CudaResult<Self> {
        let mut new = Self::allocate(self.len());
        mem::d2d_on(self, &mut new, stream)?;

        Ok(new)
    }
}

impl<T> CloneOnPool for DVec<T, PoolAllocator> {
    fn clone_on(&self, pool: bc_mem_pool, stream: bc_stream) -> CudaResult<Self> {
        let mut new = Self::allocate_on(self.len(), pool, stream);
        mem::d2d_on(self, &mut new, stream)?;

        Ok(new)
    }
}

impl<T, A: DeviceAllocator> std::ops::Deref for DVec<T, A> {
    type Target = DSlice<T>;

    #[inline]
    fn deref(&self) -> &Self::Target {
        unsafe { DSlice::from_raw_parts(self.buf.as_ptr(), self.buf.len()) }
    }
}

impl<T, A: DeviceAllocator> std::ops::DerefMut for DVec<T, A> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { DSlice::from_raw_parts_mut(self.buf.as_mut_ptr(), self.buf.len()) }
    }
}

impl<T, A: DeviceAllocator> AsRef<DSlice<T>> for DVec<T, A> {
    fn as_ref(&self) -> &DSlice<T> {
        self
    }
}

impl<T, A: DeviceAllocator> AsMut<DSlice<T>> for DVec<T, A> {
    fn as_mut(&mut self) -> &mut DSlice<T> {
        self
    }
}

impl<T, A: DeviceAllocator> std::ops::Index<usize> for DVec<T, A> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        self.get(index)
    }
}

impl<T, A: DeviceAllocator> std::ops::IndexMut<usize> for DVec<T, A> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        self.get_mut(index)
    }
}

impl<T, A: DeviceAllocator> std::ops::Index<std::ops::Range<usize>> for DVec<T, A> {
    type Output = DSlice<T>;

    fn index(&self, index: std::ops::Range<usize>) -> &Self::Output {
        &self.as_ref()[index]
    }
}

impl<T, A: DeviceAllocator> std::ops::IndexMut<std::ops::Range<usize>> for DVec<T, A> {
    fn index_mut(&mut self, index: std::ops::Range<usize>) -> &mut Self::Output {
        &mut self.as_mut()[index]
    }
}

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

    fn index(&self, index: std::ops::RangeTo<usize>) -> &Self::Output {
        &self.as_ref()[index]
    }
}

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

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

    fn index(&self, index: std::ops::RangeFrom<usize>) -> &Self::Output {
        &self.as_ref()[index]
    }
}

impl<T, A: DeviceAllocator> std::ops::IndexMut<std::ops::RangeFrom<usize>> for DVec<T, A> {
    fn index_mut(&mut self, index: std::ops::RangeFrom<usize>) -> &mut Self::Output {
        &mut self.as_mut()[index]
    }
}
