mod tile_light;

use crate::grid;
use bevy::{
    prelude::*,
    reflect::TypePath,
    render::{
        render_resource::{encase, AsBindGroup, OwnedBindingResource, ShaderRef, ShaderType},
        renderer::RenderQueue,
        Render, RenderApp, RenderSet,
    },
    sprite::{Material2d, Material2dPlugin, MaterialMesh2dBundle, RenderMaterials2d},
};
use tile_light::{extract_tile_lights, TileLight, TileLightsUniform};

pub(crate) use tile_light::EmissiveTile;

pub struct CanvasPlugin;

impl Plugin for CanvasPlugin {
    fn build(&self, app: &mut App) {
        app.add_plugins(Material2dPlugin::<CanvasMaterial>::default())
            .add_systems(Startup, spawn_canvas);

        app.sub_app_mut(RenderApp)
            .add_systems(ExtractSchedule, extract_tile_lights)
            .add_systems(
                Render,
                prepare_canvas_material.in_set(RenderSet::PrepareAssets),
            );
    }
}

const CANVAS_COLOR: Color = Color::rgb(0.05, 0.05, 0.05);

#[derive(AsBindGroup, TypePath, Asset, Clone)]
struct CanvasMaterial {
    #[uniform(0)]
    ambient_color: Color,
    #[uniform(0)]
    tile_lights: tile_light::TileLightsUniform,
    #[uniform(0)]
    number_of_lights: u32,
}

impl Default for CanvasMaterial {
    fn default() -> Self {
        Self {
            ambient_color: CANVAS_COLOR,
            number_of_lights: 0,
            tile_lights: Default::default(),
        }
    }
}

#[derive(ShaderType)]
struct CanvasMaterialUniform {
    ambient_color: Color,
    tile_lights: tile_light::TileLightsUniform,
    number_of_lights: u32,
}

impl Material2d for CanvasMaterial {
    fn fragment_shader() -> ShaderRef {
        "canvas_shader.wgsl".into()
    }
}

fn prepare_canvas_material(
    materials: Res<RenderMaterials2d<CanvasMaterial>>,
    canvas_query: Query<&Handle<CanvasMaterial>>,
    tile_light_query: Query<&TileLight>,
    render_queue: Res<RenderQueue>,
) {
    let tile_lights = tile_light_query.iter().cloned().collect::<Vec<_>>();
    let tile_lights_uniform = TileLightsUniform::from_lights(&tile_lights);

    for handle in canvas_query.iter() {
        if let Some(material) = materials.get(&handle.id()) {
            let binding = &material.bindings[0];
            if let (_, OwnedBindingResource::Buffer(cur_buffer)) = binding {
                let mut buffer = encase::UniformBuffer::new(Vec::new());
                buffer
                    .write(&CanvasMaterialUniform {
                        ambient_color: Color::RED,
                        tile_lights: tile_lights_uniform.clone(),
                        number_of_lights: tile_lights.len() as _,
                    })
                    .unwrap();
                render_queue.write_buffer(cur_buffer, 0, buffer.as_ref());
            }
        }
    }
}

fn spawn_canvas(
    mut commands: Commands,
    mut mesh_assets: ResMut<Assets<Mesh>>,
    mut canvas_material_assets: ResMut<Assets<CanvasMaterial>>,
) {
    commands
        .spawn(MaterialMesh2dBundle {
            mesh: mesh_assets.add(Mesh::from(Rectangle::default())).into(),
            material: canvas_material_assets.add(CanvasMaterial::default()),
            transform: Transform::from_scale(Vec3::splat(
                grid::SEGMENT_SIZE * f32::from(grid::SIZE),
            )),
            ..Default::default()
        })
        .insert(Name::new("Canvas"));
}