perovskite_client 0.3.2

use std::ops::Deref;
use std::sync::Arc;

use crate::{media::CacheManager, vulkan::RectF32};
use perovskite_core::protocol::entities as proto;
use perovskite_core::protocol::render::TextureTransform;
use rustc_hash::{FxHashMap, FxHashSet};
use texture_packer::Rect;

use super::{VkAllocator, VulkanContext};
use crate::media::load_or_generate_image;
use crate::vulkan::atlas::{TextureAtlas, TextureKey};
use crate::vulkan::shaders::entity_geometry::EntityVertex;
use crate::vulkan::shaders::VkDrawBufferGpu;
use anyhow::{bail, ensure, Result};
use cgmath::{Matrix3, Rad, Vector3, Zero};

/// Manages the entity definitions and their respective meshes
///
/// Does not hold the actual state of the entities being sent to the display adapter;
/// that state is in EntityState.
pub(crate) struct EntityRenderer {
    texture_atlas: TextureAtlas,
    _allocator: Arc<VkAllocator>,
    mesh_definitions: FxHashMap<u32, EntityMesh>,
    /// These include buffers that render a single entity. Later on, as we add various accelerated and
    /// instanced rendering, the renderer may build buffers with multiple entities in them.
    singleton_gpu_buffers: FxHashMap<u32, Option<VkDrawBufferGpu<EntityVertex>>>,
}
impl EntityRenderer {
    pub(crate) async fn new(
        entity_defs: Vec<proto::EntityDef>,
        cache_manager: &mut CacheManager,
        ctx: &VulkanContext,
    ) -> Result<EntityRenderer> {
        let mut all_texture_names = FxHashSet::default();

        let mut pack_textures = FxHashSet::default();
        for def in &entity_defs {
            if let Some(appearance) = &def.appearance {
                for mesh in &appearance.custom_mesh {
                    if let Some(tex) = &mesh.texture {
                        pack_textures.insert(TextureKey::from(tex));
                        all_texture_names.insert(tex.diffuse.clone());
                        if !tex.rt_specular.is_empty() {
                            all_texture_names.insert(tex.rt_specular.clone());
                        }
                        if !tex.emissive.is_empty() {
                            all_texture_names.insert(tex.emissive.clone());
                        }
                        if !tex.normal_map.is_empty() {
                            all_texture_names.insert(tex.normal_map.clone());
                        }
                        if !tex.alt_diffuse.is_empty() {
                            all_texture_names.insert(tex.alt_diffuse.clone());
                        }
                    }
                }
            }
        }

        let mut fetched_textures = FxHashMap::default();
        for x in all_texture_names {
            let texture = load_or_generate_image(cache_manager, &x).await?;
            fetched_textures.insert(x, texture);
        }
        let texture_atlas = TextureAtlas::new(&ctx, pack_textures, fetched_textures)?;
        // Entity classes may not be contiguous or ordered; hence a hashmap for now
        let mut all_meshes = FxHashMap::default();
        let mut singleton_gpu_buffers = FxHashMap::default();

        for def in entity_defs {
            if let Some(appearance) = def.appearance {
                let mesh = EntityRenderer::pre_render(
                    appearance,
                    &texture_atlas.texel_coords,
                    (texture_atlas.width, texture_atlas.height),
                    def.short_name,
                )?;
                let singleton_buffer =
                    VkDrawBufferGpu::from_buffers(&mesh.vtx, &mesh.idx, ctx.clone_allocator())?;
                singleton_gpu_buffers.insert(def.entity_class, singleton_buffer);

                all_meshes.insert(def.entity_class, mesh);
            }
        }
        Ok(EntityRenderer {
            texture_atlas,
            _allocator: ctx.clone_allocator(),
            mesh_definitions: all_meshes,
            singleton_gpu_buffers,
        })
    }

    pub(crate) fn atlas(&self) -> &TextureAtlas {
        &self.texture_atlas
    }

    /// Convert from the network mesh to a vulkan renderable mesh
    fn pre_render(
        appearance: proto::EntityAppearance,
        texture_coords: &FxHashMap<TextureKey, Rect>,
        atlas_dims: (u32, u32),
        name: String,
    ) -> Result<EntityMesh> {
        let meshes = &appearance.custom_mesh;
        let vertex_count = meshes.iter().map(|x| x.x.len()).sum();
        if vertex_count >= u32::MAX as usize {
            bail!("Too many vertices");
        }
        let mut vertices = Vec::with_capacity(vertex_count);
        let mut indices = Vec::with_capacity(meshes.iter().map(|x| x.indices.len()).sum());
        let mut aabb_min = Vector3::new(f32::INFINITY, f32::INFINITY, f32::INFINITY);
        let mut aabb_max = Vector3::new(-f32::INFINITY, -f32::INFINITY, -f32::INFINITY);
        for mesh in meshes {
            let tex_ref = mesh.texture.as_ref();
            let tex_rectangle: RectF32 = tex_ref
                .and_then(|x| texture_coords.get(&TextureKey::from(x)))
                .unwrap_or_else(|| texture_coords.get(&TextureKey::FallbackUnknownTex).unwrap())
                .into();
            let transform = tex_ref
                .map(|t| {
                    TextureTransform::try_from(t.texture_transform)
                        .unwrap_or(TextureTransform::None)
                })
                .unwrap_or(TextureTransform::None);
            if transform != TextureTransform::None {
                if let Some(t) = tex_ref {
                    if !t.normal_map.is_empty() {
                        log::warn!(
                            "Texture \"{}\" has a normal map and a texture transform; the normals \
                             in the normal map will NOT be automatically adjusted to match the new \
                             orientation",
                            t.diffuse
                        );
                    }
                }
            }
            let tex_rectangle = tex_rectangle.with_transform(transform).div(atlas_dims);
            let vertices_len = mesh.x.len();
            ensure!(mesh.y.len() == vertices_len);
            ensure!(mesh.z.len() == vertices_len);
            ensure!(mesh.u.len() == vertices_len);
            ensure!(mesh.v.len() == vertices_len);
            ensure!(mesh.nx.len() == vertices_len);
            ensure!(mesh.ny.len() == vertices_len);
            ensure!(mesh.nz.len() == vertices_len);
            for i in 0..vertices_len {
                // Bilinear interpolation using explicit corners, supporting transforms.
                // atlas_uv = tl + u*(tr - tl) + v*(bl - tl)
                let [tl_u, tl_v] = tex_rectangle.tl();
                let [tr_u, tr_v] = tex_rectangle.tr();
                let [bl_u, bl_v] = tex_rectangle.bl();
                let mu = mesh.u[i];
                let mv = mesh.v[i];
                let u = tl_u + mu * (tr_u - tl_u) + mv * (bl_u - tl_u);
                let v = tl_v + mu * (tr_v - tl_v) + mv * (bl_v - tl_v);
                vertices.push(EntityVertex {
                    position: [mesh.x[i], mesh.y[i], mesh.z[i]],
                    normal: [mesh.nx[i], mesh.ny[i], mesh.nz[i]],
                    uv_texcoord: [u, v],
                });
                aabb_min.x = aabb_min.x.min(mesh.x[i]);
                aabb_min.y = aabb_min.y.min(mesh.y[i]);
                aabb_min.z = aabb_min.z.min(mesh.z[i]);

                aabb_max.x = aabb_max.x.max(mesh.x[i]);
                aabb_max.y = aabb_max.y.max(mesh.y[i]);
                aabb_max.z = aabb_max.z.max(mesh.z[i]);
            }
            ensure!(mesh.indices.len() % 3 == 0);
            ensure!(mesh.indices.iter().all(|x| *x < vertices_len as u32));
            let offset = indices.len();
            for index in &mesh.indices {
                indices.push(index + offset as u32);
            }
        }
        Ok(EntityMesh {
            vtx: vertices,
            idx: indices,
            attach_offset: match &appearance.attachment_offset {
                Some(x) => x.try_into().unwrap(),
                None => Vector3::zero(),
            },
            attach_in_model_space: appearance.attachment_offset_in_model_space,
            aabb: (aabb_min, aabb_max),
            def: appearance.clone(),
            name,
        })
    }

    pub(crate) fn transform_position(
        &self,
        class: u32,
        pos: Vector3<f64>,
        face_dir: Rad<f32>,
        pitch: Rad<f32>,
    ) -> Vector3<f64> {
        let def = match self.mesh_definitions.get(&class) {
            Some(x) => x,
            None => return pos,
        };
        if def.attach_in_model_space {
            let rotation: Matrix3<f64> = (Matrix3::from_angle_y(face_dir)
                * Matrix3::from_angle_x(pitch))
            .cast()
            .unwrap();
            pos + (rotation * def.attach_offset)
        } else {
            def.attach_offset + pos
        }
    }

    pub(crate) fn get_singleton(&self, class: u32) -> Option<VkDrawBufferGpu<EntityVertex>> {
        self.singleton_gpu_buffers.get(&class).cloned().flatten()
    }

    pub(crate) fn client_info(&self, class: u32) -> Option<&proto::EntityAppearance> {
        self.mesh_definitions.get(&class).map(|x| &x.def)
    }

    pub(crate) fn class_name(&self, class: u32) -> Option<&str> {
        self.mesh_definitions.get(&class).map(|x| x.name.deref())
    }

    pub(crate) fn mesh_aabb(&self, class: u32) -> Option<(Vector3<f32>, Vector3<f32>)> {
        self.mesh_definitions.get(&class).map(|x| x.aabb)
    }
}

pub(crate) struct EntityMesh {
    // Note: As the shaders evolve, we will stop using CubeGeometryVertex and instead will use
    // a more capable vertex type that can express entity-specific details.
    pub(crate) vtx: Vec<EntityVertex>,
    pub(crate) idx: Vec<u32>,
    pub(crate) attach_offset: Vector3<f64>,
    pub(crate) attach_in_model_space: bool,
    // Vulkan-coordinate (Y down) bounding box. Subtlety: vertices in the mesh vertex buffer, but
    // with no indices referencing them, contribute to the bounding box.
    pub(crate) aabb: (Vector3<f32>, Vector3<f32>),
    pub(crate) name: String,
    // TODO: Is duplicating the mesh in here a waste of memory, or will it remain reasonably small
    // under normal circumstances?
    pub(crate) def: proto::EntityAppearance,
}