mod binder;
mod blas;
mod extract;
mod types;

use bevy_asset::{Assets, Handle};
use bevy_ecs::{
    resource::Resource,
    system::{Res, ResMut},
};
use bevy_shader::{Shader, load_shader_library};
pub use binder::RaytracingSceneBindings;
use blacksite_surface::SurfaceEvaluatorRegistry;
pub use types::RaytracingMesh3d;

use crate::SolariPlugins;
use bevy_app::{App, Plugin, PostUpdate};
use bevy_ecs::schedule::IntoScheduleConfigs;
use bevy_render::{
    ExtractSchedule, GpuResourceAppExt, Render, RenderApp, RenderSystems,
    extract_resource::ExtractResourcePlugin,
    mesh::{
        RenderMesh,
        allocator::{MeshAllocatorSettings, allocate_and_free_meshes},
    },
    render_asset::prepare_assets,
    render_resource::BufferUsages,
    renderer::RenderDevice,
};
use binder::prepare_raytracing_scene_bindings;
use blas::{BlasManager, compact_raytracing_blas, prepare_raytracing_blas};
use extract::{StandardMaterialAssets, SurfaceMaterialAssets, extract_raytracing_scene};
use tracing::warn;

const EVALUATOR_LIBRARY_MARKER: &str = "// BLACKSITE_SURFACE_EVALUATOR_LIBRARY";
const EVALUATOR_DISPATCH_MARKER: &str = "// BLACKSITE_SURFACE_EVALUATOR_DISPATCH";

#[derive(Resource)]
struct SceneShaderLibraryHandle(Handle<Shader>);

#[derive(Resource, Default)]
struct SceneShaderEvaluatorRevisions(Vec<(u32, u64)>);

/// Creates acceleration structures and binding arrays of resources for raytracing.
pub struct RaytracingScenePlugin;

impl Plugin for RaytracingScenePlugin {
    fn build(&self, app: &mut App) {
        load_shader_library!(app, "brdf.wgsl");
        bevy_asset::embedded_asset!(app, "raytracing_scene_bindings.wgsl");
        let scene_shader: Handle<Shader> =
            bevy_asset::load_embedded_asset!(app, "raytracing_scene_bindings.wgsl");
        app.insert_resource(SceneShaderLibraryHandle(scene_shader));
        app.init_resource::<SceneShaderEvaluatorRevisions>()
            .add_systems(PostUpdate, sync_surface_evaluator_library);
        load_shader_library!(app, "sampling.wgsl");
    }

    fn finish(&self, app: &mut App) {
        let render_app = app.sub_app_mut(RenderApp);
        let render_device = render_app.world().resource::<RenderDevice>();
        let features = render_device.features();
        if !features.contains(SolariPlugins::required_wgpu_features()) {
            warn!(
                "RaytracingScenePlugin not loaded. GPU lacks support for required features: {:?}.",
                SolariPlugins::required_wgpu_features().difference(features)
            );
            return;
        }

        app.add_plugins((
            ExtractResourcePlugin::<StandardMaterialAssets>::default(),
            ExtractResourcePlugin::<SurfaceMaterialAssets>::default(),
            ExtractResourcePlugin::<SurfaceEvaluatorRegistry>::default(),
        ));

        let render_app = app.sub_app_mut(RenderApp);

        render_app
            .world_mut()
            .resource_mut::<MeshAllocatorSettings>()
            .extra_buffer_usages |= BufferUsages::BLAS_INPUT | BufferUsages::STORAGE;

        render_app
            .init_gpu_resource::<BlasManager>()
            .init_gpu_resource::<StandardMaterialAssets>()
            .init_gpu_resource::<SurfaceMaterialAssets>()
            .init_gpu_resource::<SurfaceEvaluatorRegistry>()
            .insert_resource(RaytracingSceneBindings::new())
            .add_systems(ExtractSchedule, extract_raytracing_scene)
            .add_systems(
                Render,
                (
                    prepare_raytracing_blas
                        .in_set(RenderSystems::PrepareAssets)
                        .before(prepare_assets::<RenderMesh>)
                        .after(allocate_and_free_meshes),
                    compact_raytracing_blas
                        .in_set(RenderSystems::PrepareAssets)
                        .after(prepare_raytracing_blas),
                    prepare_raytracing_scene_bindings.in_set(RenderSystems::PrepareBindGroups),
                ),
            );
    }
}

fn sync_surface_evaluator_library(
    registry: Res<SurfaceEvaluatorRegistry>,
    handle: Res<SceneShaderLibraryHandle>,
    mut revisions: ResMut<SceneShaderEvaluatorRevisions>,
    mut shaders: ResMut<Assets<Shader>>,
) {
    let mut next_revisions = registry
        .evaluators
        .values()
        .map(|record| (record.shader_id, record.revision))
        .collect::<Vec<_>>();
    next_revisions.sort_unstable();
    if revisions.0 == next_revisions {
        return;
    }

    let source =
        generate_scene_shader_source(include_str!("raytracing_scene_bindings.wgsl"), &registry);
    let shader = Shader::from_wgsl(source, "generated://bevy_solari/scene_bindings.wgsl");
    if let Err(error) = shaders.insert(handle.0.id(), shader) {
        warn!("Solari kept the last-good Surface evaluator library: {error}");
        return;
    }
    revisions.0 = next_revisions;
}

fn generate_scene_shader_source(template: &str, registry: &SurfaceEvaluatorRegistry) -> String {
    let mut records = registry.evaluators.values().collect::<Vec<_>>();
    records.sort_unstable_by_key(|record| record.shader_id);

    let mut library = String::new();
    let mut dispatch = String::new();
    for record in records {
        let names = wgsl_declaration_names(&record.evaluator_source);
        if !names.iter().any(|name| name == "evaluate") {
            continue;
        }
        let prefix = format!("blacksite_{}_", record.shader_id);
        let renamed = rename_wgsl_identifiers(&record.evaluator_source, &names, &prefix);
        library.push_str(&renamed);
        library.push('\n');
        dispatch.push_str(&format!(
            "    if shader_id == {}u {{ return {}evaluate(input, params, samples); }}\n",
            record.shader_id, prefix
        ));
    }

    template
        .replace(EVALUATOR_LIBRARY_MARKER, &library)
        .replace(EVALUATOR_DISPATCH_MARKER, &dispatch)
}

fn wgsl_declaration_names(source: &str) -> Vec<String> {
    let bytes = source.as_bytes();
    let mut names = Vec::new();
    let mut index = 0;
    while index < bytes.len() {
        if !is_ident_start(bytes[index]) {
            index += 1;
            continue;
        }
        let token_start = index;
        index += 1;
        while index < bytes.len() && is_ident_continue(bytes[index]) {
            index += 1;
        }
        if !matches!(
            &source[token_start..index],
            "fn" | "struct" | "alias" | "const" | "override"
        ) {
            continue;
        }
        while index < bytes.len() && bytes[index].is_ascii_whitespace() {
            index += 1;
        }
        if index >= bytes.len() || !is_ident_start(bytes[index]) {
            continue;
        }
        let name_start = index;
        index += 1;
        while index < bytes.len() && is_ident_continue(bytes[index]) {
            index += 1;
        }
        let name = source[name_start..index].to_string();
        if !names.contains(&name) {
            names.push(name);
        }
    }
    names
}

fn rename_wgsl_identifiers(source: &str, names: &[String], prefix: &str) -> String {
    let bytes = source.as_bytes();
    let mut result = String::with_capacity(source.len() + names.len() * prefix.len());
    let mut index = 0;
    while index < bytes.len() {
        if !is_ident_start(bytes[index]) {
            let character = source[index..]
                .chars()
                .next()
                .expect("index is inside the source string");
            result.push(character);
            index += character.len_utf8();
            continue;
        }
        let start = index;
        index += 1;
        while index < bytes.len() && is_ident_continue(bytes[index]) {
            index += 1;
        }
        let token = &source[start..index];
        if names.iter().any(|name| name == token) {
            result.push_str(prefix);
        }
        result.push_str(token);
    }
    result
}

const fn is_ident_start(byte: u8) -> bool {
    byte == b'_' || byte.is_ascii_alphabetic()
}

const fn is_ident_continue(byte: u8) -> bool {
    is_ident_start(byte) || byte.is_ascii_digit()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn evaluator_generation_namespaces_helpers() {
        let evaluator = "const SCALE: f32 = 1.0; fn helper(v: f32) -> f32 { return v * SCALE; } fn evaluate(input: SurfaceInput, params: SurfaceParams, samples: SurfaceSamples) -> Surface { var s = surface_default(); s.metallic = helper(1.0); return s; }";
        let names = wgsl_declaration_names(evaluator);
        let generated = rename_wgsl_identifiers(evaluator, &names, "blacksite_42_");
        assert!(generated.contains("const blacksite_42_SCALE"));
        assert!(generated.contains("v * blacksite_42_SCALE"));
        assert!(generated.contains("fn blacksite_42_helper"));
        assert!(generated.contains("blacksite_42_helper(1.0)"));
        assert!(generated.contains("fn blacksite_42_evaluate"));
    }
}
