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

cuneus::uniform_params! {
    struct CliffordParams {
        a: f32, b: f32, c: f32, d: f32,
        motion_speed: f32,
        rotation_x: f32, rotation_y: f32,
        brightness: f32,
        scale: f32,
        dof_amount: f32, dof_focal_dist: f32,
        dispersion: f32, warp: f32, _pad3: f32,
        wl_center: f32, wl_spread: f32,
        symmetry: f32,
        _pad0: f32, _pad1: f32, _pad2: f32,
    }
}

struct CliffordShader {
    base: RenderKit,
    compute_shader: ComputeShader,
    current_params: CliffordParams}

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

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

        let initial_params = CliffordParams {
            a: -2.0, b: -2.0, c: -1.2, d: 2.0,
            motion_speed: 1.0,
            rotation_x: 0.0, rotation_y: 0.0,
            brightness: 2.5,
            scale: 0.6,
            dof_amount: 0.2, dof_focal_dist: 0.5,
            dispersion: 0.25, warp: 3.0, _pad3: 0.0,
            wl_center: 550.0, wl_spread: 300.0,
            symmetry: 0.15,
            _pad0: 0.0, _pad1: 0.0, _pad2: 0.0,
        };

        let mut config = ComputeShader::builder()
            .with_entry_point("Splat")
            .with_custom_uniforms::<CliffordParams>()
            .with_atomic_buffer(3)
            .with_workgroup_size([16, 16, 1])
            .with_texture_format(COMPUTE_TEXTURE_FORMAT_RGBA16)
            .with_label("Clifford Attractor Unified")
            .build();

        config.entry_points.push("main_image".to_string());

        let compute_shader = cuneus::compute_shader!(core, "shaders/cliffordcompute.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) {
        // Handle export with custom dispatch pattern for cliffordcompute
        self.compute_shader.handle_export_dispatch(
            core,
            &mut self.base,
            |shader, encoder, core| {
                shader.dispatch_stage_with_workgroups(encoder, 0, [3072, 1, 1]);
                shader.dispatch_stage(encoder, core, 1);
            },
        );
    }

    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)?;

        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 full_output = if self.base.key_handler.show_ui {
            self.base.render_ui(core, |ctx| {
                RenderKit::apply_default_style(ctx);

                egui::Window::new("Clifford Attractor")
                    .collapsible(true)
                    .resizable(true)
                    .default_width(250.0)
                    .show(ctx, |ui| {
                        egui::CollapsingHeader::new("Attractor Parameters")
                            .default_open(false)
                            .show(ui, |ui| {
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.a, -3.0..=3.0)
                                            .text("Parameter A"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.b, -3.0..=3.0)
                                            .text("Parameter B"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.c, -3.0..=3.0)
                                            .text("Parameter C"),
                                    )
                                    .changed();
                                changed |= ui
                                    .add(
                                        egui::Slider::new(&mut params.d, -3.0..=3.0)
                                            .text("Parameter D"),
                                    )
                                    .changed();
                                ui.separator();
                                ui.label("Presets:");
                                ui.horizontal_wrapped(|ui| {
                                    if ui.button("Nebula").clicked() {
                                        params.a = -2.0; params.b = -2.0; params.c = -1.2; params.d = 2.0;
                                        changed = true;
                                    }
                                    if ui.button("Ribbon").clicked() {
                                        params.a = 1.6; params.b = -2.3; params.c = 1.2; params.d = -1.7;
                                        changed = true;
                                    }
                                    if ui.button("Shell").clicked() {
                                        params.a = -1.8; params.b = 1.8; params.c = -1.9; params.d = -0.4;
                                        changed = true;
                                    }
                                });
                            });

                        egui::CollapsingHeader::new("Camera & Motion")
                            .default_open(false)
                            .show(ui, |ui| {
                                changed |= ui.add(egui::Slider::new(&mut params.rotation_x, -1.0..=1.0).text("Orbit X")).changed();
                                changed |= ui.add(egui::Slider::new(&mut params.rotation_y, -1.0..=1.0).text("Orbit Y")).changed();
                                changed |= ui.add(egui::Slider::new(&mut params.scale, 0.1..=2.0).text("Scale")).changed();
                                changed |= ui.add(egui::Slider::new(&mut params.motion_speed, 0.0..=3.0).text("Animation Speed")).changed();
                            });

                        egui::CollapsingHeader::new("Spectral Color")
                            .default_open(true)
                            .show(ui, |ui| {
                                ui.label("Color comes from a wavelength band, not RGB.");
                                changed |= ui.add(egui::Slider::new(&mut params.wl_center, 390.0..=700.0).text("Wavelength (nm)")).changed();
                                changed |= ui.add(egui::Slider::new(&mut params.wl_spread, 0.0..=320.0).text("Spread")).changed();
                                changed |= ui.add(egui::Slider::new(&mut params.symmetry, 0.0..=1.0).text("Spectral Split")).changed();
                                changed |= ui.add(egui::Slider::new(&mut params.brightness, 0.1..=6.0).logarithmic(true).text("Brightness")).changed();
                            });

                        egui::CollapsingHeader::new("Prism Optics")
                            .default_open(false)
                            .show(ui, |ui| {
                                changed |= ui.add(egui::Slider::new(&mut params.dispersion, 0.0..=0.5).text("Dispersion")).changed();
                                changed |= ui.add(egui::Slider::new(&mut params.warp, 0.0..=10.0).text("Refraction Warp")).changed();
                                ui.separator();
                                changed |= ui.add(egui::Slider::new(&mut params.dof_amount, 0.0..=3.0).text("DOF Amount")).changed();
                                changed |= ui.add(egui::Slider::new(&mut params.dof_focal_dist, 0.0..=1.0).text("Focal Distance")).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| {})
        };

        self.base.export_manager.apply_ui_request(export_request);
        if controls_request.should_clear_buffers {
            self.clear_buffers(core);
        }
        self.base.apply_control_request(controls_request);

        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);

        if changed {
            self.current_params = params;
            self.compute_shader.set_custom_params(params, &core.queue);
        }

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

        // Stage 0: Generate and splat particles (workgroup size [256, 1, 1])
        self.compute_shader
            .dispatch_stage_with_workgroups(&mut frame.encoder, 0, [3072, 1, 1]);

        // Stage 1: Render to screen (workgroup size [16, 16, 1])
        self.compute_shader.dispatch_stage(&mut frame.encoder, core, 1);

        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>> {
    env_logger::init();
    let (app, event_loop) = cuneus::ShaderApp::new("Clifford", 800, 600);

    app.run(event_loop, CliffordShader::init)
}
