use crate::ffi::vec_znx_big;
use crate::znx_base::{ZnxInfos, ZnxView};
use crate::{Backend, DataView, DataViewMut, FFT64, Module, alloc_aligned};
use std::marker::PhantomData;

// const VEC_ZNX_BIG_ROWS: usize = 1;

/// VecZnxBig is `Backend` dependent, denoted with backend generic `B`
pub struct VecZnxBig<D, B> {
    data: D,
    n: usize,
    cols: usize,
    size: usize,
    _phantom: PhantomData<B>,
}

impl<D, B> ZnxInfos for VecZnxBig<D, B> {
    fn cols(&self) -> usize {
        self.cols
    }

    fn rows(&self) -> usize {
        1
    }

    fn n(&self) -> usize {
        self.n
    }

    fn size(&self) -> usize {
        self.size
    }

    fn sl(&self) -> usize {
        self.cols() * self.n()
    }
}

impl<D, B> DataView for VecZnxBig<D, B> {
    type D = D;
    fn data(&self) -> &Self::D {
        &self.data
    }
}

impl<D, B> DataViewMut for VecZnxBig<D, B> {
    fn data_mut(&mut self) -> &mut Self::D {
        &mut self.data
    }
}

impl<D: AsRef<[u8]>> ZnxView for VecZnxBig<D, FFT64> {
    type Scalar = i64;
}

pub(crate) fn bytes_of_vec_znx_big<B: Backend>(module: &Module<B>, cols: usize, size: usize) -> usize {
    unsafe { vec_znx_big::bytes_of_vec_znx_big(module.ptr, size as u64) as usize * cols }
}

impl<D: From<Vec<u8>>, B: Backend> VecZnxBig<D, B> {
    pub(crate) fn new(module: &Module<B>, cols: usize, size: usize) -> Self {
        let data = alloc_aligned::<u8>(bytes_of_vec_znx_big(module, cols, size));
        Self {
            data: data.into(),
            n: module.n(),
            cols,
            size,
            _phantom: PhantomData,
        }
    }

    pub(crate) fn new_from_bytes(module: &Module<B>, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
        let data: Vec<u8> = bytes.into();
        assert!(data.len() == bytes_of_vec_znx_big(module, cols, size));
        Self {
            data: data.into(),
            n: module.n(),
            cols,
            size,
            _phantom: PhantomData,
        }
    }
}

impl<D, B> VecZnxBig<D, B> {
    pub(crate) fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
        Self {
            data,
            n,
            cols,
            size,
            _phantom: PhantomData,
        }
    }
}

pub type VecZnxBigOwned<B> = VecZnxBig<Vec<u8>, B>;

impl<B> VecZnxBig<Vec<u8>, B> {
    pub fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
        VecZnxBig {
            data: self.data.as_mut_slice(),
            n: self.n,
            cols: self.cols,
            size: self.size,
            _phantom: PhantomData,
        }
    }

    pub fn to_ref(&self) -> VecZnxBig<&[u8], B> {
        VecZnxBig {
            data: self.data.as_slice(),
            n: self.n,
            cols: self.cols,
            size: self.size,
            _phantom: PhantomData,
        }
    }
}