gled 2.29.0

use crate::{
    audio::sound_data::SoundData,
    pipeline::{
        gpu::{Gpu, OutputContext},
        group::Group,
        output_mix::OutputMix,
        render_to_output::{RenderToOutput, positions::Positions},
        renderer_callback::RendererCallback,
    },
    storage::{
        Animation, AssetId, Palette,
        animation::{
            config::AnimationConfig,
            renderer::{AnimationRenderer, EffectUniforms},
        },
        asset::scene::effect::ShaderCodeComplete,
        collections::Collections,
    },
    ui::action::UiAction,
    wgpu_render_state,
};
use arboard::{Clipboard, ImageData};
use egui::TextureId;
use kanal::bounded;
use wgpu::{CommandEncoder, MapMode, PollType, Queue};

#[derive(Debug, Default, PartialEq)]
pub struct EffectState {
    /// Progress in current beat.
    pub beat_progression: f32,
    /// Beats per minute
    pub beats_per_minute: f32,
    /// Current displayed framerate
    pub framerate: f32,
    /// Opacity of animation: 0.0 -> 1.0
    pub opacity: f32,
    /// Color shift in full circles. (0.5 = 180 degrees)
    pub color_shift: f32,
    /// In 2^n of bpm
    pub speed_exponent: i32,

    pub animation_config: AnimationConfig,
    pub sent_group: Option<Group>,
    /// Animation the GPU pipeline was built for. `None` until built through
    /// `Effect::ensure_*_pipeline`, after a direct `set_*_shader_code`, or after
    /// `Effect::invalidate_pipeline`.
    pub(crate) pipeline_animation: Option<Option<AssetId<Animation>>>,

    /// Thumbnail pipeline on the UI device.
    matrix: Option<MatrixPipeline>,
    /// LED output pipeline on the render thread's output device.
    output: Option<OutputPipeline>,
}

/// Renders the effect into a texture the UI shows as a thumbnail.
#[derive(Debug, PartialEq)]
struct MatrixPipeline {
    uniforms: EffectUniforms,
    renderer: AnimationRenderer,
    texture_id: OwnedTextureId,
}

/// Computes the effect at the LED positions and mixes it into the output buffer.
#[derive(Debug, PartialEq)]
struct OutputPipeline {
    uniforms: EffectUniforms,
    render_to_output: RenderToOutput,
    output_mix: OutputMix,
}

impl EffectState {
    /// Build the thumbnail pipeline on the UI device. UI thread only.
    pub fn set_matrix_shader_code(&mut self, gpu: &Gpu, shader_code: &ShaderCodeComplete) {
        let uniforms = EffectUniforms::new(&gpu.device);
        let renderer =
            AnimationRenderer::new(&gpu.device, &shader_code.fragment_render_matrix, &uniforms);
        let texture_id = OwnedTextureId(
            wgpu_render_state()
                .renderer
                .write()
                .register_native_texture(&gpu.device, renderer.view(), wgpu::FilterMode::Nearest),
        );
        self.matrix = Some(MatrixPipeline {
            uniforms,
            renderer,
            texture_id,
        });
        self.pipeline_animation = None;
    }

    /// Build the LED output pipeline on the output device. Render thread only.
    pub fn set_output_shader_code(
        &mut self,
        context: &OutputContext,
        shader_code: &ShaderCodeComplete,
    ) {
        let device = &context.gpu.device;
        let uniforms = EffectUniforms::new(device);
        let render_to_output = RenderToOutput::init(
            device,
            &shader_code.compute_render_output,
            uniforms.buffer(),
        );
        let output_mix = OutputMix::new(
            device,
            render_to_output.output_buffer(),
            &context.output_buffer,
        );
        self.output = Some(OutputPipeline {
            uniforms,
            render_to_output,
            output_mix,
        });
        self.pipeline_animation = None;
        // The freshly created pipeline has an empty positions buffer. Force the
        // next `prepare()` to upload positions, otherwise it would mix zeros
        // into the output.
        self.sent_group = None;
    }

    pub fn has_matrix_pipeline(&self) -> bool {
        self.matrix.is_some()
    }

    pub fn has_output_pipeline(&self) -> bool {
        self.output.is_some()
    }

    /// Resend positions to gpu
    pub fn send_positions(&mut self) {
        self.sent_group.take();
    }

    pub fn texture_id(&self) -> Option<TextureId> {
        self.matrix.as_ref().map(|matrix| matrix.texture_id.0)
    }

    /// Upload LED positions to the output pipeline, if there is one.
    pub fn set_positions(&self, queue: &Queue, positions: Positions) {
        if let Some(output) = self.output.as_ref() {
            output.render_to_output.set_positions(queue, positions);
        }
    }

    /// Write palette and state into the uniforms of every pipeline this state has.
    pub fn write_uniforms(
        &self,
        queue: &Queue,
        beat_progression: f32,
        palette: Option<Palette>,
        collections: &Collections,
        sound_data: &SoundData,
    ) {
        let uniforms = self
            .matrix
            .iter()
            .map(|matrix| &matrix.uniforms)
            .chain(self.output.iter().map(|output| &output.uniforms));
        for uniforms in uniforms {
            uniforms.write(
                queue,
                self,
                beat_progression,
                palette.clone(),
                collections,
                sound_data,
            );
        }
    }

    #[cfg_attr(feature = "profiling", profiling::function)]
    pub fn render(&self, encoder: &mut CommandEncoder, render_matrix: bool, send_output: bool) {
        if render_matrix && let Some(matrix) = self.matrix.as_ref() {
            matrix.renderer.render(encoder);
        }
        if send_output && let Some(output) = self.output.as_ref() {
            output.render_to_output.run(encoder);
            output.output_mix.run(encoder);
        }
    }

    pub fn write_data(
        &self,
        beat_progression: f32,
        data: &mut [u8],
        collections: &Collections,
        sound_data: &SoundData,
    ) {
        data[0..4].copy_from_slice(&rand::random::<f32>().to_le_bytes());
        data[4..8].copy_from_slice(&self.beat_progression.to_le_bytes());
        data[8..12].copy_from_slice(&self.beats_per_minute.to_le_bytes());
        data[12..16].copy_from_slice(&self.framerate.to_le_bytes());
        data[16..20].copy_from_slice(&self.opacity.to_le_bytes());
        data[20..24].copy_from_slice(&self.color_shift.to_le_bytes());
        data[24..28].copy_from_slice(&2f32.powi(self.speed_exponent).to_le_bytes());
        self.animation_config.write_data(
            &mut data[28..28 + AnimationConfig::size()],
            beat_progression,
            collections,
            sound_data,
        );

        // Write FFT data (256 frequency bins = 1024 bytes)
        let fft_data = sound_data.fft_data_u8();
        let fft_offset = 28 + AnimationConfig::size();
        data[fft_offset..fft_offset + fft_data.len()].copy_from_slice(&fft_data);
    }

    pub const fn size() -> usize {
        const SIZE: usize = 28 + AnimationConfig::size() + 256 * 4; // 64 + 1024 = 1088 bytes
        static_assertions::const_assert_eq!(SIZE % 16, 0);
        SIZE
    }

    pub fn copy_rendered_image_to_clipboard(&self) {
        let Some(renderer) = self.matrix.as_ref().map(|matrix| &matrix.renderer) else {
            return;
        };

        let texture = renderer.texture();
        let texture_size = texture.size();
        let buffer_size = (texture_size.width * texture_size.height * 4) as wgpu::BufferAddress;
        let buffer = wgpu_render_state()
            .device
            .create_buffer(&wgpu::BufferDescriptor {
                label: Some("Texture Buffer"),
                size: buffer_size,
                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
                mapped_at_creation: false,
            });
        let mut encoder =
            wgpu_render_state()
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("Copy Texture to Buffer"),
                });
        encoder.copy_texture_to_buffer(
            wgpu::TexelCopyTextureInfo {
                texture: renderer.texture(),
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::TexelCopyBufferInfo {
                buffer: &buffer,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(texture_size.width * 4),
                    rows_per_image: Some(texture_size.height),
                },
            },
            texture_size,
        );

        RendererCallback::add(encoder.finish());

        let buffer_slice = buffer.slice(..);
        let (tx, rx) = bounded(0);
        buffer_slice.map_async(MapMode::Read, move |v| {
            tx.send(v).expect("Could not send on oneshot sender")
        });
        wgpu_render_state()
            .device
            .poll(PollType::wait_indefinitely())
            .expect("Could not poll device");
        rx.recv()
            .expect("Could not receive on gpu rx")
            .expect("Error receiving answer to output_data map on gpu");
        let mut data = buffer_slice
            .get_mapped_range()
            .expect("Could not get mapped range of effect state buffer")
            .to_vec();
        buffer.unmap();

        std::thread::Builder::new()
            .name("gled:clipboard:copy".to_string())
            .spawn(move || {
                #[cfg(feature = "profiling")]
                profiling::register_thread!("copy_image_to_clipboard");

                data.chunks_exact_mut(4).for_each(|pixel| {
                    pixel.swap(0, 2);
                    pixel[3] = 255;
                });

                if let Ok(mut clipboard) = Clipboard::new() {
                    if let Err(err) = clipboard.set_image(ImageData {
                        bytes: data.into(),
                        width: texture_size.width as usize,
                        height: texture_size.height as usize,
                    }) {
                        UiAction::Error(format!(
                            "Could not copy rendered image to clipboard: {err}"
                        ))
                        .enqueue();
                    } else {
                        tracing::info!("Copied rendered image to clipboard");
                    }
                }
            })
            .ok();
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct OwnedTextureId(pub TextureId);

impl Drop for OwnedTextureId {
    fn drop(&mut self) {
        wgpu_render_state().renderer.write().free_texture(&self.0)
    }
}