forestiles/src/map/cells/grow.rs

use std::net::UdpSocket;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use bevy::asset::{embedded_asset, embedded_path};
use bevy::prelude::*;
use bevy::tasks::{AsyncComputeTaskPool, Task};
use log::info;
use rand::rngs::SmallRng;
use rand::seq::IteratorRandom;
use rand::{rng, Rng};
use voronoice::Voronoi;

use crate::map::cells::{wealth_to_unit, WealthType};
use crate::map::{CellKind, CellsEntities, MeshNeedsUpdate};

use super::Cell;

pub struct Plugin;
impl bevy::prelude::Plugin for Plugin {
    fn build(&self, app: &mut App) {
        app.add_systems(Update, rx_chunks_updates);
        // embedded_asset!(app, "../../../", "assets/models/tree.glb");
    }
}

// Background task that manages terrain
#[derive(Resource)]
pub struct GrowThread(
    pub Task<()>,
    pub async_channel::Sender<Cell>,
    pub async_channel::Receiver<Cell>,
);
impl GrowThread {
    pub fn new(voronoi: Voronoi, gt: Arc<AtomicU64>, cells: Vec<Cell>, rng: SmallRng) -> Self {
        let task_pool = AsyncComputeTaskPool::get();
        let (tx_control, rx_control) = async_channel::unbounded();
        let (tx_update, rx_update) = async_channel::unbounded();
        Self(
            task_pool.spawn(grow_thread(voronoi, gt, cells, rng, rx_control, tx_update)),
            tx_control,
            rx_update,
        )
    }
}

// #[derive(Component)]
// struct GetChunkTask(Task<Vec<Cell>>);

// const MAX_UPDATE_PER_FRAME: usize = 10000;
fn rx_chunks_updates(
    mut chunks: Query<&mut MeshNeedsUpdate>,
    mut cells: Query<(&mut Cell, &ChildOf, Entity)>,
    cells_entities: Res<CellsEntities>,
    grow_thread: Res<GrowThread>,
    mut cmds: Commands,
    asset_server: Res<AssetServer>,
) {
    while let Ok(c) = grow_thread.2.try_recv() {
        let id = c.voronoi_id;
        let (mut cell, parent, id) = cells.get_mut(cells_entities.0[id]).unwrap();
        chunks.get_mut(parent.parent()).unwrap().0 += cell.kind.diff(&c.kind);

        if matches!(c.kind, CellKind::Forest { .. })
            && !matches!(cell.kind, CellKind::Forest { .. })
        {
            // Spawn tree
            cmds.spawn((
                SceneRoot(
                    asset_server
                        .load(GltfAssetLabel::Scene(0).from_asset("embedded://models/tree.glb")),
                ),
                ChildOf(id),
                Transform::from_scale(Vec3::new(0.1, 0.1, 0.1)).with_rotation(
                    Quat::from_axis_angle(
                        Vec3::Y,
                        rng().random_range(-std::f32::consts::PI..std::f32::consts::PI),
                    ),
                ),
            ));
        }

        *cell = c;
    }
}

const GRASS_EXTEND_STEP: u64 = CellKind::GRASS.growth_duration();
const FOREST_EXTEND_STEP: u64 = CellKind::FOREST.growth_duration();
pub async fn grow_thread(
    voronoi: Voronoi,
    gt: Arc<AtomicU64>,
    cells: Vec<Cell>,
    mut rng: SmallRng,
    rx_control: async_channel::Receiver<Cell>,
    tx_update: async_channel::Sender<Cell>,
) {
    // Add "modified" flag
    let mut cells: Vec<(Cell, f32)> = cells.into_iter().map(|c| (c, 0.)).collect();
    let mut last_grass_grow = 0;
    let mut last_forest_grow = 0;
    let mut last_grass_extend = 0;
    let mut last_forest_extend = 0;
    loop {
        // Receive control from the game
        while let Ok(c) = rx_control.try_recv() {
            let id = c.voronoi_id;
            cells[id] = (c, 1.);
        }

        let gt = gt.load(Ordering::Acquire);

        // Update step
        let grass_grow = gt.saturating_sub(last_grass_grow) * WealthType::MAX as u64
            / CellKind::GRASS.growth_duration();
        let forest_grow = gt.saturating_sub(last_forest_grow) * WealthType::MAX as u64
            / CellKind::FOREST.growth_duration();
        if grass_grow > 0 || forest_grow > 0 {
            for (c, modified) in cells.iter_mut() {
                if grass_grow > 0 {
                    if let CellKind::Grass { ref mut wealth } = c.kind {
                        *wealth = wealth.saturating_add(grass_grow as WealthType);
                        *modified += grass_grow as f32 / WealthType::MAX as f32;
                    }
                }
                if forest_grow > 0 {
                    if let CellKind::Forest { ref mut wealth } = c.kind {
                        *wealth = wealth.saturating_add(forest_grow as WealthType);
                        *modified += grass_grow as f32 / WealthType::MAX as f32;
                    }
                }
            }
        // Don't make multiple updates at once to avoid problems and lags
        } else if gt - last_grass_extend > GRASS_EXTEND_STEP {
            for i in 0..cells.len() {
                if cells[i].1 < 1. {
                    if let Some(wealth) = cells[i].0.kind.grass_wealth() {
                        if rng.gen_bool(wealth_to_unit(wealth) as f64) {
                            let target = &mut cells
                                [voronoi.cell(i).iter_neighbors().choose(&mut rng).unwrap()];

                            if matches!(target.0.kind, CellKind::Dirt) {
                                target.0.kind = CellKind::Grass { wealth: 0 };
                                target.1 = 1.;
                            }
                        }
                    }
                }
            }
            last_grass_extend += GRASS_EXTEND_STEP;
        } else if gt - last_forest_extend > FOREST_EXTEND_STEP {
            for i in 0..cells.len() {
                if cells[i].1 < 1. {
                    if let CellKind::Forest { wealth } = cells[i].0.kind {
                        if rng.gen_bool(wealth_to_unit(wealth) as f64) {
                            let target = &mut cells
                                [voronoi.cell(i).iter_neighbors().choose(&mut rng).unwrap()];
                            match target.0.kind {
                                CellKind::Dirt => {
                                    target.0.kind = CellKind::Forest { wealth: 0 };
                                    target.1 = 1.;
                                }
                                CellKind::Grass { wealth: w } => {
                                    target.0.kind = CellKind::Forest {
                                        wealth: (w as f32
                                            * CellKind::GRASS.growth_duration() as f32
                                            / CellKind::FOREST.growth_duration() as f32)
                                            as WealthType,
                                    };
                                    target.1 = 1.;
                                }
                                _ => {}
                            }
                        }
                    }
                }
            }
            last_forest_extend += FOREST_EXTEND_STEP;
        }

        let grass_count = cells
            .iter()
            .filter(|(c, _)| matches!(c.kind, CellKind::Grass { .. }))
            .count();

        last_grass_grow += grass_grow * CellKind::GRASS.growth_duration() / WealthType::MAX as u64;
        last_forest_grow += forest_grow * CellKind::FOREST.growth_duration();

        // Send modifications
        let modification_treshold = 1. - (1. / (tx_update.len() as f32));
        for (c, m) in cells.iter_mut().filter(|(_, m)| *m > modification_treshold) {
            tx_update.send(c.clone()).await.unwrap();
            *m = 0.;
        }
    }
}