soundview 0.3.0

use anyhow::{bail, Context, Result};
use sdl2::video::Window;
use tracing::info;
use wgpu::util::DeviceExt;

/// Pairing of a width/height
pub struct Dimensions {
    pub width: usize,
    pub height: usize,
}

/// The render configuration to be displayed
#[derive(Clone, Debug)]
pub enum Orientation {
    /// - Analyzer at top
    /// - Voiceprint scrolling down
    /// - Low frequencies along left edge
    Vertical,

    /// - Analyzer at right
    /// - Voiceprint scrolling left
    /// - Low frequencies along bottom edge
    Horizontal,
}
/// Required for using Orientation in commandline args.
impl std::fmt::Display for Orientation {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            Orientation::Vertical => write!(f, "vertical"),
            Orientation::Horizontal => write!(f, "horizontal"),
        }
    }
}
/// Optional for better parsing of Orientation in commandline args.
impl std::str::FromStr for Orientation {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Orientation> {
        match s {
            "|" | "v" | "V" | "vert" | "Vert" | "vertical" | "Vertical" => {
                Ok(Orientation::Vertical)
            }
            "-" | "h" | "H" | "horiz" | "Horiz" | "horizontal" | "Horizontal" => {
                Ok(Orientation::Horizontal)
            }
            _ => bail!("unrecognized orientation: '{}'", s),
        }
    }
}

/// RGBA/BGRA is a u32 -> 4 bytes
pub const COLOR_BYTES: usize = 4;
/// Black in RGBA/BGRA texture format (r=0,g=0,b=0,a=255)
pub const VALUE_BLACK: u32 = 255 * 16777216;

/// Use two triangles to render our rectangular texture.
/// The referenced vertices are expected to be counter-clockwise from upper left,
/// but this otherwise doesn't need to change if we toggle horiz/vert mode.
const INDICES_RECT: &[u16] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];

//            LF     MF     HF     VHF
// +1   0  1------3------5------7------9
//         |\     |\     |\     |\     |
//         | \    | \    | \    | \    |
//         |  \   |  voiceprint |  \   |
//         |   \  |   \  |   \  |   \  |
//         |    \ |    \ |    \ |    \ |
//         |     \|   analyzer \|     \|
// -1   1  0------2------4------6------8
//     IN  0     1/32   4/32  12/32    1 (texture coords: [0,1])
// OUT    -1   -12/16  -6/16  +2/16   +1 (display coords: [-1,+1])
//
// LF:
//   1/32 of texture on 2/16 of screen: zoom 4x
//   (0 to 1/32 -> 0 to 2/16)
// MF:
//   3/32 of texture on 3/16 of screen: zoom 2x
//   (1/32 to 4/32 -> 2/16 to 5/16)
// HF:
//   8/32 of texture on 4/16 of screen: zoom 1x
//   (4/32 to 12/32 -> 5/16 to 9/16)
// VHF:
//   remaining 20/32 of texture on remaining 7/16 of screen: zoom 0.7x
//   (12/32 to 32/32 -> 9/16 to 1)

/// Horizontal mode:
/// - texture is rotated 90 degrees counter-clockwise
/// - the analyzer is at the right and low frequencies are at the bottom
const VERTICES_HORIZ: &[Vertex] = &[
    Vertex {
        // 0
        position: [1.0, -1.0],  // bot right draw out
        tex_coords: [0.0, 1.0], // bot left tex in (rotate 90 ccw)
    },
    Vertex {
        // 1
        position: [-1.0, -1.0], // bot left draw out
        tex_coords: [0.0, 0.0], // top left tex in (rotate 90 ccw)
    },
    Vertex {
        // 2
        position: [1.0, -12. / 16.],
        tex_coords: [1. / 32., 1.0],
    },
    Vertex {
        // 3
        position: [-1.0, -12. / 16.],
        tex_coords: [1. / 32., 0.0],
    },
    Vertex {
        // 4
        position: [1.0, -6. / 16.],
        tex_coords: [4. / 32., 1.0],
    },
    Vertex {
        // 5
        position: [-1.0, -6. / 16.],
        tex_coords: [4. / 32., 0.0],
    },
    Vertex {
        // 6
        position: [1.0, 2. / 16.],
        tex_coords: [12. / 32., 1.0],
    },
    Vertex {
        // 7
        position: [-1.0, 2. / 16.],
        tex_coords: [12. / 32., 0.0],
    },
    Vertex {
        // 8
        position: [1.0, 1.0],   // top right draw out (y,x)
        tex_coords: [1.0, 1.0], // bot right tex in (rotate 90 ccw)
    },
    Vertex {
        // 9
        position: [-1.0, 1.0],  // top left draw out (y,x)
        tex_coords: [1.0, 0.0], // top right tex in (rotate 90 ccw)
    },
];

/// Vertical mode:
/// - texture is mirrored vertically
/// - the analyzer is at the top and low frequencies are at the left
const VERTICES_VERT: &[Vertex] = &[
    Vertex {
        // 0
        position: [-1.0, -1.0], // bot left draw out
        tex_coords: [0.0, 0.0], // top left tex in (mirror vertically)
    },
    Vertex {
        // 1
        position: [-1.0, 1.0],  // top left draw out
        tex_coords: [0.0, 1.0], // bot left tex in (mirror vertically)
    },
    Vertex {
        // 2
        position: [-12. / 16., -1.0],
        tex_coords: [1. / 32., 0.0],
    },
    Vertex {
        // 3
        position: [-12. / 16., 1.0],
        tex_coords: [1. / 32., 1.0],
    },
    Vertex {
        // 4
        position: [-6. / 16., -1.0],
        tex_coords: [4. / 32., 0.0],
    },
    Vertex {
        // 5
        position: [-6. / 16., 1.0],
        tex_coords: [4. / 32., 1.0],
    },
    Vertex {
        // 6
        position: [2. / 16., -1.0],
        tex_coords: [12. / 32., 0.0],
    },
    Vertex {
        // 7
        position: [2. / 16., 1.0],
        tex_coords: [12. / 32., 1.0],
    },
    Vertex {
        // 8
        position: [1.0, -1.0],  // bot right draw out
        tex_coords: [1.0, 0.0], // top right tex in (mirror vertically)
    },
    Vertex {
        // 9
        position: [1.0, 1.0],   // top right draw out
        tex_coords: [1.0, 1.0], // bot right tex in (mirror vertically)
    },
];

/// The definition of a vertex for rendering the texture.
/// This is sent to the shader, see shader.wgsl.
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct Vertex {
    /// The drawing output coordinates.
    /// - x/y range: [-1.0, 1.0]
    /// - x increases rightward
    /// - y increases upward
    /// e.g. top left is [0.0, 1.0], bot right is [1.0, 0.0]
    position: [f32; 2],

    /// The texture input coordinates.
    /// - x/y range: [0.0, 1.0]
    /// - x increases rightward
    /// - y increases downward
    /// e.g. top left is [0.0, 0.0], bot right is [1.0, 1.0]
    tex_coords: [f32; 2],
}

impl Vertex {
    /// Describes to wgpu how Vertex is serialized
    pub fn desc<'a>() -> wgpu::VertexBufferLayout<'a> {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &[
                // position
                wgpu::VertexAttribute {
                    offset: 0,
                    shader_location: 0,
                    format: wgpu::VertexFormat::Float32x2,
                },
                // tex_coords
                wgpu::VertexAttribute {
                    // skip past position
                    offset: std::mem::size_of::<[f32; 2]>() as wgpu::BufferAddress,
                    shader_location: 1,
                    format: wgpu::VertexFormat::Float32x2,
                },
            ],
        }
    }
}

// Note: remaining 5/16 at the top of the spectrum is discarded

/// Rendering members that need to be recreated when window dimensions change.
pub struct WgpuSizedState {
    /// The vertexes used for transforming and rendering the texture.
    /// We use a different transform for vertical vs horizontal mode.
    vertex_buffer: wgpu::Buffer,
    /// The GPU texture where voiceprint_buf and analyzer_buf are written to before displaying to the user.
    texture: wgpu::Texture,
    /// The bind group of the GPU texture and its sampler.
    /// This needs to be reinitialized whenever the texture is resized.
    bind_group: wgpu::BindGroup,
}

impl WgpuSizedState {
    /// (Re)initializes the render state to reflect a change in window size.
    /// This avoids being a method of State to keep it clear what information is needed.
    pub fn new(
        state: &WgpuState,
        texture_dims: Dimensions,
        orientation: &Orientation,
    ) -> WgpuSizedState {
        // Inputs to vertex shader - deals with rotation/scaling of the texture content.
        let vertex_buffer = state
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("vertex_buffer"),
                contents: bytemuck::cast_slice(match orientation {
                    Orientation::Vertical => VERTICES_VERT,
                    Orientation::Horizontal => VERTICES_HORIZ,
                }),
                usage: wgpu::BufferUsages::VERTEX,
            });

        // Texture that has both voiceprint and analyzer buffers written to it
        let texture = state.device.create_texture(&wgpu::TextureDescriptor {
            label: Some("buf_texture"),
            size: wgpu::Extent3d {
                width: texture_dims.width as u32,
                height: texture_dims.height as u32,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: state.preferred_format,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[state.preferred_format],
        });

        // Reinitialize this when the texture is recreated
        let bind_group = state.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("bind_group"),
            layout: &state.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(
                        &texture.create_view(&wgpu::TextureViewDescriptor::default()),
                    ),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&state.device.create_sampler(
                        &wgpu::SamplerDescriptor {
                            address_mode_u: wgpu::AddressMode::ClampToEdge, // x
                            address_mode_v: wgpu::AddressMode::ClampToEdge, // y
                            address_mode_w: wgpu::AddressMode::ClampToEdge, // z
                            mag_filter: wgpu::FilterMode::Nearest,
                            min_filter: wgpu::FilterMode::Nearest,
                            mipmap_filter: wgpu::FilterMode::Nearest,
                            ..Default::default()
                        },
                    )),
                },
            ],
        });

        WgpuSizedState {
            vertex_buffer,
            texture,
            bind_group,
        }
    }
}

/// The rendering state.
pub struct WgpuState {
    // The wgpu output components. Initialized once on startup.
    surface: wgpu::Surface<'static>,
    device: wgpu::Device,
    queue: wgpu::Queue,
    render_pipeline: wgpu::RenderPipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    index_buffer: wgpu::Buffer,
    pub preferred_format: wgpu::TextureFormat,

    /// The wgpu window configuration, with sizes updated if the window is resized
    config: wgpu::SurfaceConfiguration,
}

impl WgpuState {
    /// Sets up a visualization rendering pipeline for the provided window.
    /// Most of this inscrutable boilerplate mess is copied from here:
    /// https://github.com/Rust-SDL2/rust-sdl2/blob/master/examples/raw-window-handle-with-wgpu/main.rs
    pub async fn new(window: &Window) -> Result<Self> {
        // The instance is a handle to our GPU
        // BackendBit::PRIMARY => Vulkan + Metal + DX12 + Browser WebGPU
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::default());
        let surface = unsafe {
            let window = wgpu::SurfaceTargetUnsafe::from_window(&window)
                .context("Couldn't cast SDL2 window as wgpu window")?;
            instance
                .create_surface_unsafe(window)
                .context("Failed to create wgpu surface with SDL2 window")?
        };
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                force_fallback_adapter: false,
                compatible_surface: Some(&surface),
            })
            .await
            .context("Failed to initialize graphics adapter. Bad driver? Reboot needed?")?;

        let adapter_info = adapter.get_info();
        info!(
            "Rendering with {:?} to: {} (driver: {} {})",
            adapter_info.backend, adapter_info.name, adapter_info.driver, adapter_info.driver_info
        );

        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    label: Some("device-descriptor"),
                    required_features: wgpu::Features::empty(),
                    required_limits: wgpu::Limits::default(),
                    memory_hints: wgpu::MemoryHints::Performance,
                },
                None, // Trace path
            )
            .await
            .context("Failed to initialize graphics device. Bad driver? Reboot needed?")?;

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        multisampled: false,
                        view_dimension: wgpu::TextureViewDimension::D2,
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
            label: Some("bind_group_layout"),
        });

        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("soundview-shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()),
        });

        let formats = surface.get_capabilities(&adapter).formats;
        let preferred_format = *formats
            .first()
            .context("No formats found for graphics adapter")?;

        let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("render_pipeline"),
            layout: Some(
                &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some("render_pipeline_layout"),
                    bind_group_layouts: &[&bind_group_layout],
                    push_constant_ranges: &[],
                }),
            ),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vertex_shader_main"),
                compilation_options: wgpu::PipelineCompilationOptions::default(),
                buffers: &[Vertex::desc()],
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fragment_shader_main"),
                compilation_options: wgpu::PipelineCompilationOptions::default(),
                targets: &[Some(wgpu::ColorTargetState {
                    format: preferred_format,
                    blend: Some(wgpu::BlendState {
                        color: wgpu::BlendComponent::REPLACE,
                        alpha: wgpu::BlendComponent::REPLACE,
                    }),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleStrip,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Cw,
                cull_mode: Some(wgpu::Face::Back),
                // Setting this to anything other than Fill requires Features::NON_FILL_POLYGON_MODE
                polygon_mode: wgpu::PolygonMode::Fill,
                // Requires Features::DEPTH_CLIP_CONTROL
                unclipped_depth: false,
                // Requires Features::CONSERVATIVE_RASTERIZATION
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            // If the pipeline will be used with a multiview render pass, this
            // indicates how many array layers the attachments will have.
            multiview: None,
            // Assume cache isn't helpful since we only create the pipeline once
            cache: None,
        });

        let index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("index_buffer"),
            contents: bytemuck::cast_slice(INDICES_RECT),
            usage: wgpu::BufferUsages::INDEX,
        });

        let (width, height) = window.size();
        let config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: preferred_format,
            width,
            height,
            present_mode: wgpu::PresentMode::Fifo,
            alpha_mode: wgpu::CompositeAlphaMode::Auto,
            view_formats: (&[preferred_format]).to_vec(),
            desired_maximum_frame_latency: 0,
        };
        // Call configure() once up-front, needed for tiling WMs like sway
        surface.configure(&device, &config);

        Ok(WgpuState {
            surface,
            device,
            queue,
            render_pipeline,
            bind_group_layout,
            index_buffer,
            preferred_format,

            config,
        })
    }

    /// Updates the render size to reflect a change in window size
    pub fn resize(&mut self, new_size: Option<Dimensions>) {
        // Update width/height in wgpu config.
        // This is separate from sized_state to avoid rebuilding the whole config.
        if let Some(new_size) = new_size {
            if new_size.width > 0 && new_size.height > 0 {
                self.config.width = new_size.width as u32;
                self.config.height = new_size.height as u32;
            }
        }
        // Update surface with the change to config.
        self.surface.configure(&self.device, &self.config);
    }

    pub fn window_dims(&self) -> Dimensions {
        Dimensions {
            width: self.config.width as usize,
            height: self.config.height as usize,
        }
    }

    /// Returns the underlying surface texture for rendering.
    pub fn surface_texture(&mut self) -> Result<wgpu::SurfaceTexture, wgpu::SurfaceError> {
        self.surface.get_current_texture()
    }

    pub fn write_texture(
        &mut self,
        sized_state: &WgpuSizedState,
        buf: &[u8],
        image_offset: usize,
        image_dims: Dimensions,
        extent_dims: Dimensions,
        origin_y: usize,
    ) {
        self.queue.write_texture(
            wgpu::ImageCopyTexture {
                aspect: wgpu::TextureAspect::All,
                texture: &sized_state.texture,
                mip_level: 0,
                origin: match origin_y {
                    0 => wgpu::Origin3d::ZERO,
                    y => wgpu::Origin3d {
                        x: 0,
                        y: y as u32,
                        z: 0,
                    },
                },
            },
            buf,
            wgpu::ImageDataLayout {
                offset: (COLOR_BYTES * image_offset) as u64,
                bytes_per_row: Some((COLOR_BYTES * image_dims.width) as u32),
                rows_per_image: Some(image_dims.height as u32),
            },
            wgpu::Extent3d {
                width: extent_dims.width as u32,
                height: extent_dims.height as u32,
                depth_or_array_layers: 1,
            },
        );
    }

    /// Re-renders the display:
    /// - Collects a batch of audio frequency data from recv_processed
    /// - For each audio entry in the batch, adds it to the voiceprint to be rendered
    /// - For the last audio entry in the batch, updates the analyzer to be rendered
    /// - Writes the resulting buffers to a texture which is then rotated/scaled and rendered to the output
    pub fn render(
        &mut self,
        sized_state: &WgpuSizedState,
        output: wgpu::SurfaceTexture,
    ) -> Result<()> {
        let view = output
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default());
        let mut command = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Render Encoder"),
            });

        // Render the texture, with vertex_buffer controlling rotation/mirroring to match the display mode.
        {
            let mut render_pass = command.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("render_pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color {
                            r: 0.1,
                            g: 0.2,
                            b: 0.3,
                            a: 1.0,
                        }),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            });
            render_pass.set_pipeline(&self.render_pipeline);
            render_pass.set_bind_group(0, &sized_state.bind_group, &[]);
            render_pass.set_vertex_buffer(0, sized_state.vertex_buffer.slice(..));
            render_pass.set_index_buffer(self.index_buffer.slice(..), wgpu::IndexFormat::Uint16);
            render_pass.draw_indexed(0..INDICES_RECT.len() as u32, 0, 0..1);
        }

        self.queue.submit(std::iter::once(command.finish()));
        output.present();

        Ok(())
    }
}