Grim_Rune_Terrain_Generator/src/app.rs

use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::thread;

use arboard::{Clipboard, ImageData};
use eframe::egui::{
    self, Align2, Color32, ColorImage, Context, CornerRadius, DragValue, Event, FontId, Frame, Key,
    Response, RichText, Sense, Slider, Stroke, StrokeKind, TextureHandle, TextureOptions, Ui, Vec2,
};
use rand::Rng;

use crate::generator::{self, GeneratedImages, GenerationParams};

#[derive(Copy, Clone, Eq, PartialEq)]
enum LeftTab {
    Generation,
    Advanced,
}

#[derive(Copy, Clone, Eq, PartialEq)]
enum PreviewFocus {
    TerrainComposite,
    TerrainColor,
    Heightmap,
    Bumpmap,
}

#[derive(Clone)]
struct TextureSet {
    terrain_composite: TextureHandle,
    terrain_color: TextureHandle,
    heightmap: TextureHandle,
    bumpmap: TextureHandle,
}

pub struct TerrainApp {
    params: GenerationParams,
    generation_tab: LeftTab,
    preview_focus: PreviewFocus,
    hovered_preview: Option<PreviewFocus>,
    textures: Option<TextureSet>,
    generated: Option<GeneratedImages>,
    generation_receiver: Option<Receiver<Result<GeneratedImages, String>>>,
    generation_in_progress: bool,
    copy_receiver: Option<Receiver<Result<String, String>>>,
    copy_in_progress: bool,
    status_message: String,
}

impl TerrainApp {
    pub fn new(_cc: &eframe::CreationContext<'_>) -> Self {
        Self {
            params: GenerationParams::default(),
            generation_tab: LeftTab::Generation,
            preview_focus: PreviewFocus::TerrainComposite,
            hovered_preview: None,
            textures: None,
            generated: None,
            generation_receiver: None,
            generation_in_progress: false,
            copy_receiver: None,
            copy_in_progress: false,
            status_message: String::from("Adjust parameters and click Generate."),
        }
    }

    fn start_generation(&mut self) {
        if self.generation_in_progress {
            return;
        }

        let params = self.params.clone();
        let (tx, rx) = mpsc::channel();
        self.generation_receiver = Some(rx);
        self.generation_in_progress = true;
        self.status_message = format!(
            "Generating {}x{} terrain...",
            params.image_width, params.image_height
        );

        thread::spawn(move || {
            let result = std::panic::catch_unwind(|| generator::generate(&params))
                .map_err(|_| String::from("Generation panicked unexpectedly."));
            let _ = tx.send(result);
        });
    }

    fn poll_generation(&mut self, ctx: &Context) {
        let Some(receiver) = &self.generation_receiver else {
            return;
        };

        match receiver.try_recv() {
            Ok(Ok(images)) => {
                self.status_message = format!(
                    "Generated {}x{} heightmap, bumpmap, terrain color, and composite.",
                    images.width, images.height
                );
                self.install_textures(ctx, &images);
                self.generated = Some(images);
                self.generation_receiver = None;
                self.generation_in_progress = false;
            }
            Ok(Err(error)) => {
                self.status_message = error;
                self.generation_receiver = None;
                self.generation_in_progress = false;
            }
            Err(TryRecvError::Empty) => {
                ctx.request_repaint();
            }
            Err(TryRecvError::Disconnected) => {
                self.status_message = String::from("Generation worker disconnected.");
                self.generation_receiver = None;
                self.generation_in_progress = false;
            }
        }
    }

    fn poll_copy(&mut self, ctx: &Context) {
        let Some(receiver) = &self.copy_receiver else {
            return;
        };

        match receiver.try_recv() {
            Ok(Ok(message)) => {
                self.status_message = message;
                self.copy_receiver = None;
                self.copy_in_progress = false;
            }
            Ok(Err(error)) => {
                self.status_message = error;
                self.copy_receiver = None;
                self.copy_in_progress = false;
            }
            Err(TryRecvError::Empty) => {
                ctx.request_repaint();
            }
            Err(TryRecvError::Disconnected) => {
                self.status_message = String::from("Copy worker disconnected.");
                self.copy_receiver = None;
                self.copy_in_progress = false;
            }
        }
    }

    fn install_textures(&mut self, ctx: &Context, images: &GeneratedImages) {
        let heightmap = gray_to_color_image(&images.heightmap_image);
        let bumpmap = gray_to_color_image(&images.bumpmap_image);
        let terrain_composite = rgba_to_color_image(&images.terrain_composite_image);
        let terrain_color = rgba_to_color_image(&images.terrain_color_image);

        self.textures = Some(TextureSet {
            terrain_composite: ctx.load_texture(
                "terrain_composite_preview",
                terrain_composite,
                TextureOptions::default(),
            ),
            terrain_color: ctx.load_texture(
                "terrain_color_preview",
                terrain_color,
                TextureOptions::default(),
            ),
            heightmap: ctx.load_texture("heightmap_preview", heightmap, TextureOptions::default()),
            bumpmap: ctx.load_texture("bumpmap_preview", bumpmap, TextureOptions::default()),
        });
    }

    fn copy_preview(&mut self, preview: PreviewFocus) {
        if self.copy_in_progress {
            self.status_message = String::from("Copy already in progress...");
            return;
        }

        let Some(generated) = &self.generated else {
            self.status_message = String::from("No generated image is available to copy.");
            return;
        };

        let image = preview_image_data(generated, preview);
        let success_message = match preview {
            PreviewFocus::TerrainComposite => {
                String::from("Terrain composite copied to clipboard.")
            }
            PreviewFocus::TerrainColor => String::from("Terrain color copied to clipboard."),
            PreviewFocus::Heightmap => String::from("Heightmap copied to clipboard."),
            PreviewFocus::Bumpmap => String::from("Bumpmap copied to clipboard."),
        };

        let (tx, rx) = mpsc::channel();
        self.copy_receiver = Some(rx);
        self.copy_in_progress = true;
        self.status_message = match preview {
            PreviewFocus::TerrainComposite => {
                String::from("Copying terrain composite to clipboard...")
            }
            PreviewFocus::TerrainColor => String::from("Copying terrain color to clipboard..."),
            PreviewFocus::Heightmap => String::from("Copying heightmap to clipboard..."),
            PreviewFocus::Bumpmap => String::from("Copying bumpmap to clipboard..."),
        };

        thread::spawn(move || {
            let result = copy_image_to_clipboard(&image).map(|_| success_message);
            let _ = tx.send(result);
        });
    }

    fn preview_context_menu(response: &Response) -> bool {
        let mut should_copy = false;
        response.context_menu(|ui| {
            if ui.button("Copy").clicked() {
                should_copy = true;
                ui.close();
            }
        });
        should_copy
    }

    fn overlay_copy_button(ui: &mut Ui, anchor_rect: egui::Rect, preview: PreviewFocus) -> bool {
        let button_size = Vec2::new(28.0, 28.0);
        let button_rect = Align2::RIGHT_TOP
            .align_size_within_rect(button_size, anchor_rect.shrink2(Vec2::new(8.0, 8.0)));
        let response = ui.interact(
            button_rect,
            ui.make_persistent_id(("main_copy_button", preview == PreviewFocus::Bumpmap)),
            Sense::click(),
        );

        let fill = if response.hovered() {
            Color32::from_rgba_premultiplied(70, 90, 120, 220)
        } else {
            Color32::from_rgba_premultiplied(40, 40, 40, 180)
        };
        ui.painter()
            .rect_filled(button_rect, CornerRadius::same(6), fill);
        ui.painter().rect_stroke(
            button_rect,
            CornerRadius::same(6),
            Stroke::new(1.0, Color32::from_rgba_premultiplied(255, 255, 255, 60)),
            StrokeKind::Outside,
        );
        ui.painter().text(
            button_rect.center(),
            Align2::CENTER_CENTER,
            "🗍",
            FontId::proportional(14.0),
            Color32::WHITE,
        );

        let clicked = response.clicked();
        response.on_hover_text("Copy");
        let _ = preview;
        clicked
    }

    fn tab_button(ui: &mut Ui, current: &mut LeftTab, tab: LeftTab, label: &str) {
        let selected = *current == tab;
        if ui.selectable_label(selected, label).clicked() {
            *current = tab;
        }
    }

    fn ui_left_panel(&mut self, ctx: &Context) {
        egui::SidePanel::left("controls_panel")
            .resizable(true)
            .default_width(320.0)
            .min_width(280.0)
            .show(ctx, |ui| {
                ui.heading("Advanced Erosion Terrain Generator");
                ui.add_space(8.0);

                ui.horizontal(|ui| {
                    Self::tab_button(
                        ui,
                        &mut self.generation_tab,
                        LeftTab::Generation,
                        "Generation",
                    );
                    Self::tab_button(ui, &mut self.generation_tab, LeftTab::Advanced, "Advanced");
                });

                ui.separator();

                match self.generation_tab {
                    LeftTab::Generation => self.ui_generation_tab(ui),
                    LeftTab::Advanced => self.ui_advanced_tab(ui),
                }

                ui.add_space(12.0);
                let button = ui.add_sized(
                    [ui.available_width(), 56.0],
                    egui::Button::new(RichText::new("Generate").size(20.0)),
                );
                if button.clicked() {
                    self.start_generation();
                }

                if self.generation_in_progress {
                    ui.add_space(8.0);
                    ui.spinner();
                }
                if self.copy_in_progress {
                    ui.add_space(4.0);
                    ui.spinner();
                }

                ui.add_space(8.0);
                ui.label(&self.status_message);
            });
    }

    fn ui_generation_tab(&mut self, ui: &mut Ui) {
        labeled_int_slider(ui, "Image Width", &mut self.params.image_width, 1..=5000);
        labeled_int_slider(ui, "Image Height", &mut self.params.image_height, 1..=5000);
        labeled_float_slider(
            ui,
            "Generation Scale",
            &mut self.params.generation_scale,
            0.1..=50.0,
            0.1,
        );
        labeled_float_slider(
            ui,
            "Erosion Scale",
            &mut self.params.erosion_scale,
            0.08..=0.25,
            0.01,
        );

        ui.label("Random Seed");
        ui.horizontal(|ui| {
            ui.add(
                DragValue::new(&mut self.params.random_seed)
                    .range(0..=u64::MAX)
                    .speed(1.0),
            );

            if ui.button("Random").clicked() {
                self.params.random_seed = rand::rng().random::<u64>();
            }
        });
    }

    fn ui_advanced_tab(&mut self, ui: &mut Ui) {
        labeled_float_slider(
            ui,
            "Strength",
            &mut self.params.erosion_strength,
            0.01..=0.10,
            0.01,
        );
        labeled_float_slider(
            ui,
            "Gully Weight",
            &mut self.params.erosion_gully_weight,
            0.0..=1.0,
            0.01,
        );
        labeled_float_slider(
            ui,
            "Detail",
            &mut self.params.erosion_detail,
            0.7..=3.0,
            0.1,
        );
        labeled_int_slider(ui, "Octaves", &mut self.params.erosion_octaves, 1..=8);
        labeled_int_slider(
            ui,
            "Snowline Height",
            &mut self.params.snowline_height,
            0..=255,
        );
        labeled_int_slider(
            ui,
            "Water Table Height",
            &mut self.params.water_table_height,
            0..=255,
        );
        labeled_float_slider(
            ui,
            "Rock Slope Angle",
            &mut self.params.rock_slope_angle_degrees,
            0.0..=89.0,
            1.0,
        );
    }

    fn ui_right_panel(&mut self, ctx: &Context) {
        let mut pending_copy = None;
        let mut pending_focus = None;
        let mut pending_hover = None;
        let preview_focus = self.preview_focus;
        let textures = self.textures.clone();
        egui::CentralPanel::default().show(ctx, |ui| {
            let Some(textures) = &textures else {
                ui.centered_and_justified(|ui| {
                    ui.label("No generated images yet.");
                });
                return;
            };
            let terrain_composite_texture = textures.terrain_composite.clone();
            let terrain_color_texture = textures.terrain_color.clone();
            let heightmap_texture = textures.heightmap.clone();
            let bumpmap_texture = textures.bumpmap.clone();

            let available = ui.available_size();
            let bottom_strip_height = (available.y * 0.24).clamp(120.0, 220.0);
            let main_height = (available.y - bottom_strip_height - 8.0).max(100.0);

            let main_texture = match preview_focus {
                PreviewFocus::TerrainComposite => &terrain_composite_texture,
                PreviewFocus::TerrainColor => &terrain_color_texture,
                PreviewFocus::Heightmap => &heightmap_texture,
                PreviewFocus::Bumpmap => &bumpmap_texture,
            };

            let main_response =
                render_clickable_image(ui, main_texture, Vec2::new(available.x, main_height), true);
            if main_response.hovered() {
                pending_hover = Some(preview_focus);
            }
            if Self::preview_context_menu(&main_response)
                || Self::overlay_copy_button(ui, main_response.rect, preview_focus)
            {
                pending_copy = Some(preview_focus);
            }
            ui.add_space(8.0);

            let thumb_width = (available.x * 0.26).clamp(120.0, 220.0);
            let panel_width = (thumb_width + 32.0).min(available.x);
            let secondary_previews = secondary_previews(preview_focus);

            ui.horizontal_centered(|ui| {
                ui.horizontal(|ui| {
                    for secondary_focus in secondary_previews {
                        let secondary_texture = match secondary_focus {
                            PreviewFocus::TerrainComposite => &terrain_composite_texture,
                            PreviewFocus::TerrainColor => &terrain_color_texture,
                            PreviewFocus::Heightmap => &heightmap_texture,
                            PreviewFocus::Bumpmap => &bumpmap_texture,
                        };
                        let secondary_label = preview_label(secondary_focus);

                        Frame::group(ui.style()).show(ui, |ui| {
                            ui.set_min_width(panel_width);
                            ui.set_max_width(panel_width);
                            ui.vertical_centered(|ui| {
                                ui.label(RichText::new(secondary_label).strong());
                                ui.add_space(6.0);
                                let secondary_response = render_clickable_image(
                                    ui,
                                    secondary_texture,
                                    Vec2::new(thumb_width, bottom_strip_height - 32.0),
                                    true,
                                );
                                if secondary_response.hovered() {
                                    pending_hover = Some(secondary_focus);
                                }
                                if Self::preview_context_menu(&secondary_response) {
                                    pending_copy = Some(secondary_focus);
                                }
                                if secondary_response.clicked() {
                                    pending_focus = Some(secondary_focus);
                                }
                            });
                        });
                    }
                });
            });
        });

        if let Some(preview) = pending_copy {
            self.copy_preview(preview);
        }
        if let Some(preview) = pending_focus {
            self.preview_focus = preview;
        }
        self.hovered_preview = pending_hover;
    }
}

impl eframe::App for TerrainApp {
    fn update(&mut self, ctx: &Context, _frame: &mut eframe::Frame) {
        self.hovered_preview = None;
        self.poll_generation(ctx);
        self.poll_copy(ctx);
        self.ui_left_panel(ctx);
        self.ui_right_panel(ctx);

        let wants_keyboard_input = ctx.wants_keyboard_input();
        let should_copy = ctx.input(|input| {
            !wants_keyboard_input
                && input.events.iter().any(|event| match event {
                    Event::Copy => true,
                    Event::Key {
                        key: Key::C,
                        pressed: true,
                        modifiers,
                        ..
                    } => modifiers.ctrl || modifiers.command,
                    _ => false,
                })
        });

        if should_copy {
            let preview = self.hovered_preview.unwrap_or(self.preview_focus);
            self.copy_preview(preview);
        }
    }
}

fn gray_to_color_image(image: &image::GrayImage) -> ColorImage {
    let size = [image.width() as usize, image.height() as usize];
    let mut pixels = Vec::with_capacity(size[0] * size[1]);
    for pixel in image.pixels() {
        let value = pixel.0[0];
        pixels.push(Color32::from_gray(value));
    }
    ColorImage::new(size, pixels)
}

fn rgba_to_color_image(image: &image::RgbaImage) -> ColorImage {
    let size = [image.width() as usize, image.height() as usize];
    let mut pixels = Vec::with_capacity(size[0] * size[1]);
    for pixel in image.pixels() {
        let [r, g, b, a] = pixel.0;
        pixels.push(Color32::from_rgba_unmultiplied(r, g, b, a));
    }
    ColorImage::new(size, pixels)
}

fn preview_image_data(images: &GeneratedImages, preview: PreviewFocus) -> ImageData<'static> {
    match preview {
        PreviewFocus::TerrainComposite => {
            rgba_image_to_clipboard_data(&images.terrain_composite_image)
        }
        PreviewFocus::TerrainColor => rgba_image_to_clipboard_data(&images.terrain_color_image),
        PreviewFocus::Heightmap => gray_image_to_clipboard_data(&images.heightmap_image),
        PreviewFocus::Bumpmap => gray_image_to_clipboard_data(&images.bumpmap_image),
    }
}

fn preview_label(preview: PreviewFocus) -> &'static str {
    match preview {
        PreviewFocus::TerrainComposite => "Composite",
        PreviewFocus::TerrainColor => "Terrain Color",
        PreviewFocus::Heightmap => "Heightmap",
        PreviewFocus::Bumpmap => "Bumpmap",
    }
}

fn secondary_previews(preview: PreviewFocus) -> [PreviewFocus; 3] {
    match preview {
        PreviewFocus::TerrainComposite => [
            PreviewFocus::TerrainColor,
            PreviewFocus::Heightmap,
            PreviewFocus::Bumpmap,
        ],
        PreviewFocus::TerrainColor => [
            PreviewFocus::TerrainComposite,
            PreviewFocus::Heightmap,
            PreviewFocus::Bumpmap,
        ],
        PreviewFocus::Heightmap => [
            PreviewFocus::TerrainComposite,
            PreviewFocus::TerrainColor,
            PreviewFocus::Bumpmap,
        ],
        PreviewFocus::Bumpmap => [
            PreviewFocus::TerrainComposite,
            PreviewFocus::TerrainColor,
            PreviewFocus::Heightmap,
        ],
    }
}

fn labeled_int_slider<Num: egui::emath::Numeric>(
    ui: &mut Ui,
    label: &str,
    value: &mut Num,
    range: std::ops::RangeInclusive<Num>,
) {
    ui.label(label);
    ui.add(Slider::new(value, range));
}

fn labeled_float_slider(
    ui: &mut Ui,
    label: &str,
    value: &mut f32,
    range: std::ops::RangeInclusive<f32>,
    step: f32,
) {
    ui.label(label);
    let slider = Slider::new(value, range).step_by(step as f64);
    ui.add(slider);
}

fn render_clickable_image(
    ui: &mut Ui,
    texture: &TextureHandle,
    max_size: Vec2,
    clickable: bool,
) -> Response {
    let image_size = texture.size_vec2();
    let scale = (max_size.x / image_size.x).min(max_size.y / image_size.y);
    let desired_size = image_size * scale.max(0.01);
    let image = egui::Image::new(texture)
        .fit_to_exact_size(desired_size)
        .sense(if clickable {
            Sense::click()
        } else {
            Sense::hover()
        });
    let response = ui.add(image);

    if response.hovered() {
        ui.painter().rect_stroke(
            response.rect.expand(4.0),
            CornerRadius::same(8),
            Stroke::new(2.0, Color32::from_rgba_premultiplied(180, 220, 255, 180)),
            StrokeKind::Outside,
        );
    }

    response
}

fn gray_image_to_clipboard_data(image: &image::GrayImage) -> ImageData<'static> {
    let mut rgba = Vec::with_capacity((image.width() * image.height() * 4) as usize);
    for pixel in image.pixels() {
        let gray = pixel.0[0];
        rgba.extend_from_slice(&[gray, gray, gray, 255]);
    }

    ImageData {
        width: image.width() as usize,
        height: image.height() as usize,
        bytes: rgba.into(),
    }
}

fn rgba_image_to_clipboard_data(image: &image::RgbaImage) -> ImageData<'static> {
    let mut rgba = Vec::with_capacity((image.width() * image.height() * 4) as usize);
    for pixel in image.pixels() {
        rgba.extend_from_slice(&pixel.0);
    }

    ImageData {
        width: image.width() as usize,
        height: image.height() as usize,
        bytes: rgba.into(),
    }
}

fn copy_image_to_clipboard(image: &ImageData<'static>) -> Result<(), String> {
    let mut clipboard = Clipboard::new().map_err(|error| error.to_string())?;
    clipboard
        .set_image(image.clone())
        .map_err(|error| error.to_string())
}