roast2d_internal 0.4.0

use glam::UVec2;

use crate::commands::{Command, RemoveTexture, ResourceCommand};
use crate::engine::Engine;
use crate::mesh3d::{
    Draw3DPbr, Draw3DSkinned, InstanceData, Mesh3DInstancedShader, Mesh3DPbrShader,
    ResolvedDraw3DSkinned, SkinnedMesh3DShader,
};
use crate::platform::types::TextureResource;
use crate::render::BackendState;
use crate::renderer::resource::TextureCache;
use crate::renderer::shader2d::runner::Renderer2DRunner;
use crate::renderer::shader3d::draw3d::{Draw3D, InstancedDraw3D};
use crate::renderer::shader3d::runner::{ResolvedDraw3D, Shader3DRunner};
use crate::skybox::SkyboxShader;

pub struct RenderDriver {
    device: wgpu::Device,
    queue: wgpu::Queue,
    surface_view_format: wgpu::TextureFormat,
    depth_format: wgpu::TextureFormat,
    post_sampler: wgpu::Sampler,
    textures: TextureCache,
    renderer2d: Renderer2DRunner,
    renderer_ui: Renderer2DRunner,
    // Cached offscreen
    render_target: Option<wgpu::Texture>,
    depth_target: Option<wgpu::Texture>,
    render_target_size: UVec2,
    renderer3d: Option<Shader3DRunner>,
    renderer3d_instanced: Option<Mesh3DInstancedShader>,
    renderer3d_pbr: Option<Mesh3DPbrShader>,
    renderer3d_skinned: Option<SkinnedMesh3DShader>,
    skybox_shader: Option<SkyboxShader>,
}

impl RenderDriver {
    pub fn new(backend: &BackendState) -> Self {
        let device = backend.device.clone();
        let queue = backend.queue.clone();
        let surface_view_format = backend.surface_view_format;
        let depth_format = wgpu::TextureFormat::Depth32Float;
        let renderer2d = Renderer2DRunner::setup(backend, false, depth_format);
        let renderer_ui = Renderer2DRunner::setup(backend, true, depth_format);
        let post_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("Post Render Sampler"),
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            address_mode_w: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Nearest,
            min_filter: wgpu::FilterMode::Nearest,
            mipmap_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });
        Self {
            device,
            queue,
            surface_view_format,
            depth_format,
            post_sampler,
            textures: TextureCache::default(),
            renderer2d,
            renderer_ui,
            render_target: None,
            depth_target: None,
            render_target_size: UVec2::new(0, 0),
            renderer3d: None,
            renderer3d_instanced: None,
            renderer3d_pbr: None,
            renderer3d_skinned: None,
            skybox_shader: None,
        }
    }

    pub fn resize(&mut self) {
        // Drop cached offscreen so it will be recreated with the new size
        self.render_target = None;
        self.depth_target = None;
        self.render_target_size = UVec2::new(0, 0);
    }

    pub fn apply_resource_cmds(&mut self, cmds: Vec<ResourceCommand>) {
        for cmd in cmds {
            match cmd {
                ResourceCommand::CreateTexture(c) => {
                    let tex = TextureCache::create_texture(&self.device, &self.queue, c.data);
                    self.textures.insert(c.handle.id(), tex);
                }
                ResourceCommand::CreateRawTexture(c) => {
                    let tex = TextureCache::create_raw_texture(&self.device, c.texture);
                    self.textures.insert(c.handle.id(), tex);
                }
                ResourceCommand::RemoveTexture(RemoveTexture(id)) => {
                    self.textures.remove(&id);
                }
            }
        }
    }

    fn ensure_render_target(
        &mut self,
        g: &Engine,
        frame: &wgpu::SurfaceTexture,
    ) -> (wgpu::Texture, wgpu::TextureView) {
        let desired = g.render.logical_size().as_uvec2();
        let need_recreate = self
            .render_target
            .as_ref()
            .map(|_| self.render_target_size != desired)
            .unwrap_or(true);
        if need_recreate {
            let size = wgpu::Extent3d {
                width: desired.x,
                height: desired.y,
                depth_or_array_layers: 1,
            };
            let dimension = frame.texture.dimension();
            let texture = self.device.create_texture(&wgpu::TextureDescriptor {
                label: Some("Render Texture"),
                size,
                mip_level_count: 1,
                sample_count: 1,
                dimension,
                format: self.surface_view_format,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT
                    | wgpu::TextureUsages::TEXTURE_BINDING,
                view_formats: &[],
            });
            self.render_target = Some(texture);

            let depth_texture = self.device.create_texture(&wgpu::TextureDescriptor {
                label: Some("depth_texture"),
                size,
                mip_level_count: 1,
                sample_count: 1,
                dimension,
                format: self.depth_format,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
                view_formats: &[self.depth_format],
            });
            self.depth_target = Some(depth_texture);
            self.render_target_size = desired;
        }
        (
            self.render_target.as_ref().unwrap().clone(),
            self.depth_target
                .as_ref()
                .unwrap()
                .create_view(&wgpu::TextureViewDescriptor::default()),
        )
    }

    // NOTE: texture resolution is handled in record_frame when splitting draw commands.

    #[allow(clippy::too_many_arguments)]
    pub fn record_frame(
        &mut self,
        g: &Engine,
        draw_commands: Vec<Command>,
        draws3d: Vec<Draw3D>,
        instanced_draws3d: Vec<InstancedDraw3D>,
        draws3d_pbr: Vec<Draw3DPbr>,
        draws3d_skinned: Vec<Draw3DSkinned>,
        encoder: &mut wgpu::CommandEncoder,
        frame: &wgpu::SurfaceTexture,
    ) {
        let (render_target, depth_view) = self.ensure_render_target(g, frame);
        // Split draws into world/ui and resolve textures once
        let mut world_draws: Vec<(crate::commands::Draw2D, Option<TextureResource>)> = Vec::new();
        let mut ui_draws: Vec<(crate::commands::Draw2D, Option<TextureResource>)> = Vec::new();
        for cmd in draw_commands.into_iter() {
            match cmd {
                Command::Draw(d) => {
                    let tex = d
                        .content
                        .texture_handle()
                        .and_then(|h| self.textures.get(&h.id()))
                        .cloned();
                    world_draws.push((d, tex));
                }
                Command::DrawUi(d) => {
                    let tex = d
                        .content
                        .texture_handle()
                        .and_then(|h| self.textures.get(&h.id()))
                        .cloned();
                    ui_draws.push((d, tex));
                }
            }
        }
        // Resolve 3D textures
        let resolved_3d: Vec<ResolvedDraw3D> = draws3d
            .into_iter()
            .map(|d| {
                let texture = d
                    .texture
                    .as_ref()
                    .and_then(|h| self.textures.get(&h.id()))
                    .cloned();
                ResolvedDraw3D { draw: d, texture }
            })
            .collect();
        // Resolve instanced 3D textures
        let resolved_instanced: Vec<(
            &crate::mesh3d::Mesh3D,
            &[InstanceData],
            Option<TextureResource>,
            &crate::light3d::Material3D,
        )> = instanced_draws3d
            .iter()
            .map(|d| {
                let texture = d
                    .texture
                    .as_ref()
                    .and_then(|h| self.textures.get(&h.id()))
                    .cloned();
                (&d.mesh, d.instances.as_slice(), texture, &d.material)
            })
            .collect();

        // Resolve PBR 3D textures
        let resolved_pbr: Vec<(&Draw3DPbr, Option<TextureResource>)> = draws3d_pbr
            .iter()
            .map(|d| {
                let texture = d
                    .texture
                    .as_ref()
                    .and_then(|h| self.textures.get(&h.id()))
                    .cloned();
                (d, texture)
            })
            .collect();

        // Resolve skinned 3D textures
        let resolved_skinned: Vec<ResolvedDraw3DSkinned> = draws3d_skinned
            .into_iter()
            .map(|d| {
                let texture_bind_group = d
                    .texture
                    .as_ref()
                    .and_then(|h| self.textures.get(&h.id()))
                    .map(|tex| {
                        if self.renderer3d_skinned.is_none() {
                            self.renderer3d_skinned = Some(SkinnedMesh3DShader::new(g));
                        }
                        self.renderer3d_skinned
                            .as_mut()
                            .unwrap()
                            .get_or_create_texture_bind_group(g, tex)
                            .clone()
                    });
                ResolvedDraw3DSkinned {
                    mesh: d.mesh,
                    model: d.model,
                    color: d.color,
                    material: d.material,
                    bone_matrices: d.bone_matrices,
                    texture_bind_group,
                }
            })
            .collect();

        let has_skybox = g.skybox.is_some();
        let has_3d = !resolved_3d.is_empty();
        let has_instanced_3d = !resolved_instanced.is_empty();
        let has_pbr_3d = !resolved_pbr.is_empty();
        let has_skinned_3d = !resolved_skinned.is_empty();
        let has_world_2d = !world_draws.is_empty();
        let has_ui = !ui_draws.is_empty();

        // First renderer clears the target
        let mut cleared = false;

        // Render skybox first (at infinity, behind everything)
        if has_skybox {
            if self.skybox_shader.is_none() {
                self.skybox_shader = Some(SkyboxShader::new(g));
            }
            if let Some(skybox) = &g.skybox {
                self.skybox_shader.as_ref().unwrap().record(
                    g,
                    &render_target,
                    &depth_view,
                    skybox,
                    encoder,
                    !cleared,
                );
                cleared = true;
            }
        }

        if has_3d {
            if self.renderer3d.is_none() {
                self.renderer3d = Some(Shader3DRunner::new(g));
            }
            self.renderer3d.as_mut().unwrap().record(
                g,
                &render_target,
                &depth_view,
                &resolved_3d,
                encoder,
                !cleared,
            );
            cleared = true;
        }

        // Render instanced 3D draws
        if has_instanced_3d {
            if self.renderer3d_instanced.is_none() {
                self.renderer3d_instanced = Some(Mesh3DInstancedShader::new(g));
            }
            let size = render_target.size();
            let aspect = if size.height == 0 {
                1.0
            } else {
                size.width as f32 / size.height as f32
            };
            let view_proj = g.camera3d().view_proj(aspect);

            // Convert to the format expected by the shader
            let draws_ref: Vec<(
                &crate::mesh3d::Mesh3D,
                &[InstanceData],
                Option<&TextureResource>,
                &crate::light3d::Material3D,
            )> = resolved_instanced
                .iter()
                .map(|(mesh, instances, tex, mat)| (*mesh, *instances, tex.as_ref(), *mat))
                .collect();

            self.renderer3d_instanced.as_mut().unwrap().record(
                g,
                &render_target,
                &depth_view,
                view_proj,
                &draws_ref,
                encoder,
                !cleared,
            );
            cleared = true;
        }

        // Render PBR 3D draws
        if has_pbr_3d {
            if self.renderer3d_pbr.is_none() {
                self.renderer3d_pbr = Some(Mesh3DPbrShader::new(g));
            }
            let size = render_target.size();
            let aspect = if size.height == 0 {
                1.0
            } else {
                size.width as f32 / size.height as f32
            };
            let view_proj = g.camera3d().view_proj(aspect);

            // Convert to the format expected by the shader
            let draws_ref: Vec<(
                &crate::mesh3d::Mesh3D,
                glam::Mat4,
                Option<&TextureResource>,
                &crate::light3d::PbrMaterial,
            )> = resolved_pbr
                .iter()
                .map(|(draw, tex)| (&draw.mesh, draw.model, tex.as_ref(), &draw.material))
                .collect();

            self.renderer3d_pbr.as_mut().unwrap().record(
                g,
                &render_target,
                &depth_view,
                view_proj,
                &draws_ref,
                encoder,
                !cleared,
            );
            cleared = true;
        }

        // Render skinned 3D draws
        if has_skinned_3d {
            if self.renderer3d_skinned.is_none() {
                self.renderer3d_skinned = Some(SkinnedMesh3DShader::new(g));
            }
            let color_view = render_target.create_view(&wgpu::TextureViewDescriptor::default());
            self.renderer3d_skinned.as_mut().unwrap().record(
                g,
                &color_view,
                &depth_view,
                &resolved_skinned,
                encoder,
                !cleared,
            );
            cleared = true;
        }

        let first_clears = !cleared;

        if has_world_2d {
            self.renderer2d.record(
                g,
                &render_target,
                &depth_view,
                world_draws,
                encoder,
                first_clears,
            );
        }

        if has_ui {
            // If 3D was rendered, clear depth before UI (3D uses reverse-Z, incompatible with 2D)
            if has_3d || has_instanced_3d || has_pbr_3d || has_skinned_3d {
                encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                    label: Some("UI Depth Clear Pass"),
                    color_attachments: &[],
                    depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                        view: &depth_view,
                        depth_ops: Some(wgpu::Operations {
                            load: wgpu::LoadOp::Clear(0.0),
                            store: wgpu::StoreOp::Store,
                        }),
                        stencil_ops: None,
                    }),
                    occlusion_query_set: None,
                    timestamp_writes: None,
                });
            }
            let ui_clears = first_clears && !has_world_2d;
            self.renderer_ui
                .record(g, &render_target, &depth_view, ui_draws, encoder, ui_clears);
        }

        // Post pass to present to the frame texture
        let frame_view = frame.texture.create_view(&wgpu::TextureViewDescriptor {
            label: Some("Swapchain sRGB View"),
            format: Some(self.surface_view_format),
            ..Default::default()
        });
        let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("Post Render Pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: &frame_view,
                resolve_target: None,
                ops: wgpu::Operations {
                    // Clear the swapchain so transparent offscreen regions don't preserve the prior frame
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                    store: wgpu::StoreOp::Store,
                },
                depth_slice: None,
            })],
            depth_stencil_attachment: None,
            occlusion_query_set: None,
            timestamp_writes: None,
        });
        // Update post renderer per-frame state here (unified spot)
        g.render.post_shader.frame_pass(g, &mut render_pass);
        let view = render_target.create_view(&wgpu::TextureViewDescriptor::default());
        let tex = TextureResource {
            texture: render_target,
            view,
            sampler: self.post_sampler.clone(),
        };
        g.render.post_shader.draw(&tex, &mut render_pass);
        g.render.post_shader.frame_end(&mut render_pass);
    }
}