codecraft 0.1.2

use std::sync::Arc;

use winit::window::Window;

#[derive(Debug)]
pub enum FrameError {
    Skip,
    Outdated,
    Lost,
    Fatal,
}

enum Drawn {
    Swapchain(wgpu::SurfaceTexture),
    Offscreen(wgpu::Texture),
}

/// An acquired image, ready to be rendered into.
pub struct Frame {
    drawn: Drawn,
    pub view: wgpu::TextureView,
}

impl Frame {
    /// The image being drawn into, for copying it back out.
    pub fn texture(&self) -> &wgpu::Texture {
        match &self.drawn {
            Drawn::Swapchain(frame) => &frame.texture,
            Drawn::Offscreen(texture) => texture,
        }
    }
}

enum Target {
    Window {
        surface: wgpu::Surface<'static>,
        config: wgpu::SurfaceConfiguration,
    },
    Offscreen {
        format: wgpu::TextureFormat,
        width: u32,
        height: u32,
        in_flight: Option<wgpu::SubmissionIndex>,
    },
}

/// Bare device bootstrap shared by every renderer in this workspace.
pub struct GpuContext {
    pub device: wgpu::Device,
    pub queue: wgpu::Queue,
    target: Target,
    submitted: Option<wgpu::SubmissionIndex>,
}

impl GpuContext {
    pub fn new(window: Arc<Window>) -> Self {
        let size = window.inner_size();

        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: wgpu::Backends::PRIMARY,
            ..wgpu::InstanceDescriptor::new_without_display_handle()
        });

        let surface = instance.create_surface(window).unwrap();

        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::default(),
            compatible_surface: Some(&surface),
            ..Default::default()
        }))
        .expect("no suitable GPU adapter found");

        let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
            label: Some("device"),
            required_features: wgpu::Features::empty(),
            required_limits: wgpu::Limits::default(),
            ..Default::default()
        }))
        .expect("failed to create device");

        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 {
            // COPY_SRC so a frame can be read back for a screenshot.
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
            format: surface_format,
            width: size.width.max(1),
            height: size.height.max(1),
            present_mode: surface_caps.present_modes[0],
            alpha_mode: surface_caps.alpha_modes[0],
            view_formats: vec![],
            desired_maximum_frame_latency: 2,
            color_space: wgpu::SurfaceColorSpace::Auto,
        };
        surface.configure(&device, &config);

        Self {
            device,
            queue,
            target: Target::Window { surface, config },
            submitted: None,
        }
    }

    /// A context with no window: frames are drawn into a plain texture.
    pub fn offscreen(width: u32, height: u32) -> Self {
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: wgpu::Backends::PRIMARY,
            ..wgpu::InstanceDescriptor::new_without_display_handle()
        });

        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::default(),
            compatible_surface: None,
            ..Default::default()
        }))
        .expect("no suitable GPU adapter found");

        let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
            label: Some("device"),
            required_features: wgpu::Features::empty(),
            required_limits: wgpu::Limits::default(),
            ..Default::default()
        }))
        .expect("failed to create device");

        Self {
            device,
            queue,
            target: Target::Offscreen {
                // An sRGB target, to match what a surface would give.
                format: wgpu::TextureFormat::Rgba8UnormSrgb,
                width: width.max(1),
                height: height.max(1),
                in_flight: None,
            },
            submitted: None,
        }
    }

    pub fn width(&self) -> u32 {
        match &self.target {
            Target::Window { config, .. } => config.width,
            Target::Offscreen { width, .. } => *width,
        }
    }

    pub fn height(&self) -> u32 {
        match &self.target {
            Target::Window { config, .. } => config.height,
            Target::Offscreen { height, .. } => *height,
        }
    }

    pub fn format(&self) -> wgpu::TextureFormat {
        match &self.target {
            Target::Window { config, .. } => config.format,
            Target::Offscreen { format, .. } => *format,
        }
    }

    pub fn resize(&mut self, width: u32, height: u32) {
        if width == 0 || height == 0 {
            return;
        }
        match &mut self.target {
            Target::Window { surface, config } => {
                config.width = width;
                config.height = height;
                surface.configure(&self.device, config);
            }
            Target::Offscreen {
                width: w,
                height: h,
                ..
            } => {
                *w = width;
                *h = height;
            }
        }
    }

    pub fn begin_frame(&self) -> Result<Frame, FrameError> {
        let Target::Window { surface, .. } = &self.target else {
            // Offscreen: a fresh texture each frame, so a screenshot cannot race the frame after it.
            let texture = self.device.create_texture(&wgpu::TextureDescriptor {
                label: Some("offscreen frame"),
                size: wgpu::Extent3d {
                    width: self.width(),
                    height: self.height(),
                    depth_or_array_layers: 1,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: self.format(),
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
                view_formats: &[],
            });
            let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
            return Ok(Frame {
                drawn: Drawn::Offscreen(texture),
                view,
            });
        };

        let texture = match surface.get_current_texture() {
            wgpu::CurrentSurfaceTexture::Success(texture) => texture,
            wgpu::CurrentSurfaceTexture::Suboptimal(texture) => texture,
            wgpu::CurrentSurfaceTexture::Timeout | wgpu::CurrentSurfaceTexture::Occluded => {
                return Err(FrameError::Skip);
            }
            wgpu::CurrentSurfaceTexture::Outdated => return Err(FrameError::Outdated),
            wgpu::CurrentSurfaceTexture::Lost => return Err(FrameError::Lost),
            wgpu::CurrentSurfaceTexture::Validation => return Err(FrameError::Fatal),
        };
        let view = texture
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default());
        Ok(Frame {
            drawn: Drawn::Swapchain(texture),
            view,
        })
    }

    pub fn end_frame(&mut self, frame: Frame) {
        match (frame.drawn, &mut self.target) {
            (Drawn::Swapchain(texture), _) => self.queue.present(texture),
            // No swapchain to throttle the CPU offscreen: wait for the previous frame, or textures and staging pile up until out of memory.
            (Drawn::Offscreen(_), Target::Offscreen { in_flight, .. }) => {
                if let Some(previous) = in_flight.take() {
                    let _ = self.device.poll(wgpu::PollType::Wait {
                        submission_index: Some(previous),
                        timeout: None,
                    });
                }
                *in_flight = self.submitted.take();
            }
            (Drawn::Offscreen(_), Target::Window { .. }) => {
                unreachable!("an offscreen frame from a window")
            }
        }
    }

    pub fn create_encoder(&self, label: &str) -> wgpu::CommandEncoder {
        self.device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some(label) })
    }

    pub fn submit(&mut self, encoder: wgpu::CommandEncoder) {
        self.submitted = Some(self.queue.submit(std::iter::once(encoder.finish())));
    }
}