ruisa 0.0.1

// Copyright 2006 The Android Open Source Project
// Copyright 2020 Yevhenii Reizner
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use alloc::vec;
use alloc::vec::Vec;
use image::{ImageBuffer, Rgba, GenericImageView, ImageFormat};

use core::convert::TryFrom;
use core::num::NonZeroUsize;
use std::println;

use ruisa_path::{IntSize, ScreenIntRect};

use crate::{Color, IntRect};

use crate::color::PremultipliedColorU8;
use image::{DynamicImage, ImageError};

/// Number of bytes per pixel.
pub const BYTES_PER_PIXEL: usize = 4;

/// A container that owns premultiplied RGBA pixels.
///
/// The data is not aligned, therefore width == stride.
#[derive(Clone, PartialEq)]
pub struct Canvas {
    data: Vec<u8>,
    size: IntSize,
}

impl Canvas {
    /// Allocates a new canvas.
    ///
    /// A canvas is filled with transparent black by default, aka (0, 0, 0, 0).
    ///
    /// Zero size in an error.
    ///
    /// Canvas's width is limited by i32::MAX/4.
    pub fn new(width: u32, height: u32) -> Option<Self> {
        let size = IntSize::from_wh(width, height)?;
        let data_len = data_len_for_size(size)?;

        // We cannot check that allocation was successful yet.
        // We have to wait for https://github.com/rust-lang/rust/issues/48043

        Some(Canvas {
            data: vec![0; data_len],
            size,
        })
    }

    ///
    /// Dump debug info.
    /// 
    pub fn dump(&self) {
        println!("Canvas width:{}, height: {}, data_len: {}", self.width(), self.height(), self.data.len());
    }

    ///
    /// 把一个图片对象转化成一个CanvasRef
    /// 
    pub fn from(buffer: &ImageBuffer<Rgba<u8>, Vec<u8>>) -> Self {
        Canvas {
            data: buffer.to_vec(),
            size:IntSize::from_wh(buffer.width(), buffer.height()).unwrap(),
        }
    }

    /// Creates a new canvas by taking ownership over an image buffer
    /// (premultiplied RGBA pixels).
    ///
    /// The size needs to match the data provided.
    ///
    /// Canvas's width is limited by i32::MAX/4.
    pub fn from_vec(data: Vec<u8>, size: IntSize) -> Option<Self> {
        let data_len = data_len_for_size(size)?;
        if data.len() != data_len {
            return None;
        }

        Some(Canvas { data, size })
    }

    /// Decodes a PNG data into a `Canvas`.
    ///
    /// Only 8-bit images are supported.
    /// Index PNGs are not supported.
    pub fn from_image(data: &[u8]) -> Result<Self, ImageError> {
        let image: DynamicImage = image::load_from_memory(data).unwrap();

        let buffer = ImageBuffer::from_fn(image.width(), image.height(), |x, y|{
            image.get_pixel(x, y)
        });

        Ok(Canvas {
            data: buffer.to_vec(),
            size:IntSize::from_wh(image.width(), image.height()).unwrap(),
        })
    }

    /// Loads a PNG file into a `Canvas`.
    ///
    /// Only 8-bit images are supported.
    /// Index PNGs are not supported.
    pub fn from_image_file<P: AsRef<std::path::Path>>(path: P) -> Result<Self, ImageError> {
        let image = image::open(path).unwrap();
        let buffer = ImageBuffer::from_fn(image.width(), image.height(), |x, y|{
            image.get_pixel(x, y)
        });

        Ok(Canvas {
            data: buffer.to_vec(),
            size:IntSize::from_wh(image.width(), image.height()).unwrap(),
        })
    }

    /// Encodes canvas into a PNG data.
    pub fn to_png(&self) -> Result<Vec<u8>, ImageError> {
        self.as_ref().to_png()
    }

    /// Saves canvas as a PNG file.
    pub fn save_png<P: AsRef<std::path::Path>>(&self, path: P) -> Result<(), ImageError> {
        self.as_ref().save_png(path)
    }

    /// Saves canvas as a PNG file.
    pub fn save_image<P: AsRef<std::path::Path>>(&self, path: P, format: ImageFormat) -> Result<(), ImageError> {
        self.as_ref().save_image(path, format)
    }

    /// Returns a container that references Canvas's data.
    pub fn as_ref(&self) -> CanvasRef {
        CanvasRef {
            data: &self.data,
            size: self.size,
        }
    }

    /// Returns a container that references Canvas's data.
    pub fn as_mut(&mut self) -> CanvasMut {
        CanvasMut {
            data: &mut self.data,
            size: self.size,
        }
    }

    /// Returns canvas's width.
    #[inline]
    pub fn width(&self) -> u32 {
        self.size.width()
    }

    /// Returns canvas's height.
    #[inline]
    pub fn height(&self) -> u32 {
        self.size.height()
    }

    /// Returns canvas's size.
    #[allow(dead_code)]
    pub(crate) fn size(&self) -> IntSize {
        self.size
    }

    /// Fills the entire canvas with a specified color.
    pub fn fill(&mut self, color: Color) {
        let c = color.premultiply().to_color_u8();
        for p in self.as_mut().pixels_mut() {
            *p = c;
        }
    }

    /// Returns the internal data.
    ///
    /// Byteorder: RGBA
    pub fn data(&self) -> &[u8] {
        self.data.as_slice()
    }

    /// Returns the mutable internal data.
    ///
    /// Byteorder: RGBA
    pub fn data_mut(&mut self) -> &mut [u8] {
        self.data.as_mut_slice()
    }

    /// Returns a pixel color.
    ///
    /// Returns `None` when position is out of bounds.
    pub fn pixel(&self, x: u32, y: u32) -> Option<PremultipliedColorU8> {
        let idx = self.width().checked_mul(y)?.checked_add(x)?;
        self.pixels().get(idx as usize).cloned()
    }

    /// Returns a mutable slice of pixels.
    pub fn pixels_mut(&mut self) -> &mut [PremultipliedColorU8] {
        bytemuck::cast_slice_mut(self.data_mut())
    }

    /// Returns a slice of pixels.
    pub fn pixels(&self) -> &[PremultipliedColorU8] {
        let ans = bytemuck::cast_slice(self.data());
        ans
    }

    /// Consumes the internal data.
    ///
    /// Byteorder: RGBA
    pub fn take(self) -> Vec<u8> {
        self.data
    }

    /// Returns a copy of the canvas that intersects the `rect`.
    ///
    /// Returns `None` when `Canvas`'s rect doesn't contain `rect`.
    pub fn clone_rect(&self, rect: IntRect) -> Option<Canvas> {
        self.as_ref().clone_rect(rect)
    }
}

impl core::fmt::Debug for Canvas {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Canvas")
            .field("data", &"...")
            .field("width", &self.size.width())
            .field("height", &self.size.height())
            .finish()
    }
}

/// A container that references premultiplied RGBA pixels.
///
/// Can be created from `Canvas` or from a user provided data.
///
/// The data is not aligned, therefore width == stride.
#[derive(Clone, Copy, PartialEq)]
pub struct CanvasRef<'a> {
    data: &'a [u8],
    size: IntSize,
}

impl<'a> CanvasRef<'a> {
    /// Creates a new `CanvasRef` from bytes.
    ///
    /// The size must be at least `size.width() * size.height() * BYTES_PER_PIXEL`.
    /// Zero size in an error. Width is limited by i32::MAX/4.
    ///
    /// The `data` is assumed to have premultiplied RGBA pixels (byteorder: RGBA).
    pub fn from_bytes(data: &'a [u8], width: u32, height: u32) -> Option<Self> {
        let size = IntSize::from_wh(width, height)?;
        let data_len = data_len_for_size(size)?;
        if data.len() < data_len {
            return None;
        }

        Some(CanvasRef { data, size })
    }

    /// Creates a new `Canvas` from the current data.
    ///
    /// Clones the underlying data.
    pub fn to_owned(&self) -> Canvas {
        Canvas {
            data: self.data.to_vec(),
            size: self.size,
        }
    }

    /// Returns canvas's width.
    #[inline]
    pub fn width(&self) -> u32 {
        self.size.width()
    }

    /// Returns canvas's height.
    #[inline]
    pub fn height(&self) -> u32 {
        self.size.height()
    }

    /// Returns canvas's size.
    pub(crate) fn size(&self) -> IntSize {
        self.size
    }

    /// Returns canvas's rect.
    pub(crate) fn rect(&self) -> ScreenIntRect {
        self.size.to_screen_int_rect(0, 0)
    }

    /// Returns the internal data.
    ///
    /// Byteorder: RGBA
    pub fn data(&self) -> &'a [u8] {
        self.data
    }

    /// Returns a pixel color.
    ///
    /// Returns `None` when position is out of bounds.
    pub fn pixel(&self, x: u32, y: u32) -> Option<PremultipliedColorU8> {
        let idx = self.width().checked_mul(y)?.checked_add(x)?;
        self.pixels().get(idx as usize).cloned()
    }

    /// Returns a slice of pixels.
    pub fn pixels(&self) -> &'a [PremultipliedColorU8] {
        bytemuck::cast_slice(self.data())
    }

    // TODO: add rows() iterator

    /// Returns a copy of the canvas that intersects the `rect`.
    ///
    /// Returns `None` when `Canvas`'s rect doesn't contain `rect`.
    pub fn clone_rect(&self, rect: IntRect) -> Option<Canvas> {
        // TODO: to ScreenIntRect?

        let rect = self.rect().to_int_rect().intersect(&rect)?;
        let mut new = Canvas::new(rect.width(), rect.height())?;
        {
            let old_pixels = self.pixels();
            let mut new_mut = new.as_mut();
            let new_pixels = new_mut.pixels_mut();

            // TODO: optimize
            for y in 0..rect.height() {
                for x in 0..rect.width() {
                    let old_idx = (y + rect.y() as u32) * self.width() + (x + rect.x() as u32);
                    let new_idx = y * rect.width() + x;
                    new_pixels[new_idx as usize] = old_pixels[old_idx as usize];
                }
            }
        }

        Some(new)
    }

    /// Encodes canvas into a PNG data.
    pub fn to_png(&self) -> Result<Vec<u8>, ImageError> {
        // Skia uses skcms here, which is somewhat similar to RasterPipeline.

        // Sadly, we have to copy the canvas here, because of demultiplication.
        // Not sure how to avoid this.
        // TODO: remove allocation
        let image: DynamicImage = image::load_from_memory(&self.data).unwrap();
        
        Ok(image.to_bytes())
    }

    /// Saves canvas as a PNG file.
    pub fn save_png<P: AsRef<std::path::Path>>(&self, path: P) -> Result<(), ImageError> {
        self.save_image(path, image::ImageFormat::Png)
    }

    /// Saves canvas as a image file.
    pub fn save_image<P: AsRef<std::path::Path>>(&self, path: P, format: image::ImageFormat) -> Result<(), ImageError> {
        let buffer = ImageBuffer::from_fn(self.width(), self.height(), |x,y|{
            let color = self.pixel(x, y).unwrap();
            color.to_rgba()
        });
        let image: DynamicImage = DynamicImage::ImageRgba8(buffer);
        image.save_with_format(path, format)
    }
}

impl core::fmt::Debug for CanvasRef<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("CanvasRef")
            .field("data", &"...")
            .field("width", &self.size.width())
            .field("height", &self.size.height())
            .finish()
    }
}

/// A container that references mutable premultiplied RGBA pixels.
///
/// Can be created from `Canvas` or from a user provided data.
///
/// The data is not aligned, therefore width == stride.
#[derive(PartialEq)]
pub struct CanvasMut<'a> {
    data: &'a mut [u8],
    size: IntSize,
}

impl<'a> CanvasMut<'a> {
    /// Creates a new `CanvasMut` from bytes.
    ///
    /// The size must be at least `size.width() * size.height() * BYTES_PER_PIXEL`.
    /// Zero size in an error. Width is limited by i32::MAX/4.
    ///
    /// The `data` is assumed to have premultiplied RGBA pixels (byteorder: RGBA).
    pub fn from_bytes(data: &'a mut [u8], width: u32, height: u32) -> Option<Self> {
        let size = IntSize::from_wh(width, height)?;
        let data_len = data_len_for_size(size)?;
        if data.len() < data_len {
            return None;
        }

        Some(CanvasMut { data, size })
    }

    /// Creates a new `Canvas` from the current data.
    ///
    /// Clones the underlying data.
    pub fn to_owned(&self) -> Canvas {
        Canvas {
            data: self.data.to_vec(),
            size: self.size,
        }
    }

    /// Returns a container that references Canvas's data.
    pub fn as_ref(&self) -> CanvasRef {
        CanvasRef {
            data: self.data,
            size: self.size,
        }
    }

    /// Returns canvas's width.
    #[inline]
    pub fn width(&self) -> u32 {
        self.size.width()
    }

    /// Returns canvas's height.
    #[inline]
    pub fn height(&self) -> u32 {
        self.size.height()
    }

    /// Returns canvas's size.
    pub(crate) fn size(&self) -> IntSize {
        self.size
    }

    /// Fills the entire canvas with a specified color.
    pub fn fill(&mut self, color: Color) {
        let c = color.premultiply().to_color_u8();
        for p in self.pixels_mut() {
            *p = c;
        }
    }

    /// Returns the mutable internal data.
    ///
    /// Byteorder: RGBA
    pub fn data_mut(&mut self) -> &mut [u8] {
        self.data
    }

    /// Returns a mutable slice of pixels.
    pub fn pixels_mut(&mut self) -> &mut [PremultipliedColorU8] {
        bytemuck::cast_slice_mut(self.data_mut())
    }

    /// Creates `SubCanvasMut` that contains the whole `CanvasMut`.
    pub(crate) fn as_subcanvas(&mut self) -> SubCanvasMut {
        SubCanvasMut {
            size: self.size(),
            real_width: self.width() as usize,
            data: &mut self.data,
        }
    }

    /// Returns a mutable reference to the canvas region that intersects the `rect`.
    ///
    /// Returns `None` when `Canvas`'s rect doesn't contain `rect`.
    pub(crate) fn subcanvas(&mut self, rect: IntRect) -> Option<SubCanvasMut> {
        let rect = self.size.to_int_rect(0, 0).intersect(&rect)?;
        let row_bytes = self.width() as usize * BYTES_PER_PIXEL;
        let offset = rect.top() as usize * row_bytes + rect.left() as usize * BYTES_PER_PIXEL;

        Some(SubCanvasMut {
            size: rect.size(),
            real_width: self.width() as usize,
            data: &mut self.data[offset..],
        })
    }
}

impl core::fmt::Debug for CanvasMut<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("CanvasMut")
            .field("data", &"...")
            .field("width", &self.size.width())
            .field("height", &self.size.height())
            .finish()
    }
}

/// A `CanvasMut` subregion.
///
/// Unlike `CanvasMut`, contains `real_width` which references the parent `CanvasMut` width.
/// This way we can operate on a `CanvasMut` subregion without reallocations.
/// Primarily required because of `DrawTiler`.
///
/// We cannot implement it in `CanvasMut` directly, because it will brake `fill`, `data_mut`
/// `pixels_mut` and other similar methods.
/// This is because `SubCanvasMut.data` references more "data" than it actually allowed to access.
/// On the other hand, `CanvasMut.data` can access all it's data and it's stored linearly.
pub struct SubCanvasMut<'a> {
    pub data: &'a mut [u8],
    pub size: IntSize,
    pub real_width: usize,
}

impl<'a> SubCanvasMut<'a> {
    /// Returns a mutable slice of pixels.
    pub fn pixels_mut(&mut self) -> &mut [PremultipliedColorU8] {
        bytemuck::cast_slice_mut(self.data)
    }
}

/// Returns minimum bytes per row as usize.
///
/// Canvas's maximum value for row bytes must fit in 31 bits.
fn min_row_bytes(size: IntSize) -> Option<NonZeroUsize> {
    let w = i32::try_from(size.width()).ok()?;
    let w = w.checked_mul(BYTES_PER_PIXEL as i32)?;
    NonZeroUsize::new(w as usize)
}

/// Returns storage size required by pixel array.
fn compute_data_len(size: IntSize, row_bytes: usize) -> Option<usize> {
    let h = size.height().checked_sub(1)?;
    let h = (h as usize).checked_mul(row_bytes)?;

    let w = (size.width() as usize).checked_mul(BYTES_PER_PIXEL)?;

    h.checked_add(w)
}

fn data_len_for_size(size: IntSize) -> Option<usize> {
    let row_bytes = min_row_bytes(size)?;
    compute_data_len(size, row_bytes.get())
}