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 {
    Heightmap,
    Bumpmap,
}

struct TexturePair {
    heightmap: TextureHandle,
    bumpmap: TextureHandle,
}

pub struct TerrainApp {
    params: GenerationParams,
    generation_tab: LeftTab,
    preview_focus: PreviewFocus,
    hovered_preview: Option<PreviewFocus>,
    textures: Option<TexturePair>,
    generated: Option<GeneratedImages>,
    generation_receiver: Option<Receiver<Result<GeneratedImages, String>>>,
    generation_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::Bumpmap,
            hovered_preview: None,
            textures: None,
            generated: None,
            generation_receiver: None,
            generation_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 and bumpmap.",
                    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 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);

        self.textures = Some(TexturePair {
            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) {
        let Some(generated) = &self.generated else {
            self.status_message = String::from("No generated image is available to copy.");
            return;
        };

        let image = match preview {
            PreviewFocus::Heightmap => &generated.heightmap_image,
            PreviewFocus::Bumpmap => &generated.bumpmap_image,
        };

        match copy_gray_image_to_clipboard(image) {
            Ok(()) => {
                self.status_message = match preview {
                    PreviewFocus::Heightmap => String::from("Heightmap copied to clipboard."),
                    PreviewFocus::Bumpmap => String::from("Bumpmap copied to clipboard."),
                };
            }
            Err(error) => {
                self.status_message = format!("Failed to copy image: {error}");
            }
        }
    }

    fn preview_context_menu(&mut self, response: &Response, preview: PreviewFocus) {
        response.context_menu(|ui| {
            if ui.button("Copy").clicked() {
                self.copy_preview(preview);
                ui.close();
            }
        });
    }

    fn overlay_copy_button(&mut self, ui: &mut Ui, anchor_rect: egui::Rect, preview: PreviewFocus) {
        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,
        );

        if response.clicked() {
            self.copy_preview(preview);
        }

        response.on_hover_text("Copy");
    }

    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();
                }

                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);
    }

    fn ui_right_panel(&mut self, ctx: &Context) {
        egui::CentralPanel::default().show(ctx, |ui| {
            let Some(textures) = &self.textures else {
                ui.centered_and_justified(|ui| {
                    ui.label("No generated images yet.");
                });
                return;
            };
            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).max(120.0).min(220.0);
            let main_height = (available.y - bottom_strip_height - 8.0).max(100.0);

            let main_texture = match self.preview_focus {
                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() {
                self.hovered_preview = Some(self.preview_focus);
            }
            self.preview_context_menu(&main_response, self.preview_focus);
            self.overlay_copy_button(ui, main_response.rect, self.preview_focus);
            ui.add_space(8.0);

            let thumb_width = (available.x * 0.28).max(120.0).min(240.0);
            let panel_width = (thumb_width + 32.0).min(available.x);
            let secondary_texture = match self.preview_focus {
                PreviewFocus::Heightmap => &bumpmap_texture,
                PreviewFocus::Bumpmap => &heightmap_texture,
            };
            let secondary_focus = match self.preview_focus {
                PreviewFocus::Heightmap => PreviewFocus::Bumpmap,
                PreviewFocus::Bumpmap => PreviewFocus::Heightmap,
            };
            let secondary_label = match secondary_focus {
                PreviewFocus::Heightmap => "Heightmap",
                PreviewFocus::Bumpmap => "Bumpmap",
            };

            ui.horizontal_centered(|ui| {
                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() {
                            self.hovered_preview = Some(secondary_focus);
                        }
                        self.preview_context_menu(&secondary_response, secondary_focus);
                        if secondary_response.clicked() {
                            self.preview_focus = secondary_focus;
                        }
                    });
                });
            });
        });
    }
}

impl eframe::App for TerrainApp {
    fn update(&mut self, ctx: &Context, _frame: &mut eframe::Frame) {
        self.hovered_preview = None;
        self.poll_generation(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 labeled_int_slider(
    ui: &mut Ui,
    label: &str,
    value: &mut usize,
    range: std::ops::RangeInclusive<usize>,
) {
    ui.label(label);
    let mut slider_value = *value as u32;
    if ui
        .add(Slider::new(
            &mut slider_value,
            *range.start() as u32..=*range.end() as u32,
        ))
        .changed()
    {
        *value = slider_value as usize;
    }
}

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::click()
        });
    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 copy_gray_image_to_clipboard(image: &image::GrayImage) -> Result<(), String> {
    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]);
    }

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

    let mut clipboard = Clipboard::new().map_err(|error| error.to_string())?;
    clipboard.set_image(data).map_err(|error| error.to_string())
}