mspaint/src/paint.rs

use ::image::{ImageBuffer, ImageError, Rgba};
use iced::Theme;
use iced::padding;
use iced::widget::container;
use iced::widget::{Column, button, column, image, pick_list, row, text};
use iced::{Border, Color, Element, Length, Point, Renderer, Task};
use iced::{Subscription, color, event, mouse};
use std::thread;

use crate::image_button::image_button;
use crate::mouse_area::mouse_area;
use crate::mscanvas::MSCanvas;

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

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
enum BrushKind {
    #[default]
    Circle,
    Square,
    Slash,
    Backslash,
}

impl BrushKind {
    const ALL: [BrushKind; 4] = [
        BrushKind::Circle,
        BrushKind::Square,
        BrushKind::Slash,
        BrushKind::Backslash,
    ];
}

impl std::fmt::Display for BrushKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}",
            match self {
                BrushKind::Circle => "Circle",
                BrushKind::Square => "Square",
                BrushKind::Slash => "Slash",
                BrushKind::Backslash => "Backslash",
            }
        )
    }
}

#[derive(Debug, Clone, Copy)]
enum Message {
    MousePressed(Point),
    MouseReleased(Point),
    MouseMoved(Point),

    Clear,
    SavePNG,
    /// 内部消息：请求刷新图像
    RefreshImage,

    ClickTool(Tool),

    Increment(ConfigOption),
    Decrement(ConfigOption),

    BrushSelected(BrushKind),

    /// 全局鼠标释放
    WindowMouseRelease,
}

#[derive(Clone, Copy, Debug)]
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)]
pub enum ConfigOption {
    EraserWidth,
    LineWidth,
}

#[derive(Debug, Clone, Copy)]
struct Config {
    eraser_width: i32,
    line_width: i32,
}

impl Default for Config {
    fn default() -> Self {
        Self::DEFAULT.clone()
    }
}

impl Config {
    const DEFAULT: Config = Config {
        line_width: 1,
        eraser_width: 4,
    };
    fn incr(&mut self, option: ConfigOption, step: i32) {
        match option {
            ConfigOption::EraserWidth => {
                self.eraser_width += step;
                if self.eraser_width < Self::DEFAULT.eraser_width {
                    self.eraser_width = Self::DEFAULT.eraser_width;
                }
            }
            ConfigOption::LineWidth => {
                self.line_width += step;
                if self.line_width < Self::DEFAULT.line_width {
                    self.line_width = Self::DEFAULT.line_width;
                }
            }
        }
    }
}

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

    canvas: MSCanvas,

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

    begin_point: Point,

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

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

    config: Config,

    brush_selected: Option<BrushKind>,
}

impl PaintApp {
    // region iced application

    pub fn new() -> Self {
        let mut canvas = MSCanvas::new(WIDTH as i32, HEIGHT as i32);
        let (width, height) = canvas.size();
        let pixels = canvas.get_pixels();
        let config = Config::default();
        canvas.set_line_width(config.line_width);
        Self {
            tool_states: [false; Tool::Count as usize],
            tool_selected: Tool::Count,
            canvas,
            is_drawing: false,
            begin_point: Point::ORIGIN,
            image_handle: image::Handle::from_rgba(width as u32, height as u32, pixels),
            dirty: false,
            config,
            brush_selected: None,
        }
    }

    pub fn view(&self) -> Column<'_, Message> {
        // 创建显示图像的 Widget
        let (width, height) = self.canvas.size();
        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 canvas_area = container(canvas_area)
            .width(Length::Fill)
            .height(Length::Fill)
            .padding(padding::left(5).top(5).right(5))
            .style(|_| container::Style {
                background: Some(color!(0x808080).into()),
                ..container::Style::default()
            });

        let mut columns: Vec<Element<'_, Message, Theme, Renderer>> = Vec::new();
        for i in (0..(Tool::Count as usize)).step_by(2) {
            let btn1 = 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::from(i)));
            let btn2 = image_button(
                format!("image/normal/normal_{:02}.jpg", i + 2),
                format!("image/selected/selected_{:02}.jpg", i + 2),
                self.tool_states[i + 1],
            )
            .on_press(Message::ClickTool(Tool::from(i + 1)));
            columns.push(row![btn1, btn2].into());
        }
        let tool_config_area = container("").width(90).height(200).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()
            }
        });
        let tool_config_area = container(tool_config_area).padding(padding::top(5).left(5));
        columns.push(tool_config_area.into());
        let grid = Column::from_vec(columns);

        let tool_area = container(grid)
            .padding(padding::top(5).left(5).right(5).bottom(10))
            .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()
                }
            });

        let mut debug_area = row![
            column![
                text("Eraser Width "),
                row![
                    button("+").on_press(Message::Increment(ConfigOption::EraserWidth)),
                    text(self.config.eraser_width).size(20).center(),
                    button("-").on_press(Message::Decrement(ConfigOption::EraserWidth)),
                ],
            ],
            column![
                text("Line Width "),
                row![
                    button("+").on_press(Message::Increment(ConfigOption::LineWidth)),
                    text(self.config.line_width).size(20).center(),
                    button("-").on_press(Message::Decrement(ConfigOption::LineWidth)),
                ],
            ],
            pick_list(
                &BrushKind::ALL[..],
                self.brush_selected,
                Message::BrushSelected,
            )
            .placeholder("Brush..."),
        ];
        debug_area = debug_area.padding(padding::top(10).left(5).bottom(10));

        column![
            row![
                button("CLEAR").on_press(Message::Clear),
                button("PNG").on_press(Message::SavePNG),
            ],
            row![tool_area, canvas_area],
            debug_area,
        ]
    }

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

        match message {
            Message::Clear => {
                // 重置为白色
                self.canvas.clear();
                self.dirty = true;
            }
            Message::SavePNG => {
                let (width, height) = self.canvas.size();
                let pixels = self.canvas.get_pixels();
                let path = "mspaint.png";
                save_pixels_async(pixels, width as u32, height as u32, 4, path.to_string());
            }
            Message::RefreshImage => {
                if self.dirty {
                    self.update_image_handle();
                    self.dirty = false;
                }
            }
            Message::ClickTool(tool) => {
                self.update_tool_states(tool);
            }
            Message::Increment(opt) => {
                self.config.incr(opt, 1);
                if opt == ConfigOption::LineWidth {
                    self.canvas.set_line_width(self.config.line_width);
                }
            }
            Message::Decrement(opt) => {
                self.config.incr(opt, -1);
                if opt == ConfigOption::LineWidth {
                    self.canvas.set_line_width(self.config.line_width);
                }
            }
            Message::BrushSelected(kind) => {
                self.brush_selected = Some(kind);
            }
            Message::WindowMouseRelease => {
                // 处理鼠标在 canvas_area 外面释放
                self.is_drawing = false;
            }
            _ => {}
        }

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

        Task::none()
    }

    pub fn subscription(&self) -> Subscription<Message> {
        event::listen().filter_map(|event| match event {
            // 只关心鼠标释放事件
            iced::Event::Mouse(mouse::Event::ButtonReleased(mouse::Button::Left)) => {
                Some(Message::WindowMouseRelease)
            }
            _ => None,
        })
    }

    // endregion

    // region tool update

    pub fn update_with_eraser(&mut self, message: Message) {
        match message {
            Message::MousePressed(pos) => {
                self.is_drawing = true;
                self.canvas.clear_rect(
                    (pos.x as i32) - self.config.eraser_width / 2,
                    (pos.y as i32) - self.config.eraser_width / 2,
                    self.config.eraser_width,
                    self.config.eraser_width,
                );
                self.dirty = true;
            }
            Message::MouseReleased(_) => {
                self.is_drawing = false;
            }
            Message::MouseMoved(pos) => {
                if self.is_drawing {
                    self.canvas.clear_rect(
                        (pos.x as i32) - self.config.eraser_width / 2,
                        (pos.y as i32) - self.config.eraser_width / 2,
                        self.config.eraser_width,
                        self.config.eraser_width,
                    );
                    self.dirty = true;
                }
            }
            _ => {}
        }
    }

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

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

    pub fn update_with_brush(&mut self, message: Message) {
        let mut brush_fn: fn(&mut MSCanvas, Point) = MSCanvas::brush_circle;
        if let Some(kind) = self.brush_selected {
            brush_fn = match kind {
                BrushKind::Circle => MSCanvas::brush_circle,
                BrushKind::Square => MSCanvas::brush_square,
                BrushKind::Slash => MSCanvas::brush_slash,
                BrushKind::Backslash => MSCanvas::brush_backslash,
            }
        }
        match message {
            Message::MousePressed(pos) => {
                self.is_drawing = true;
                brush_fn(&mut self.canvas, pos);
                self.begin_point = pos;
                self.dirty = true
            }
            Message::MouseReleased(pos) => {
                self.is_drawing = false;
                self.begin_point = pos;
            }
            Message::MouseMoved(pos) => {
                if self.is_drawing {
                    self.canvas.draw_line_with(self.begin_point, pos, brush_fn);
                    self.begin_point = pos;
                    self.dirty = true;
                }
            }
            _ => {}
        }
    }

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

    pub fn update_with_rectangle(&mut self, message: Message) {
        match message {
            Message::MousePressed(pos) => {
                self.is_drawing = true;
                self.canvas.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.canvas.restore_pixels();
                    self.canvas.stroke_rect1(self.begin_point, pos);
                    self.dirty = true;
                }
            }
            _ => {}
        }
    }

    // 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.canvas.get_pixels();

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

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

fn save_rgba_to_png(
    rgba_data: Vec<u8>,
    width: u32,
    height: u32,
    path: &str,
) -> Result<(), ImageError> {
    // Each pixel is 4 bytes (R, G, B, A)
    let img: ImageBuffer<Rgba<u8>, Vec<u8>> = ImageBuffer::from_raw(width, height, rgba_data)
        .expect("Failed to create ImageBuffer: data size mismatch");

    img.save(path)?;
    Ok(())
}

fn scale_pixels(pixels: Vec<u8>, width: u32, height: u32, scale: u32) -> Vec<u8> {
    if scale <= 1 {
        return pixels;
    }
    let dst_width = width * scale;
    let dst_height = height * scale;
    let mut dst = vec![0; (dst_width * dst_height * 4) as usize]; // RGBA

    for y in 0..height {
        for x in 0..width {
            // 源像素索引
            let src_idx = ((y * width + x) * 4) as usize;

            // 源像素颜色
            let r = pixels[src_idx];
            let g = pixels[src_idx + 1];
            let b = pixels[src_idx + 2];
            let a = pixels[src_idx + 3];

            // 在目标图像中填充 scale×scale 区域
            for dy in 0..scale {
                for dx in 0..scale {
                    let dst_x = x * scale + dx;
                    let dst_y = y * scale + dy;
                    let dst_idx = ((dst_y * dst_width + dst_x) * 4) as usize;

                    dst[dst_idx] = r;
                    dst[dst_idx + 1] = g;
                    dst[dst_idx + 2] = b;
                    dst[dst_idx + 3] = a;
                }
            }
        }
    }

    dst
}

fn save_pixels_async(pixels: Vec<u8>, width: u32, height: u32, scale: u32, path: String) {
    // 用 Arc 共享数据（或直接 move 进去）
    thread::spawn(move || {
        let pixels = scale_pixels(pixels, width, height, scale);
        match save_rgba_to_png(pixels, width, height, &path) {
            Ok(()) => println!("✅ Image saved to {}", path),
            Err(e) => eprintln!("❌ Failed to save image: {}", e),
        }
    });
}