spectral-render/src/renderer.rs

use std::sync::Arc;

use wgpu;
use bytemuck::{Pod, Zeroable};

use spectral_core::Terminal;
use spectral_font::{GlyphAtlas, GlyphKey};

use crate::shaders::{BG_VERTEX, BG_FRAGMENT, TEXT_VERTEX, TEXT_FRAGMENT};

#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
pub struct BgInstance {
    pub col: u32, pub row: u32,
    pub bg_r: u32, pub bg_g: u32, pub bg_b: u32, pub bg_a: u32,
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
pub struct TextInstance {
    pub col: f32, pub row: f32,
    pub bearing_x: f32, pub bearing_y: f32,
    pub glyph_w: f32, pub glyph_h: f32,
    pub atlas_x: f32, pub atlas_y: f32,
    pub atlas_w: f32, pub atlas_h: f32,
    pub fg_r: u32, pub fg_g: u32, pub fg_b: u32,
    pub flags: u32, pub thickening: f32,
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
pub struct Uniforms {
    pub cell_width: f32, pub cell_height: f32, pub atlas_size: f32,
    pub screen_cols: u32, pub screen_rows: u32,
    _pad: u32, _pad2: u32,
}

pub struct Renderer {
    pub device: wgpu::Device,
    pub queue: wgpu::Queue,
    pub surface: wgpu::Surface<'static>,
    pub surface_config: wgpu::SurfaceConfiguration,
    pub bg_pipeline: wgpu::RenderPipeline,
    pub text_pipeline: wgpu::RenderPipeline,
    pub bg_instance_buf: wgpu::Buffer,
    pub text_instance_buf: wgpu::Buffer,
    pub uniform_buf: wgpu::Buffer,
    pub atlas_texture: wgpu::Texture,
    pub atlas_view: wgpu::TextureView,
    pub atlas_sampler: wgpu::Sampler,
    pub bg_bind_group: wgpu::BindGroup,
    pub text_uniform_bind_group: wgpu::BindGroup,
    pub text_texture_bind_group: wgpu::BindGroup,
    pub glyph_atlas: GlyphAtlas,
    pub bg_instances: Vec<BgInstance>,
    pub text_instances: Vec<TextInstance>,
    pub needs_atlas_upload: bool,
}

impl Renderer {
    pub async fn new(window: Arc<winit::window::Window>, term_cols: u32, term_rows: u32, atlas_size: u16) -> Self {
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: wgpu::Backends::all(), ..Default::default()
        });
        let surface = instance.create_surface(window).unwrap();
        let adapter = instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::HighPerformance,
            compatible_surface: Some(&surface), force_fallback_adapter: false,
        }).await.unwrap();

        let (device, queue) = adapter.request_device(&wgpu::DeviceDescriptor {
            required_features: wgpu::Features::empty(),
            required_limits: wgpu::Limits::default(),
            label: Some("Spectral Device"),
            memory_hints: wgpu::MemoryHints::Performance,
        }, None).await.unwrap();

        let surface_caps = surface.get_capabilities(&adapter);
        let surface_format = surface_caps.formats.iter().copied()
            .find(|f| f.is_srgb()).unwrap_or(surface_caps.formats[0]);

        let config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: surface_format, width: 1280, height: 720,
            present_mode: wgpu::PresentMode::AutoNoVsync,
            desired_maximum_frame_latency: 1,
            alpha_mode: wgpu::CompositeAlphaMode::Auto,
            view_formats: vec![],
        };
        surface.configure(&device, &config);

        let glyph_atlas = GlyphAtlas::new(atlas_size);
        let max_instances = (term_cols * term_rows) as usize;
        let atlas_size_px = 2048u32;

        let bg_instance_buf = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("bg_instances"),
            size: (max_instances * std::mem::size_of::<BgInstance>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let text_instance_buf = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("text_instances"),
            size: (max_instances * std::mem::size_of::<TextInstance>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let uniform_buf = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("uniforms"),
            size: std::mem::size_of::<Uniforms>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let atlas_texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("atlas_texture"),
            size: wgpu::Extent3d { width: atlas_size_px, height: atlas_size_px, depth_or_array_layers: 1 },
            mip_level_count: 1, sample_count: 1, dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8UnormSrgb,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });

        let atlas_view = atlas_texture.create_view(&wgpu::TextureViewDescriptor::default());
        let atlas_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("atlas_sampler"),
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            ..Default::default()
        });

        let uniform_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("uniform_layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0, visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false, min_binding_size: None,
                    }, count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1, visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false, min_binding_size: None,
                    }, count: None,
                },
            ],
        });

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

        let bg_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("bg_pipeline_layout"),
            bind_group_layouts: &[&uniform_layout], push_constant_ranges: &[],
        });

        let text_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("text_pipeline_layout"),
            bind_group_layouts: &[&uniform_layout, &text_texture_layout], push_constant_ranges: &[],
        });

        let bg_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("bg_shader"), source: wgpu::ShaderSource::Wgsl(BG_VERTEX.into()),
        });
        let bg_frag = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("bg_frag"), source: wgpu::ShaderSource::Wgsl(BG_FRAGMENT.into()),
        });
        let text_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("text_shader"), source: wgpu::ShaderSource::Wgsl(TEXT_VERTEX.into()),
        });
        let text_frag = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("text_frag"), source: wgpu::ShaderSource::Wgsl(TEXT_FRAGMENT.into()),
        });

        let bg_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("bg_pipeline"),
            layout: Some(&bg_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &bg_shader, entry_point: Some("main"),
                buffers: &[], compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &bg_frag, entry_point: Some("main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_format, blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState::default(), depth_stencil: None,
            multisample: wgpu::MultisampleState::default(), multiview: None, cache: None,
        });

        let text_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("text_pipeline"),
            layout: Some(&text_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &text_shader, entry_point: Some("main"),
                buffers: &[], compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &text_frag, entry_point: Some("main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_format, blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState::default(), depth_stencil: None,
            multisample: wgpu::MultisampleState::default(), multiview: None, cache: None,
        });

        let bg_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("bg_bind_group"), layout: &uniform_layout,
            entries: &[
                wgpu::BindGroupEntry { binding: 0, resource: bg_instance_buf.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 1, resource: uniform_buf.as_entire_binding() },
            ],
        });

        let text_uniform_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("text_uniform_bind_group"), layout: &uniform_layout,
            entries: &[
                wgpu::BindGroupEntry { binding: 0, resource: text_instance_buf.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 1, resource: uniform_buf.as_entire_binding() },
            ],
        });

        let text_texture_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("text_texture_bind_group"), layout: &text_texture_layout,
            entries: &[
                wgpu::BindGroupEntry { binding: 2, resource: wgpu::BindingResource::Sampler(&atlas_sampler) },
                wgpu::BindGroupEntry { binding: 3, resource: wgpu::BindingResource::TextureView(&atlas_view) },
            ],
        });

        Self {
            device, queue, surface, surface_config: config,
            bg_pipeline, text_pipeline,
            bg_instance_buf, text_instance_buf, uniform_buf,
            atlas_texture, atlas_view, atlas_sampler,
            bg_bind_group, text_uniform_bind_group, text_texture_bind_group,
            glyph_atlas,
            bg_instances: Vec::with_capacity(max_instances),
            text_instances: Vec::with_capacity(max_instances),
            needs_atlas_upload: false,
        }
    }

    pub fn resize(&mut self, width: u32, height: u32) {
        if width == 0 || height == 0 { return; }
        self.surface_config.width = width;
        self.surface_config.height = height;
        self.surface.configure(&self.device, &self.surface_config);
    }

    pub fn update_atlas(&mut self) {
        if !self.needs_atlas_upload { return; }
        for (page_idx, page_data) in self.glyph_atlas.texture_data.iter().enumerate() {
            let page_size = self.glyph_atlas.page_size as u32;
            self.queue.write_texture(
                wgpu::ImageCopyTexture {
                    texture: &self.atlas_texture, mip_level: 0,
                    origin: wgpu::Origin3d { x: 0, y: 0, z: page_idx as u32 },
                    aspect: wgpu::TextureAspect::All,
                },
                page_data,
                wgpu::ImageDataLayout { offset: 0, bytes_per_row: Some(page_size * 4), rows_per_image: Some(page_size) },
                wgpu::Extent3d { width: page_size, height: page_size, depth_or_array_layers: 1 },
            );
        }
        self.needs_atlas_upload = false;
    }

    pub fn update_instances(&mut self, terminal: &Terminal) {
        self.bg_instances.clear();
        self.text_instances.clear();
        let cell_width = 8.0f32;
        let cell_height = 16.0f32;
        let thickening = if terminal.config.font_thicken { terminal.config.font_thicken_strength } else { 0.0 };

        for (row_idx, line) in terminal.grid.lines.iter().enumerate() {
            for (col_idx, cell) in line.cells.iter().enumerate() {
                self.bg_instances.push(BgInstance {
                    col: col_idx as u32, row: row_idx as u32,
                    bg_r: ((cell.bg >> 16) & 0xFF) as u32,
                    bg_g: ((cell.bg >> 8) & 0xFF) as u32,
                    bg_b: (cell.bg & 0xFF) as u32,
                    bg_a: ((cell.bg >> 24) & 0xFF) as u32,
                });
                if cell.ch == ' ' || cell.ch == '\0' { continue; }
                let key = GlyphKey { font_id: 0, glyph_id: cell.ch as u16, size_bucket: (cell_height * 4.0) as u16, style_flags: 0 };
                if let Some((_page, slot)) = self.glyph_atlas.get(&key) {
                    self.text_instances.push(TextInstance {
                        col: col_idx as f32, row: row_idx as f32,
                        bearing_x: 0.0, bearing_y: 0.0,
                        glyph_w: slot.w as f32, glyph_h: slot.h as f32,
                        atlas_x: slot.x as f32, atlas_y: slot.y as f32,
                        atlas_w: slot.w as f32, atlas_h: slot.h as f32,
                        fg_r: ((cell.fg >> 16) & 0xFF) as u32,
                        fg_g: ((cell.fg >> 8) & 0xFF) as u32,
                        fg_b: (cell.fg & 0xFF) as u32,
                        flags: cell.flags.bits() as u32,
                        thickening,
                    });
                }
            }
        }

        if !self.bg_instances.is_empty() {
            self.queue.write_buffer(&self.bg_instance_buf, 0, bytemuck::cast_slice(&self.bg_instances));
        }
        if !self.text_instances.is_empty() {
            self.queue.write_buffer(&self.text_instance_buf, 0, bytemuck::cast_slice(&self.text_instances));
        }

        let uniforms = Uniforms {
            cell_width, cell_height,
            atlas_size: self.glyph_atlas.page_size as f32,
            screen_cols: terminal.grid.cols as u32,
            screen_rows: terminal.grid.rows as u32,
            _pad: 0, _pad2: 0,
        };
        self.queue.write_buffer(&self.uniform_buf, 0, bytemuck::bytes_of(&uniforms));
    }

    pub fn render(&mut self) {
        let output = match self.surface.get_current_texture() { Ok(t) => t, Err(_) => return };
        let view = output.texture.create_view(&wgpu::TextureViewDescriptor::default());
        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("render_encoder") });

        {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("bg_pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &view, resolve_target: None,
                    ops: wgpu::Operations { load: wgpu::LoadOp::Clear(wgpu::Color::BLACK), store: wgpu::StoreOp::Store },
                })],
                depth_stencil_attachment: None, timestamp_writes: None, occlusion_query_set: None,
            });
            pass.set_pipeline(&self.bg_pipeline);
            pass.set_bind_group(0, &self.bg_bind_group, &[]);
            pass.draw(0..6, 0..self.bg_instances.len() as u32);
        }

        {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("text_pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &view, resolve_target: None,
                    ops: wgpu::Operations { load: wgpu::LoadOp::Load, store: wgpu::StoreOp::Store },
                })],
                depth_stencil_attachment: None, timestamp_writes: None, occlusion_query_set: None,
            });
            pass.set_pipeline(&self.text_pipeline);
            pass.set_bind_group(0, &self.text_uniform_bind_group, &[]);
            pass.set_bind_group(1, &self.text_texture_bind_group, &[]);
            pass.draw(0..6, 0..self.text_instances.len() as u32);
        }

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