kahlo 0.0.2

use bytemuck::cast_slice_mut;

use crate::{
    bitmap::{BitmapAccess, BitmapAccessMut, BitmapRawData, BitmapRawDataMut},
    colour::{BlendMode, Colour},
    formats::{self, PixelFormatSpec},
    math::{clip_to, PixelBox, PixelPoint, PixelSize},
};

mod dispatch;
use dispatch::{BinaryDispatch, GenericBinary, GenericUnary, UnaryDispatch, DISPATCH};

/// AVX helper types
#[cfg(target_arch = "x86_64")]
mod avx;

/// Common operational elements
#[cfg(target_arch = "x86_64")]
mod common_x86_64;

// operation implementations
mod composite;
mod copy_from;
mod fill;
mod fill_region;
mod fill_region_masked;
mod rectangle;

#[cfg(any(test, feature = "benchmark"))]
mod generate;

pub trait UnaryOpSpec: GenericUnary<Self::Params<'static>> {
    type Params<'l>: std::fmt::Debug;

    fn build<'l>() -> UnaryDispatch<Self::Params<'l>>;
}

pub trait BinaryOpSpec: GenericBinary<Self::Params<'static>> {
    type Params<'l>: std::fmt::Debug;

    fn build<'l>() -> BinaryDispatch<Self::Params<'l>>;
}

pub trait KahloOps<Format: PixelFormatSpec>: BitmapAccessMut<Format> {
    fn composite<SrcFormat: PixelFormatSpec>(
        &mut self,
        src: &dyn BitmapAccess<SrcFormat>,
        src_area: PixelBox,
        dst: PixelPoint,
        blend: BlendMode,
    ) where
        Self: Sized,
    {
        DISPATCH.composite.select::<Format, SrcFormat>()(self, src, &(src_area, dst, blend))
    }

    /// Copy pixels from one bitmap view to another, converting formats if needed.
    fn copy_from<SrcFormat: PixelFormatSpec>(
        &mut self,
        src: &dyn BitmapAccess<SrcFormat>,
        src_area: PixelBox,
        dst: PixelPoint,
    ) where
        Self: Sized,
    {
        DISPATCH.copy_from.select::<Format, SrcFormat>()(self, src, &(src_area, dst));
    }

    /// Fill an entire bitmap view with a solid colour.
    fn fill(&mut self, colour: Colour)
    where
        Self: Sized,
    {
        DISPATCH.fill.select::<Format>()(self, &(colour,));
    }

    /// Fill part of a bitmap view with a solid colour.
    fn fill_region(&mut self, region: PixelBox, colour: Colour)
    where
        Self: Sized,
    {
        DISPATCH.fill_region.select::<Format>()(self, &(region, colour));
    }

    /// Fill part of a bitmap view with a solid colour, masked by an alphamap.
    fn fill_region_masked<'l>(
        &mut self,
        mask: &'l dyn BitmapAccess<formats::A8>,
        src_area: PixelBox,
        dst: PixelPoint,
        colour: Colour,
        blend: BlendMode,
    ) where
        Self: Sized,
    {
        DISPATCH
            .use_lifetime()
            .fill_region_masked
            .select::<Format>()(self, &(mask, src_area, dst, colour, blend))
    }

    /// Draw a rectangle outline on a bitmap view.
    fn rectangle(&mut self, area: PixelBox, width: usize, colour: Colour)
    where
        Self: Sized,
    {
        DISPATCH.rectangle.select::<Format>()(self, &(area, width, colour));
    }
}

impl<Format: PixelFormatSpec, T: BitmapAccessMut<Format>> KahloOps<Format> for T {}

#[inline]
fn binary_map_pixel_data_with<
    const GROUP_SIZE: usize,
    Format: PixelFormatSpec,
    SrcFormat: PixelFormatSpec,
>(
    dst_data: &mut dyn BitmapRawDataMut,
    src_data: &dyn BitmapRawData,
    size: PixelSize,
    dst_start: PixelPoint,
    src_start: PixelPoint,
    map: impl Fn(&mut [u8], &[u8]),
) {
    let Some((write_box, read_box)) = clip_to(
        PixelBox::from_size(dst_data.size()),
        PixelBox::from_size(src_data.size()),
        PixelBox::from_origin_and_size(dst_start, size),
        PixelBox::from_origin_and_size(src_start, size),
    ) else {
        return;
    };

    let read_stride = src_data.stride();
    let write_stride = dst_data.stride();
    let read_data = src_data.data();
    let write_data = dst_data.data_mut();
    for y in 0..read_box.height() {
        let y_offset_r = (y + read_box.min.y) as usize * read_stride;
        let y_offset_w = (y + write_box.min.y) as usize * write_stride;
        let x_offset_r = read_box.min.x as usize * SrcFormat::PIXEL_STRIDE;
        let x_offset_w = write_box.min.x as usize * Format::PIXEL_STRIDE;

        let x_len_r = read_box.width() as usize * SrcFormat::PIXEL_STRIDE;
        let x_len_w = read_box.width() as usize * Format::PIXEL_STRIDE;
        // arbitrary group size?
        if GROUP_SIZE == 0 {
            (map)(
                &mut write_data[y_offset_w + x_offset_w..y_offset_w + x_offset_w + x_len_w],
                &read_data[y_offset_r + x_offset_r..y_offset_r + x_offset_r + x_len_r],
            );
        } else {
            let mut w_off = y_offset_w + x_offset_w;
            let mut r_off = y_offset_r + x_offset_r;

            let w_step = GROUP_SIZE * Format::PIXEL_STRIDE;
            let r_step = GROUP_SIZE * SrcFormat::PIXEL_STRIDE;

            let w_end = w_off + x_len_w;
            let r_end = r_off + x_len_r;
            // println!("\tconsidering full-size chunk with w_off={w_off} and w_end={w_end}");
            while w_off + w_step <= w_end {
                map(&mut write_data[w_off..], &read_data[r_off..]);

                w_off += w_step;
                r_off += r_step;
            }

            // do we need to handle the last bit specially?
            if w_off != w_end {
                // this next bit is a hack because we can't declare an array of size GROUP_SIZE*4.
                let mut dst_buf_raw = [0u32; GROUP_SIZE];
                let mut src_buf_raw = [0u32; GROUP_SIZE];
                let dst_buf = cast_slice_mut::<u32, u8>(&mut dst_buf_raw);
                let src_buf = cast_slice_mut::<u32, u8>(&mut src_buf_raw);

                let dst_len = w_end - w_off;
                let src_len = r_end - r_off;

                (&mut dst_buf[..dst_len]).copy_from_slice(&write_data[w_off..w_end]);
                (&mut src_buf[..src_len]).copy_from_slice(&read_data[r_off..r_end]);

                (map)(
                    &mut dst_buf[..Format::PIXEL_STRIDE * GROUP_SIZE],
                    &src_buf[..SrcFormat::PIXEL_STRIDE * GROUP_SIZE],
                );

                (&mut write_data[w_off..w_end]).copy_from_slice(&mut dst_buf[..dst_len]);
            }
        }
    }
}

#[cfg(test)]
mod fuzz;

#[cfg(feature = "benchmark")]
pub(crate) mod benchmark;