use std::borrow::Cow;

use arrow::array::NullBufferBuilder;
use common_error::DaftResult;
use common_image::CowImage;
use num_traits::FromPrimitive;

use crate::{array::image_array::ImageArraySidecarData, prelude::*};

#[allow(clippy::len_without_is_empty)]
pub trait AsImageObj {
    fn name(&self) -> &str;
    fn len(&self) -> usize;
    fn as_image_obj(&self, idx: usize) -> Option<CowImage<'_>>;
}

impl AsImageObj for ImageArray {
    fn len(&self) -> usize {
        ImageArray::len(self)
    }

    fn name(&self) -> &str {
        ImageArray::name(self)
    }

    fn as_image_obj(&self, idx: usize) -> Option<CowImage<'_>> {
        assert!(idx < self.len());
        if !self.physical.is_valid(idx) {
            return None;
        }

        let da = self.data_array();

        let offsets = da.offsets();

        let start = *offsets.get(idx).unwrap() as usize;
        let end = *offsets.get(idx + 1).unwrap() as usize;

        let values = da.flat_child.u8().unwrap().as_slice();
        let slice_data = Cow::Borrowed(&values[start..end]);

        let c = self.channels().values()[idx];
        let h = self.heights().values()[idx];
        let w = self.widths().values()[idx];
        let m: ImageMode = ImageMode::from_u8(self.modes().values()[idx]).unwrap();
        assert_eq!(m.num_channels(), c);
        let result = CowImage::from_raw(&m, w, h, slice_data);

        assert_eq!(result.height(), h);
        assert_eq!(result.width(), w);
        Some(result)
    }
}

impl AsImageObj for FixedShapeImageArray {
    fn len(&self) -> usize {
        FixedShapeImageArray::len(self)
    }

    fn name(&self) -> &str {
        FixedShapeImageArray::name(self)
    }

    fn as_image_obj(&self, idx: usize) -> Option<CowImage<'_>> {
        assert!(idx < self.len());
        if !self.physical.is_valid(idx) {
            return None;
        }

        match self.data_type() {
            DataType::FixedShapeImage(mode, height, width) => {
                let values = self
                    .physical
                    .flat_child
                    .downcast::<UInt8Array>()
                    .unwrap()
                    .as_slice();
                let num_channels = mode.num_channels();
                let size = height * width * u32::from(num_channels);
                let start = idx * size as usize;
                let end = (idx + 1) * size as usize;
                let slice_data = Cow::Borrowed(&values[start..end]);
                let result = CowImage::from_raw(mode, *width, *height, slice_data);

                assert_eq!(result.height(), *height);
                assert_eq!(result.width(), *width);
                Some(result)
            }
            dt => panic!(
                "FixedShapeImageArray should always have DataType::FixedShapeImage() as it's dtype, but got {dt}"
            ),
        }
    }
}

pub fn image_array_from_img_buffers<'a, I>(
    name: &str,
    inputs: I,
    image_mode: Option<ImageMode>,
) -> DaftResult<ImageArray>
where
    I: ExactSizeIterator<Item = Option<CowImage<'a>>>,
{
    use CowImage::{L, LA, RGB, RGBA};

    let mut data = Vec::with_capacity(inputs.len());
    let mut heights = Vec::with_capacity(inputs.len());
    let mut channels = Vec::with_capacity(inputs.len());
    let mut modes = Vec::with_capacity(inputs.len());
    let mut widths = Vec::with_capacity(inputs.len());
    let mut offsets = Vec::with_capacity(inputs.len() + 1);
    offsets.push(0i64);
    let mut null_builder = NullBufferBuilder::new(inputs.len());

    for ib in inputs {
        null_builder.append(ib.is_some());
        match ib {
            Some(ib) => {
                assert!(matches!(&ib, L(..) | LA(..) | RGB(..) | RGBA(..)));
                let height = ib.height();
                let width = ib.width();
                let mode = ib.mode();
                let buffer = ib.as_u8_slice();
                heights.push(height);
                widths.push(width);
                modes.push(mode as u8);
                channels.push(mode.num_channels());
                data.extend_from_slice(buffer);
                offsets.push(offsets.last().unwrap() + buffer.len() as i64);
            }
            None => {
                heights.push(0u32);
                widths.push(0u32);
                modes.push(ImageMode::L as u8);
                channels.push(ImageMode::L.num_channels());
                offsets.push(*offsets.last().unwrap());
            }
        }
    }

    let nulls = null_builder.finish();
    ImageArray::from_vecs(
        name,
        DataType::Image(image_mode),
        data,
        offsets,
        ImageArraySidecarData {
            channels,
            heights,
            widths,
            modes,
            nulls,
        },
    )
}

pub fn fixed_image_array_from_img_buffers(
    name: &str,
    inputs: &[Option<CowImage<'_>>],
    image_mode: &ImageMode,
    height: u32,
    width: u32,
) -> DaftResult<FixedShapeImageArray> {
    use CowImage::{L, LA, RGB, RGBA};
    let is_all_u8 = inputs
        .iter()
        .filter_map(|b| b.as_ref())
        .all(|b| matches!(b, L(..) | LA(..) | RGB(..) | RGBA(..)));
    assert!(is_all_u8);

    let num_channels = image_mode.num_channels();
    let mut data_ref = Vec::with_capacity(inputs.len());
    let mut null_builder = NullBufferBuilder::new(inputs.len());
    let list_size = (height * width * u32::from(num_channels)) as usize;
    let null_list = vec![0u8; list_size];
    for ib in inputs {
        null_builder.append(ib.is_some());
        let buffer = match ib {
            Some(ib) => ib.as_u8_slice(),
            None => null_list.as_slice(),
        };
        data_ref.push(buffer);
    }
    let data = data_ref.concat();
    let nulls = null_builder.finish();

    let flat_child = UInt8Array::from_vec("data", data).into_series();
    let daft_dtype = DataType::FixedSizeList(Box::new(DataType::UInt8), list_size);
    let physical_array = FixedSizeListArray::new(Field::new(name, daft_dtype), flat_child, nulls);
    let logical_dtype = DataType::FixedShapeImage(*image_mode, height, width);
    Ok(FixedShapeImageArray::new(
        Field::new(name, logical_dtype),
        physical_array,
    ))
}
