wde_terrain_editor/
lib.rs

1//! Terrain editor plugin for WaterDropEngine.
2//!
3//! WIP.
4use wde_logger::prelude::*;
5
6use bevy::prelude::*;
7use wde_terrain::prelude::*;
8
9use crate::{
10    paint::{brush::PaintBrush, paint_manager::PaintManagerPlugin},
11    processor::PaintProcessorPlugin,
12    ui_painter::TerrainEditorUIPlugin
13};
14
15mod paint;
16mod processor;
17mod ui_painter;
18
19pub struct TerrainEditorPlugin;
20impl Plugin for TerrainEditorPlugin {
21    fn build(&self, app: &mut App) {
22        app.add_plugins((
23            PaintManagerPlugin,
24            PaintProcessorPlugin,
25            TerrainEditorUIPlugin
26        ))
27        .init_resource::<SaveManager>()
28        .add_systems(Startup, init)
29        .add_systems(Update, save_extracted_tiles)
30        .add_message::<ExtractedTileMessage>();
31    }
32}
33
34fn init(mut commands: Commands, asset_server: Res<AssetServer>) {
35    // Spawn a terrain
36    commands.spawn((
37        Name::new("Terrain"),
38        Terrain::load("tests/terrain"),
39        TerrainRenderer::new(&asset_server),
40        TerrainPhysics::default()
41    ));
42
43    // Spawn a default brush for testing
44    commands.spawn((Name::new("Terrain Default Brush"), PaintBrush::default()));
45}
46
47#[derive(Resource, Default)]
48struct SaveManager {
49    // True if we are currently saving the terrain
50    saving: bool,
51    // List of tiles to save, with their position and map type (0 for heightmap, 1 for splatmap)
52    tiles_to_save: Vec<(ChunkPos, u32, u32)>
53}
54
55fn save_extracted_tiles(
56    terrain: Query<&Terrain>,
57    mut save_manager: ResMut<SaveManager>,
58    mut extracted_tiles: MessageReader<ExtractedTileMessage>
59) {
60    // If we are not currently saving, ignore the extracted tiles
61    if !save_manager.saving {
62        return;
63    }
64
65    // Get the terrain from main world
66    let terrain = match terrain.single() {
67        Ok(terrain) => terrain,
68        Err(_) => return
69    };
70
71    // Process each extracted tile message
72    for message in extracted_tiles.read() {
73        let ExtractedTileMessage {
74            pos,
75            map_type,
76            splat_map_index,
77            data
78        } = message;
79
80        // Save the extracted tile data to a file
81        let cur_dir = std::env::current_dir().unwrap();
82        let full_path = format!("{}/res/{}", cur_dir.display(), terrain.path);
83        match map_type {
84            0 => {
85                let path = format!("{}/heightmap_{}_{}.png", full_path, pos.x, pos.y);
86                if let Err(e) = save_png_from_channels(
87                    &path,
88                    data,
89                    1,
90                    (CHUNK_RENDER_SUBDIVISIONS, CHUNK_RENDER_SUBDIVISIONS)
91                ) {
92                    error!(
93                        "Failed to save heightmap for tile ({}, {}): {}",
94                        pos.x, pos.y, e
95                    );
96                }
97            }
98            1 => {
99                let path = format!(
100                    "{}/splatmap_{}_{}-{}.png",
101                    full_path, pos.x, pos.y, splat_map_index
102                );
103                if let Err(e) = save_png_from_channels(
104                    &path,
105                    data,
106                    4,
107                    (CHUNK_RENDER_SUBDIVISIONS, CHUNK_RENDER_SUBDIVISIONS)
108                ) {
109                    error!(
110                        "Failed to save splatmap for tile ({}, {}): {}",
111                        pos.x, pos.y, e
112                    );
113                }
114            }
115            _ => continue
116        };
117        save_manager
118            .tiles_to_save
119            .retain(|(p, t, s)| !(p == pos && t == map_type && s == splat_map_index));
120    }
121
122    // If all tiles have been saved, mark saving as false
123    if save_manager.tiles_to_save.is_empty() {
124        info!("Finished saving terrain.");
125        save_manager.saving = false;
126    }
127}
128
129fn save_png_from_channels(
130    path: &str,
131    data: &[u8],
132    channels: usize,
133    dimensions: (u32, u32)
134) -> Result<(), String> {
135    // Convert the raw channel data to RGBA format for saving
136    let rgba_data = match channels {
137        1 => data.iter().flat_map(|&v| vec![v, v, v, 255]).collect(),
138        4 => data.to_vec(),
139        _ => return Err("Unsupported number of channels".to_string())
140    };
141
142    // Save the RGBA data as a PNG file using png crate
143    let file =
144        std::fs::File::create(path).map_err(|e| format!("Failed to create file {path}: {e}"))?;
145    let w = std::io::BufWriter::new(file);
146    let mut encoder = png::Encoder::new(w, dimensions.0, dimensions.1);
147    encoder.set_color(png::ColorType::Rgba);
148    encoder.set_depth(png::BitDepth::Eight);
149    let mut writer = encoder
150        .write_header()
151        .map_err(|e| format!("Failed to write PNG header for {path}: {e}"))?;
152    writer
153        .write_image_data(&rgba_data)
154        .map_err(|e| format!("Failed to write PNG data for {path}: {e}"))
155}