use cuneus::compute::*;
use cuneus::prelude::*;

cuneus::uniform_params! {
    struct SplattingParams {
    animation_speed: f32,
    splat_size: f32,
    particle_spread: f32,
    intensity: f32,
    particle_density: f32,
    brightness: f32,
    physics_strength: f32,
    trail_length: f32,
    trail_decay: f32,
    flow_strength: f32,
    _padding1: f32,
    _padding2: u32}
}

struct ColorProjection {
    base: RenderKit,
    compute_shader: ComputeShader,
    current_params: SplattingParams}

impl ColorProjection {
    fn clear_buffers(&mut self, core: &Core) {
        self.compute_shader.clear_all_buffers(core);
    }
}

impl ShaderManager for ColorProjection {
    fn init(core: &Core) -> Self {
        let base = RenderKit::new(core);

        let initial_params = SplattingParams {
            animation_speed: 1.0,
            splat_size: 0.8,
            particle_spread: 0.0,
            intensity: 2.0,
            particle_density: 0.4,
            brightness: 36.0,
            physics_strength: 0.5,
            trail_length: 0.0,
            trail_decay: 0.95,
            flow_strength: 1.0,
            _padding1: 0.0,
            _padding2: 0};

        // Define the multi-stage passes
        let passes = vec![
            PassDescription::new("clear_buffer", &[]), // Stage 0: Clear atomic buffer
            PassDescription::new("project_colors", &[]), // Stage 1: Project colors to 3D space
            PassDescription::new("generate_image", &[]), // Stage 2: Generate final image
        ];

        let config = ComputeShader::builder()
            .with_entry_point("clear_buffer")
            .with_multi_pass(&passes)
            .with_input_texture() // Enable input texture support
            .with_custom_uniforms::<SplattingParams>()
            .with_storage_buffer(StorageBufferSpec::new(
                "atomic_buffer",
                (core.size.width * core.size.height * 4 * 4) as u64,
            ))
            .with_workgroup_size([16, 16, 1])
            .with_texture_format(COMPUTE_TEXTURE_FORMAT_RGBA16)
            .with_label("Particle Splatting Multi-Pass")
            .build();

        let compute_shader = cuneus::compute_shader!(core, "shaders/computecolors.wgsl", config);

        compute_shader.set_custom_params(initial_params, &core.queue);

        Self {
            base,
            compute_shader,
            current_params: initial_params}
    }

    fn update(&mut self, core: &Core) {
        // Update time
        let current_time = self.base.controls.get_time(&self.base.start_time);
        let delta = 1.0 / 60.0;
        self.compute_shader
            .set_time(current_time, delta, &core.queue);

        // Update input textures for media processing
        self.base.update_current_texture(core, &core.queue);
        if let Some(texture_manager) = self.base.get_current_texture_manager() {
            self.compute_shader.update_input_texture(
                &texture_manager.view,
                &texture_manager.sampler,
                &core.device,
            );
        }
    }

    fn resize(&mut self, core: &Core) {
        self.base.default_resize(core, &mut self.compute_shader);
    }

    fn render(&mut self, core: &Core) -> Result<(), cuneus::SurfaceError> {
        let mut frame = self.base.begin_frame(core)?;

        // Handle UI and controls
        let mut params = self.current_params;
        let mut changed = false;
        let mut should_start_export = false;
        let mut export_request = self.base.export_manager.get_ui_request();
        let mut controls_request = self
            .base
            .controls
            .get_ui_request(&self.base.start_time, &core.size, self.base.fps_tracker.fps());

        let using_video_texture = self.base.using_video_texture;
        let using_hdri_texture = self.base.using_hdri_texture;
        let using_webcam_texture = self.base.using_webcam_texture;
        let video_info = self.base.get_video_info();
        let hdri_info = self.base.get_hdri_info();
        let webcam_info = self.base.get_webcam_info();
        let full_output = if self.base.key_handler.show_ui {
            self.base.render_ui(core, |ctx| {
                RenderKit::apply_default_style(ctx);

                egui::Window::new("Particle Splatting")
                    .collapsible(true)
                    .resizable(true)
                    .default_width(250.0)
                    .show(ctx, |ui| {
                        ShaderControls::render_media_panel(
                            ui,
                            &mut controls_request,
                            using_video_texture,
                            video_info,
                            using_hdri_texture,
                            hdri_info,
                            using_webcam_texture,
                            webcam_info,
                        );

                        ui.separator();

                        egui::CollapsingHeader::new("Particles")
                            .default_open(true)
                            .show(ui, |ui| {
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.particle_density, 0.1..=1.0)
                                            .text("Density"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.splat_size, 0.1..=2.0)
                                            .text("Splat Size"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.intensity, 0.1..=6.0)
                                            .text("Intensity"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.brightness, 36.0..=48.0)
                                            .text("Brightness"),
                                    )
                                    .changed();
                            });

                        egui::CollapsingHeader::new("Effects")
                            .default_open(true)
                            .show(ui, |ui| {
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.animation_speed, 0.0..=3.0)
                                            .text("Speed"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.particle_spread, 0.0..=1.0)
                                            .text("Scramble Amount"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.physics_strength, 0.0..=12.0)
                                            .text("Return Force"),
                                    )
                                    .changed();
                            });

                        egui::CollapsingHeader::new("Flow Trails")
                            .default_open(true)
                            .show(ui, |ui| {
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.trail_length, 0.0..=2.0)
                                            .text("Trail Length"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.trail_decay, 0.8..=1.0)
                                            .text("Trail Decay"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.flow_strength, 0.0..=3.0)
                                            .text("Flow Strength"),
                                    )
                                    .changed();
                            });

                        ui.separator();

                        ShaderControls::render_controls_widget(ui, &mut controls_request);

                        ui.separator();

                        should_start_export =
                            ExportManager::render_export_ui_widget(ui, &mut export_request);
                    });
            })
        } else {
            self.base.render_ui(core, |_ctx| {})
        };

        // Apply controls
        self.base.export_manager.apply_ui_request(export_request);
        if controls_request.should_clear_buffers {
            self.clear_buffers(core);
        }
        self.base.apply_media_requests(core, &controls_request);

        if should_start_export {
            self.base.export_manager.start_export();
        }

        if changed {
            self.current_params = params;
            self.compute_shader.set_custom_params(params, &core.queue);
        }
        // Handle export
        self.compute_shader.handle_export(core, &mut self.base);

        // Color projection multi-stage dispatch - run all stages every frame for animation

        // Stage 0: Clear atomic buffer (16x16 workgroups)
        self.compute_shader.dispatch_stage(&mut frame.encoder, core, 0);

        // Stage 1: Project colors to 3D space (uses input texture dimensions)
        if let Some(texture_manager) = self.base.get_current_texture_manager() {
            let input_workgroups = [
                texture_manager.texture.width().div_ceil(16),
                texture_manager.texture.height().div_ceil(16),
                1,
            ];
            self.compute_shader
                .dispatch_stage_with_workgroups(&mut frame.encoder, 1, input_workgroups);
        } else {
            // Fallback to screen size if no input texture
            self.compute_shader.dispatch_stage(&mut frame.encoder, core, 1);
        }

        // Stage 2: Generate final image (16x16 workgroups, screen size)
        self.compute_shader.dispatch_stage(&mut frame.encoder, core, 2);

        self.base.renderer.render_to_view(&mut frame.encoder, &frame.view, &self.compute_shader.get_output_texture().bind_group);

        self.base.end_frame(core, frame, full_output);

        Ok(())
    }

    fn handle_input(&mut self, core: &Core, event: &WindowEvent) -> bool {
        self.base.default_handle_input(core, event)
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    cuneus::gst::init()?;
    env_logger::init();
    let (app, event_loop) = cuneus::ShaderApp::new("Particle Splatting", 800, 600);
    app.run(event_loop, ColorProjection::init)
}
