use std::sync::Arc;
use bevy::ecs::prelude::*;
use bevy::pbr::{MeshMaterial3d, StandardMaterial};
use bevy::asset::{AssetId, Assets};
use crate::felling::{CutFace, Felled, TreeStump, calm_params};
use crate::params::TreeParams;
use crate::plugin::{Paint, ProceduralTree, TreeRender, repaint_flora, set_color};
use crate::pool::{PoolColors, RecolorPool, TreePool};
use crate::shading::{FloraMaterial, FloraMaterialSupport, FloraWindBuffer, SharedFlora, TextureLibrary};
#[derive(Message, Clone, Copy, Debug, PartialEq)]
pub struct RecolorTree {
pub tree: Entity,
pub colors: PoolColors,
}
#[derive(Component, Clone, Copy, Debug, PartialEq)]
pub(crate) struct TreeColors(pub(crate) PoolColors);
#[derive(Component, Clone)]
pub(crate) struct CutWood {
pub(crate) params: Arc<TreeParams>,
pub(crate) unlit: bool,
pub(crate) calm: bool,
}
#[derive(Component, Clone, Copy, Debug, PartialEq, Eq)]
pub struct FromPool {
pub pool: AssetId<TreePool>,
pub look: u32,
}
pub(crate) fn recolor_from_pools(
mut pools: MessageReader<RecolorPool>,
mut trees: MessageWriter<RecolorTree>,
assets: Option<Res<bevy::asset::Assets<TreePool>>>,
left: Query<(Entity, &FromPool, Option<&TreeStump>)>,
) {
for recolor in pools.read().filter(|recolor| recolor.glide <= 0.0) {
let looks = assets.as_deref().and_then(|assets| assets.get(recolor.pool)).map_or(u32::MAX, |pool| pool.looks.len().max(1) as u32);
let look = recolor.look % looks;
for (entity, from, stump) in &left {
if (from.pool, from.look) != (recolor.pool, look) || stump.and_then(|s| s.tree).is_some_and(|tree| left.contains(tree)) {
continue;
}
trees.write(RecolorTree { tree: entity, colors: recolor.colors });
}
}
}
#[allow(clippy::type_complexity, clippy::too_many_arguments)]
pub(crate) fn recolor_trees(
mut commands: Commands,
mut messages: MessageReader<RecolorTree>,
mut materials: ResMut<Assets<StandardMaterial>>,
mut flora: ResMut<Assets<FloraMaterial>>,
mut shared: ResMut<SharedFlora>,
mut textures: TextureLibrary,
mut cut_face: ResMut<CutFace>,
flora_support: Option<Res<FloraMaterialSupport>>,
wind: Option<Res<FloraWindBuffer>>,
mut trees: Query<(&ProceduralTree, &mut TreeRender, Option<&mut Felled>, Option<&Children>)>,
mut cuts: Query<(&mut CutWood, &Children)>,
plain: Query<&MeshMaterial3d<StandardMaterial>>,
flora_children: Query<(), With<MeshMaterial3d<FloraMaterial>>>,
) {
let wind = wind.as_deref().filter(|_| flora_support.is_some());
let mut work: Vec<(Entity, PoolColors)> = messages.read().map(|m| (m.tree, m.colors)).collect();
let mut next = 0;
while next < work.len() {
let (entity, colors) = work[next];
next += 1;
let unlit = colors.unlit;
if let Ok((tree, mut render, mut felled, children)) = trees.get_mut(entity) {
if matches!(render.paint, Paint::Growth(_)) {
continue;
}
commands.entity(entity).insert(TreeColors(colors));
let painted = colors.apply(&tree.0);
let standard = match &render.paint {
Paint::Standard { trunk, foliage } => Some((trunk.clone(), foliage.clone())),
_ => None,
};
if let Some((trunk, foliage)) = standard {
set_color(&mut materials, &trunk, colors.wood, unlit);
set_color(&mut materials, &foliage, colors.leaves, unlit);
} else if let Some(wind) = wind {
let mut resolve = |params: &TreeParams| {
(
shared.resolve(&mut flora, &textures.wood(params).with_unlit(unlit), wind),
shared.resolve(&mut flora, &textures.leaves(params).with_unlit(unlit), wind),
)
};
if felled.as_deref().is_some_and(Felled::landed) {
let (wood, leaves) = resolve(&calm_params(&painted));
let level0 = render.foliage.iter().chain(&render.flowers).map(|child| child.0);
for child in level0 {
commands.entity(child).insert(MeshMaterial3d(leaves.clone()));
}
commands.entity(render.trunk).insert(MeshMaterial3d(wood));
} else {
let (wood, leaves) = resolve(&painted);
repaint_flora(&mut commands, &mut render, wood, leaves, felled.is_none());
if let Some(felled) = felled.as_deref_mut() {
let (wood, leaves) = resolve(&calm_params(&painted));
felled.restyle_still(render.trunk, &wood, &leaves);
}
}
}
for child in children.into_iter().flat_map(|children| children.iter()) {
if let Ok(current) = plain.get(child)
&& cut_face.is_face(¤t.0)
{
let face = cut_face.get(&mut textures, &mut materials, unlit);
if current.0 != face {
commands.entity(child).insert(MeshMaterial3d(face));
}
}
}
if let Some(felled) = felled.as_deref() {
work.push((felled.stump, colors));
}
} else if let Ok((mut cut, children)) = cuts.get_mut(entity) {
let painted = colors.apply(&cut.params);
let params = if cut.calm { calm_params(&painted) } else { painted.clone() };
cut.params = Arc::new(painted);
cut.unlit = unlit;
let wood = wind
.map(|wind| shared.resolve(&mut flora, &textures.wood(¶ms).with_unlit(unlit), wind));
for child in children.iter() {
if flora_children.contains(child) {
if let Some(wood) = &wood {
commands.entity(child).insert(MeshMaterial3d(wood.clone()));
}
} else if let Ok(current) = plain.get(child) {
if cut_face.is_face(¤t.0) {
let face = cut_face.get(&mut textures, &mut materials, unlit);
if current.0 != face {
commands.entity(child).insert(MeshMaterial3d(face));
}
} else {
set_color(&mut materials, ¤t.0.clone(), colors.wood, unlit);
}
}
}
}
}
}