mspaint/src/mscanvas.rs

use iced::Point;

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct MSColor {
    r: u8,
    g: u8,
    b: u8,
    a: u8,
}

#[allow(unused)]
impl MSColor {
    fn new(r: u8, g: u8, b: u8, a: u8) -> Self {
        Self { r, g, b, a }
    }

    pub const ZERO: MSColor = MSColor {
        r: 0,
        g: 0,
        b: 0,
        a: 0,
    };

    pub const BLACK: MSColor = MSColor {
        r: 0,
        g: 0,
        b: 0,
        a: 255,
    };

    pub const WHITE: MSColor = MSColor {
        r: 255,
        g: 255,
        b: 255,
        a: 255,
    };
}

pub struct MSCanvas {
    width: i32,
    height: i32,

    /// 原始像素数据：RGBA 格式
    /// 长度 = width * height * 4
    pixels: Vec<u8>,
    pixels_bak: Vec<u8>,

    /// 当前笔画颜色
    color: MSColor,

    line_width: i32,
}

#[allow(unused)]
impl MSCanvas {
    pub fn new(width: i32, height: i32) -> Self {
        Self {
            width,
            height,
            pixels: vec![255; (width * height * 4) as usize],
            pixels_bak: Vec::new(),
            color: MSColor::BLACK,
            line_width: 1,
        }
    }

    pub fn set_color(&mut self, color: MSColor) {
        self.color = color;
    }

    pub fn set_line_width(&mut self, line_width: i32) {
        self.line_width = line_width;
    }
}

#[allow(unused)]
impl MSCanvas {
    pub fn size(&self) -> (i32, i32) {
        (self.width, self.height)
    }

    pub fn get_pixels(&self) -> Vec<u8> {
        self.pixels.clone()
    }

    pub fn pixel_at(&self, x: i32, y: i32) -> MSColor {
        // 边界检查
        if x < 0 || x >= self.width || y < 0 || y >= self.height as i32 {
            return MSColor::ZERO;
        }

        let index = ((y * self.width + x) * 4) as usize;

        MSColor::new(
            self.pixels[index],
            self.pixels[index + 1],
            self.pixels[index + 2],
            self.pixels[index + 3],
        )
    }

    pub fn draw_pixel_at(&mut self, point: Point) {
        let Point { x, y } = point;
        let x = x as i32;
        let y = y as i32;
        // 边界检查
        if x < 0 || x >= self.width || y < 0 || y >= self.height {
            return;
        }
        // 计算索引：(y * width + x) * 4
        let index = ((y * self.width + x) * 4) as usize;

        // 写入 RGBA 数据
        // 注意：Color 的 r, g, b, a 是 0.0 - 1.0，需要转为 0 - 255
        self.pixels[index] = self.color.r; // R
        self.pixels[index + 1] = self.color.g; // G
        self.pixels[index + 2] = self.color.b; // B
        self.pixels[index + 3] = self.color.a; // A
    }

    fn draw_pixel_at1(&mut self, x: i32, y: i32) {
        // 边界检查
        if x < 0 || x >= self.width || y < 0 || y >= self.height {
            return;
        }
        // 计算索引：(y * width + x) * 4
        let index = ((y * self.width + x) * 4) as usize;

        // 写入 RGBA 数据
        // 注意：Color 的 r, g, b, a 是 0.0 - 1.0，需要转为 0 - 255
        self.pixels[index] = self.color.r; // R
        self.pixels[index + 1] = self.color.g; // G
        self.pixels[index + 2] = self.color.b; // B
        self.pixels[index + 3] = self.color.a; // A
    }

    pub fn draw_pixel_row(&mut self, xs: i32, xe: i32, y: i32) {
        if y < 0 || y >= self.height {
            return;
        }

        let xs = xs.clamp(0, self.width - 1);
        let xe = xe.clamp(0, self.width as i32);
        let y = y;
        for x in xs..xe {
            let index = ((y * self.width + x) * 4) as usize;

            // 写入 RGBA 数据
            // 注意：Color 的 r, g, b, a 是 0.0 - 1.0，需要转为 0 - 255
            self.pixels[index] = self.color.r; // R
            self.pixels[index + 1] = self.color.g; // G
            self.pixels[index + 2] = self.color.b; // B
            self.pixels[index + 3] = self.color.a; // A
        }
    }

    pub fn draw_pixels(&mut self, points: Vec<Point>) {
        for point in points {
            self.draw_pixel_at(point);
        }
    }

    pub fn brush_circle(&mut self, center: Point) {
        if self.line_width <= 1 {
            self.draw_pixel_at(center);
            return;
        }
        let square = (self.line_width as f32) * (self.line_width as f32) / 4.0;
        let l = self.line_width / 2;
        let r = self.line_width - l;
        for dy in -l..r {
            for dx in -l..r {
                if (dx * dx + dy * dy) as f32 <= square {
                    self.draw_pixel_at(Point::new(center.x + dx as f32, center.y + dy as f32));
                }
            }
        }
    }

    pub fn brush_square(&mut self, center: Point) {
        if self.line_width <= 1 {
            self.draw_pixel_at(center);
            return;
        }
        let l = self.line_width / 2;
        let r = self.line_width - l;
        for dy in -l..r {
            for dx in -l..r {
                self.draw_pixel_at(Point::new(center.x + dx as f32, center.y + dy as f32));
            }
        }
    }

    fn bresenham_line(&mut self, begin: Point, end: Point) -> Vec<Point> {
        let x1 = begin.x;
        let y1 = begin.y;
        let x2 = end.x;
        let y2 = end.y;

        let dx = (x2 - x1);
        let dy = (y2 - y1);
        let dx1 = dx.abs();
        let dy1 = dy.abs();
        let mut px = 2.0 * dy1 - dx1;
        let mut py = 2.0 * dx1 - dy1;

        let mut x;
        let mut y;
        let xe;
        let ye;

        let mut points = Vec::new();
        if dy1 <= dx1 {
            if dx >= 0.0 {
                x = x1;
                y = y1;
                xe = x2;
                ye = y2;
            } else {
                x = x2;
                y = y2;
                xe = x1;
                ye = y1;
            }
            points.push(Point::new(x, y));
            while x < xe {
                x += 1.0;
                if px < 0.0 {
                    px = px + 2.0 * dy1;
                } else {
                    if (dx < 0.0 && dy < 0.0) || (dx > 0.0 && dy > 0.0) {
                        y = y + 1.0;
                    } else {
                        y = y - 1.0;
                    }
                    px = px + 2.0 * (dy1 - dx1);
                }
                points.push(Point::new(x, y));
            }
        } else {
            if dy >= 0.0 {
                x = x1;
                y = y1;
                ye = y2;
                xe = x2;
            } else {
                x = x2;
                y = y2;
                ye = y1;
                xe = x1;
            }
            points.push(Point::new(x, y));
            while y < ye {
                y = y + 1.0;
                if py <= 0.0 {
                    py = py + 2.0 * dx1;
                } else {
                    if (dx < 0.0 && dy < 0.0) || (dx > 0.0 && dy > 0.0) {
                        x = x + 1.0;
                    } else {
                        x = x - 1.0;
                    }
                    py = py + 2.0 * (dx1 - dy1);
                }
                points.push(Point::new(x, y));
            }
        }
        points.push(Point::new(xe, ye));

        points
    }

    pub fn draw_line(&mut self, begin: Point, end: Point) {
        let points = self.bresenham_line(begin, end);
        for point in points {
            self.draw_pixel_at(point);
        }
    }

    pub fn draw_line_with_circle_brush(&mut self, begin: Point, end: Point) {
        // let points = self.bresenham_line(begin, end);
        // for point in points {
        //     self.draw_brush_at(point);
        // }
        self.draw_line_with(begin, end, MSCanvas::brush_circle);
    }

    pub fn draw_line_with_square_brush(&mut self, begin: Point, end: Point) {
        self.draw_line_with(begin, end, MSCanvas::brush_square);
    }

    pub fn draw_line_with<F>(&mut self, begin: Point, end: Point, brush_fn: F)
    where
        F: Fn(&mut MSCanvas, Point) -> (),
    {
        let points = self.bresenham_line(begin, end);
        for point in points {
            brush_fn(self, point);
        }
    }

    pub fn save_pixels(&mut self) {
        self.pixels_bak = self.pixels.clone();
    }

    pub fn restore_pixels(&mut self) {
        self.pixels = self.pixels_bak.clone();
    }

    pub fn fill_slow(&mut self, begin: Point) -> (i32, i32) {
        let start_x = begin.x as i32;
        let start_y = begin.y as i32;
        let target_color = self.pixel_at(start_x, start_y);
        if target_color == self.color {
            return (0, 0);
        }

        let mut scan_points = vec![(start_x, start_y)];
        let width = self.width;
        let height = self.height;
        let mut iter_count = 0;
        let mut fill_count = 0;
        while let Some((x, y)) = scan_points.pop() {
            iter_count += 1;

            if x < 0 || x >= self.width || y < 0 || y >= self.height {
                continue;
            }
            if self.pixel_at(x, y) == target_color {
                self.draw_pixel_at1(x, y);
                fill_count += 1;

                let p1 = (x - 1, y);
                let p2 = (x + 1, y);
                let p3 = (x, y - 1);
                let p4 = (x, y + 1);
                scan_points.push(p1);
                scan_points.push(p2);
                scan_points.push(p3);
                scan_points.push(p4);
            }
        }
        (iter_count, fill_count)
    }

    pub fn fill_scanline(&mut self, begin: Point) -> (i32, i32) {
        let start_x = begin.x as i32;
        let start_y = begin.y as i32;
        let width = self.width;
        let height = self.height;

        // 边界检查
        if start_x < 0 || start_x >= width || start_y < 0 || start_y >= height {
            return (0, 0);
        }

        let target_color = self.pixel_at(start_x, start_y);
        if target_color == (self.color) {
            return (0, 0);
        }

        // 栈中存储 (y, x1, x2)：表示第 y 行从 x1 到 x2（含）需要向上/下扫描
        let mut stack = vec![(start_y, start_x, start_x)];
        let mut iter_count = 0;
        let mut fill_count = 0;

        while let Some((y, mut lx, mut rx)) = stack.pop() {
            iter_count += 1;

            // 向左扩展 lx
            while lx - 1 >= 0 && self.pixel_at(lx - 1, y) == target_color {
                lx -= 1;
            }
            // 向右扩展 rx
            while rx + 1 < width && self.pixel_at(rx + 1, y) == target_color {
                rx += 1;
            }

            // 填充当前行 [lx, rx]
            for x in lx..=rx {
                self.draw_pixel_at1(x, y);
                fill_count += 1;
            }

            // 检查上一行 (y - 1)
            if y - 1 >= 0 {
                let mut x = lx;
                while x <= rx {
                    if self.pixel_at(x, y - 1) == target_color {
                        let span_start = x;
                        // 跳过连续的目标色块
                        while x <= rx && self.pixel_at(x, y - 1) == target_color {
                            x += 1;
                        }
                        // 将这个 span 入栈（用于后续处理上一行的上一行）
                        stack.push((y - 1, span_start, x - 1));
                    } else {
                        x += 1;
                    }
                }
            }

            // 检查下一行 (y + 1)
            if y + 1 < height {
                let mut x = lx;
                while x <= rx {
                    if self.pixel_at(x, y + 1) == target_color {
                        let span_start = x;
                        while x <= rx && self.pixel_at(x, y + 1) == target_color {
                            x += 1;
                        }
                        stack.push((y + 1, span_start, x - 1));
                    } else {
                        x += 1;
                    }
                }
            }
        }

        (iter_count, fill_count)
    }

    pub fn clear_rect(&mut self, x: i32, y: i32, width: i32, height: i32) {
        for yi in y..(y + height) {
            self.clear_row(x, x + width, yi);
        }
    }

    pub fn fill_rect(&mut self, x: i32, y: i32, width: i32, height: i32) {
        for yi in y..(y + height) {
            self.draw_pixel_row(x, x + width, yi);
        }
    }

    pub fn stroke_rect(&mut self, x: i32, y: i32, width: i32, height: i32) {
        self.fill_rect(x, y, width, self.line_width);
        self.fill_rect(x, y + height - self.line_width, width, self.line_width);
        self.fill_rect(x, y, self.line_width, height);
        self.fill_rect(x + width - self.line_width, y, self.line_width, height);
    }

    pub fn stroke_rect1(&mut self, p1: Point, p2: Point) {
        let mut x = p1.x;
        let mut y = p1.y;
        let mut width = (p2.x - p1.x);
        let mut height = (p2.y - p1.y);
        if width < 0.0 && height < 0.0 {
            x = p2.x;
            y = p2.y;
            width = (p1.x - p2.x);
            height = (p1.y - p2.y);
        } else if width < 0.0 {
            x += width;
            width = -width;
        } else if height < 0.0 {
            y += height;
            height = -height;
        }
        self.stroke_rect(x as i32, y as i32, width as i32, height as i32);
    }

    pub fn clear(&mut self) {
        self.pixels.fill(255);
    }

    pub fn clear_row(&mut self, xs: i32, xe: i32, y: i32) {
        if y < 0 || y >= self.height {
            return;
        }

        let xs = xs.clamp(0, self.width - 1);
        let xe = xe.clamp(0, self.width as i32);
        for x in xs..xe {
            let index = ((y * self.width + x) * 4) as usize;

            // 写入 RGBA 数据
            // 注意：Color 的 r, g, b, a 是 0.0 - 1.0，需要转为 0 - 255
            self.pixels[index] = 255; // R
            self.pixels[index + 1] = 255; // G
            self.pixels[index + 2] = 255; // B
            self.pixels[index + 3] = 255; // A
        }
    }
}