mspaint/src/paint.rs

use iced::Theme;
use iced::padding;
use iced::widget::container;
use iced::widget::{Column, Grid, button, column, image, row};
use iced::{Border, Color, Length, Point, Task};

use crate::image_button::image_button;
use crate::mouse_area::mouse_area;

const WIDTH: u32 = 800;
const HEIGHT: u32 = 600;

#[derive(Debug, Clone, Copy)]
pub enum Message {
    MousePressed(Point),
    MouseReleased(Point),
    MouseMoved(Point),
    Clear,
    // 内部消息：请求刷新图像
    RefreshImage,

    ClickTool(Tool),
}

#[derive(Clone, Copy, Debug)]
pub enum Tool {
    FreeFormSelect,
    Select,
    Eraser,
    FillWithColor,
    PickColor,
    Magnifier,
    Pencil,
    Brush,
    Airbrush,
    Text,
    Line,
    Curve,
    Rectangle,
    Polygon,
    Ellipse,
    RoundedRectangle,

    Count,
}

impl From<usize> for Tool {
    fn from(tool: usize) -> Self {
        match tool {
            0 => Self::FreeFormSelect,
            1 => Self::Select,
            2 => Self::Eraser,
            3 => Self::FillWithColor,
            4 => Self::PickColor,
            5 => Self::Magnifier,
            6 => Self::Pencil,
            7 => Self::Brush,
            8 => Self::Airbrush,
            9 => Self::Text,
            10 => Self::Line,
            11 => Self::Curve,
            12 => Self::Rectangle,
            13 => Self::Polygon,
            14 => Self::Ellipse,
            15 => Self::RoundedRectangle,
            _ => Self::Count,
        }
    }
}

impl From<Tool> for usize {
    fn from(tool: Tool) -> usize {
        match tool {
            Tool::FreeFormSelect => 0,
            Tool::Select => 1,
            Tool::Eraser => 2,
            Tool::FillWithColor => 3,
            Tool::PickColor => 4,
            Tool::Magnifier => 5,
            Tool::Pencil => 6,
            Tool::Brush => 7,
            Tool::Airbrush => 8,
            Tool::Text => 9,
            Tool::Line => 10,
            Tool::Curve => 11,
            Tool::Rectangle => 12,
            Tool::Polygon => 13,
            Tool::Ellipse => 14,
            Tool::RoundedRectangle => 15,
            Tool::Count => 16,
        }
    }
}

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

impl ColorU8 {
    fn new(r: u8, g: u8, b: u8, a: u8) -> Self {
        Self { r, g, b, a }
    }

    fn is_same(&self, color: Color) -> bool {
        self.r == (color.r * 255.0) as u8
            && self.g == (color.g * 255.0) as u8
            && self.b == (color.b * 255.0) as u8
            && self.a == (color.a * 255.0) as u8
    }
}

struct Paint {
    tool_states: [bool; Tool::Count as usize],
    tool_selected: Tool,

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

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

    // 是否正在绘制
    is_drawing: bool,

    begin_point: Point,

    // 用于显示的图像句柄缓存
    // 每次像素变化后需要重新生成
    image_handle: image::Handle,

    // brush 大小
    brush_radius: i32,

    // 标记像素是否被修改，用于优化图像句柄的生成
    dirty: bool,
}

impl Paint {
    // region iced application (new view update)

    pub fn new() -> Self {
        // 初始化全白背景 (R=255, G=255, B=255, A=255)
        let pixels = vec![255u8; (WIDTH * HEIGHT * 4) as usize];
        let data = pixels.clone();

        Self {
            tool_states: [false; Tool::Count as usize],
            tool_selected: Tool::Count,
            pixels,
            pixels_bak: Vec::new(),
            color: Color::BLACK,
            is_drawing: false,
            begin_point: Point::ORIGIN,
            image_handle: image::Handle::from_rgba(WIDTH, HEIGHT, data),
            brush_radius: 1,
            dirty: false,
        }
    }

    pub fn view(&self) -> Column<'_, Message> {
        // 创建显示图像的 Widget
        // 如果 handle 还没准备好，显示一个占位符
        let image_widget = image(self.image_handle.clone())
            .width(Length::Fixed(WIDTH as f32))
            .height(Length::Fixed(HEIGHT as f32));

        let canvas_area = mouse_area(image_widget)
            .on_press(|pos| Message::MousePressed(pos)) // 占位，实际逻辑在 on_drag 或自定义
            .on_release(|pos| Message::MouseReleased(pos))
            .on_move(|pos| Message::MouseMoved(pos));
        // 注意：mouse_area 的 on_move 给出的坐标通常是相对于 widget 左上角的，这正是我们需要的！

        let mut grid = Grid::new();
        grid = grid.columns(2).width(100);

        for i in 0..(Tool::Count as usize) {
            let tool = Tool::from(i);
            let btn = image_button(
                format!("image/normal/normal_{:02}.jpg", i + 1),
                format!("image/selected/selected_{:02}.jpg", i + 1),
                self.tool_states[i],
            )
            .on_press(Message::ClickTool(tool));
            grid = grid.push(btn);
        }

        let tool_area = container(grid)
            .padding(padding::top(5).left(5).right(5).bottom(100))
            .style(|theme: &Theme| {
                let palette = theme.extended_palette();

                container::Style {
                    background: Some(Color::from_rgb8(192, 192, 192).into()),
                    border: Border {
                        width: 1.0,
                        radius: 5.0.into(),
                        color: palette.background.weak.color,
                    },
                    ..container::Style::default()
                }
            });

        // We use a column: a simple vertical layout
        column![
            button("CLEAR").on_press(Message::Clear),
            row![tool_area, canvas_area],
        ]
    }

    pub fn update(&mut self, message: Message) -> Task<Message> {
        match self.tool_selected {
            Tool::FillWithColor => {
                self.update_with_fill_with_color(message);
            }
            Tool::Pencil => {
                self.update_with_pencil(message);
            }
            Tool::Line => {
                self.update_with_line(message);
            }
            _ => {}
        }

        match message {
            Message::Clear => {
                // 重置为白色
                self.pixels.fill(255);
                self.dirty = true;
            }
            Message::RefreshImage => {
                if self.dirty {
                    self.update_image_handle();
                    self.dirty = false;
                }
            }
            Message::ClickTool(tool) => {
                self.update_tool_states(tool);
            }
            _ => {}
        }

        // 如果像素被修改了，我们需要触发一次 RefreshImage 来更新 UI
        // 在实际复杂应用中，可能需要防抖或异步处理，这里为了实时性直接同步触发
        if self.dirty {
            // 像素变了，安排下一帧刷新图像句柄
            // 注意：频繁生成 Handle 可能消耗 CPU，生产环境建议加节流
            return Task::perform(async { Message::RefreshImage }, |msg| msg);
        }

        Task::none()
    }

    // endregion

    // region tool update

    pub fn update_with_fill_with_color(&mut self, message: Message) {
        match message {
            Message::MousePressed(pos) => {
                self.fill_scanline(pos);
            }
            _ => {}
        }
    }

    pub fn update_with_pencil(&mut self, message: Message) {
        match message {
            Message::MousePressed(pos) => {
                println!("pressed: {:?}", pos);
                self.is_drawing = true;
                self.draw_pixel_at1(pos);
                self.begin_point = pos;
            }
            Message::MouseReleased(pos) => {
                self.is_drawing = false;
                self.begin_point = pos;
            }
            Message::MouseMoved(pos) => {
                if self.is_drawing {
                    self.draw_line(self.begin_point, pos);
                    self.begin_point = pos;
                }
            }
            _ => {}
        }
    }

    pub fn update_with_line(&mut self, message: Message) {
        match message {
            Message::MousePressed(pos) => {
                self.is_drawing = true;
                self.save_pixels();
                self.begin_point = pos;
            }
            Message::MouseReleased(pos) => {
                self.is_drawing = false;
                self.begin_point = pos;
            }
            Message::MouseMoved(pos) => {
                if self.is_drawing {
                    self.restore_pixels();
                    self.draw_line(self.begin_point, pos);
                }
            }
            _ => {}
        }
    }

    // endregion

    pub fn update_tool_states(&mut self, tool: Tool) {
        let idx = tool as usize;
        if idx >= self.tool_states.len() {
            return;
        }
        let old_value = self.tool_states[idx];
        for i in 0..(Tool::Count as usize) {
            self.tool_states[i] = false;
        }
        self.tool_states[idx] = !old_value;
        self.tool_selected = idx.into();
    }

    /// 将原始字节转换为 Iced 的图像句柄
    fn update_image_handle(&mut self) {
        // 克隆数据以避免所有权问题，或者使用 Arc 如果数据量大
        // 这里为了简单直接 clone，对于 800x600 (约 2MB) 来说很快
        let data = self.pixels.clone();

        self.image_handle = image::Handle::from_rgba(WIDTH, HEIGHT, data);
    }
}

/// draw method
#[allow(unused)]
impl Paint {
    fn pixel_at(&self, x: i32, y: i32) -> ColorU8 {
        // 边界检查
        if x < 0 || x >= WIDTH as i32 || y < 0 || y >= HEIGHT as i32 {
            return ColorU8::new(0, 0, 0, 0);
        }

        let x = x as u32;
        let y = y as u32;

        let index = ((y * WIDTH + x) * 4) as usize;

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

    fn pixel_at_raw(&self, x: i32, y: i32) -> ColorU8 {
        let x = x as u32;
        let y = y as u32;
        let index = ((y * WIDTH + x) * 4) as usize;
        ColorU8::new(
            self.pixels[index],
            self.pixels[index + 1],
            self.pixels[index + 2],
            self.pixels[index + 3],
        )
    }

    fn pixel_at1(&self, point: Point) -> ColorU8 {
        self.pixel_at(point.x as i32, point.y as i32)
    }

    fn draw_brush_at(&mut self, center_x: i32, center_y: i32) {
        let r = self.brush_radius;
        for dy in -r..=r {
            for dx in -r..=r {
                if dx * dx + dy * dy <= r * r {
                    self.draw_pixel_at(center_x + dx, center_y + dy);
                }
            }
        }
    }

    fn draw_pixel_at_raw(&mut self, x: u32, y: u32) {
        // 计算索引：(y * width + x) * 4
        let index = ((y * WIDTH + x) * 4) as usize;

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

        self.dirty = true;
    }

    /// 核心绘图逻辑：直接在字节数组上操作
    fn draw_pixel_at(&mut self, x: i32, y: i32) {
        // 边界检查
        if x < 0 || x >= WIDTH as i32 || y < 0 || y >= HEIGHT as i32 {
            return;
        }
        self.draw_pixel_at_raw(x as u32, y as u32);
    }

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

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

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

        self.dirty = true;
    }

    fn draw_pixel_at1(&mut self, pos: Point) {
        self.draw_pixel_at(pos.x as i32, pos.y as i32)
    }

    fn draw_lines(&mut self, points: &[Point]) {
        if points.is_empty() {
            return;
        }
        if points.len() == 1 {
            self.draw_pixel_at1(points[0]);
            return;
        }

        let mut begin = points[0];
        for point in points.iter().skip(1) {
            self.draw_line(begin, point.clone());
            begin = point.clone();
        }
    }

    /// Bresenham's line drawing algorithm
    fn draw_line(&mut self, begin: Point, end: Point) {
        let x1 = begin.x as i32;
        let y1 = begin.y as i32;
        let x2 = end.x as i32;
        let y2 = end.y as i32;

        // draw start end point, 防止多条线段在连接点出现断开（比如 ab bc）
        self.draw_pixel_at(x1, y1);
        self.draw_pixel_at(x2, y2);

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

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

        if dy1 <= dx1 {
            if dx >= 0 {
                x = x1;
                y = y1;
                xe = x2;
            } else {
                x = x2;
                y = y2;
                xe = x1;
            }
            self.draw_pixel_at(x, y);
            while x < xe {
                x += 1;
                if px < 0 {
                    px = px + 2 * dy1;
                } else {
                    if (dx < 0 && dy < 0) || (dx > 0 && dy > 0) {
                        y = y + 1;
                    } else {
                        y = y - 1;
                    }
                    px = px + 2 * (dy1 - dx1);
                }
                self.draw_pixel_at(x, y);
            }
        } else {
            if dy >= 0 {
                x = x1;
                y = y1;
                ye = y2;
            } else {
                x = x2;
                y = y2;
                ye = y1;
            }
            self.draw_pixel_at(x, y);
            while y < ye {
                y = y + 1;
                if py <= 0 {
                    py = py + 2 * dx1;
                } else {
                    if (dx < 0 && dy < 0) || (dx > 0 && dy > 0) {
                        x = x + 1;
                    } else {
                        x = x - 1;
                    }
                    py = py + 2 * (dx1 - dy1);
                }
                self.draw_pixel_at(x, y);
            }
        }
    }

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

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

    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.is_same(self.color) {
            return (0, 0);
        }

        let mut scan_points = vec![(start_x, start_y)];
        let width = WIDTH as i32;
        let height = HEIGHT as i32;
        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 >= width || y < 0 || y >= height {
                continue;
            }
            if self.pixel_at_raw(x, y) == target_color {
                self.draw_pixel_at_raw(x as u32, y as u32);
                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)
    }

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

        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.is_same(self.color) {
            return (0, 0);
        }

        let mut stack = vec![(start_x, start_y)];
        let mut iter_count = 0;
        let mut fill_count = 0;

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

            if x < 0 || x >= width || y < 0 || y >= height {
                continue;
            }

            if self.pixel_at_raw(x, y) != target_color {
                continue;
            }

            self.draw_pixel_at_raw(x as u32, y as u32);
            fill_count += 1;

            // 提前检查边界再入栈
            if x > 0 {
                stack.push((x - 1, y));
            }
            if x + 1 < width {
                stack.push((x + 1, y));
            }
            if y > 0 {
                stack.push((x, y - 1));
            }
            if y + 1 < height {
                stack.push((x, y + 1));
            }
        }

        (iter_count, fill_count)
    }

    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 = WIDTH as i32;
        let height = HEIGHT as i32;

        // 边界检查
        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.is_same(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_raw(lx - 1, y) == target_color {
                lx -= 1;
            }
            // 向右扩展 rx
            while rx + 1 < width && self.pixel_at_raw(rx + 1, y) == target_color {
                rx += 1;
            }

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

            // 检查上一行 (y - 1)
            if y - 1 >= 0 {
                let mut x = lx;
                while x <= rx {
                    if self.pixel_at_raw(x, y - 1) == target_color {
                        let span_start = x;
                        // 跳过连续的目标色块
                        while x <= rx && self.pixel_at_raw(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_raw(x, y + 1) == target_color {
                        let span_start = x;
                        while x <= rx && self.pixel_at_raw(x, y + 1) == target_color {
                            x += 1;
                        }
                        stack.push((y + 1, span_start, x - 1));
                    } else {
                        x += 1;
                    }
                }
            }
        }

        (iter_count, fill_count)
    }
}

pub fn main() -> iced::Result {
    iced::application(Paint::new, Paint::update, Paint::view)
        .theme(Theme::CatppuccinMocha)
        .run()
}