Trac3r-rust/src/canvas/canvas_state.rs

use vulkano::command_buffer::{AutoCommandBufferBuilder, DynamicState};
use std::collections::{HashMap, HashSet};
use vulkano::buffer::{BufferAccess, BufferUsage, ImmutableBuffer, CpuAccessibleBuffer};
use std::sync::Arc;
use vulkano::format::{ClearValue, Format, R8Unorm, ClearValuesTuple};
use vulkano::framebuffer::{FramebufferAbstract, Framebuffer, RenderPass, RenderPassAbstract};
use vulkano::device::{Device, Queue};
use vulkano::instance::PhysicalDevice;
use vulkano::image::immutable::ImmutableImage;
use vulkano::image::{Dimensions, ImageAccess, ImageDimensions, SwapchainImage, ImageUsage, AttachmentImage, ImageLayout};
use vulkano::sampler::{Sampler, SamplerAddressMode, MipmapMode, Filter};
use vulkano::descriptor::DescriptorSet;
use vulkano::descriptor::descriptor_set::PersistentDescriptorSet;
use std::path::PathBuf;
use image::GenericImageView;
use std::iter::FromIterator;
use vulkano::swapchain::Capabilities;
use vulkano::pipeline::viewport::Viewport;
use vulkano::descriptor::descriptor::DescriptorDescTy::TexelBuffer;
use crate::canvas::canvas_frame::CanvasFrame;
use std::hash::Hash;
use vulkano::pipeline::depth_stencil::{StencilFaceFlags, DynamicStencilValue};
use vulkano::memory::pool::PotentialDedicatedAllocation::Generic;
use std::borrow::Borrow;
use std::fs::File;
use std::io::Read;
use rusttype::{Font, PositionedGlyph, Scale, Rect, point, GlyphId, Line, Curve, Segment};
use vulkano::pipeline::vertex::{VertexDefinition, Vertex};
use crate::canvas::managed::shader::dynamic_vertex::RuntimeVertexDef;
use crate::canvas::managed::handles::{CanvasTextureHandle, CanvasImageHandle, CanvasFontHandle, CompiledShaderHandle, Handle, DrawableHandle};
use crate::canvas::managed::gpu_buffers::{CanvasImage, CanvasTexture, CanvasFont};
use crate::canvas::managed::shader::shader_common::CompiledShader;
use crate::canvas::managed::shader::generic_shader::GenericShader;
use crate::VertexType;
use crate::util::vertex::{TextVertex3D, TextureVertex3D, ImageVertex3D, ColorVertex3D, CanvasFrameAllocation};
use shade_runner::Input;
use winit::window::Window;


/// Canvas state is used for storage of texture and image buffers in addition to vertex buffers
/// Canvas state also contains logic for writing the stored buffers to the command_buffer
#[derive(Clone)]
pub struct CanvasState {
    /// Generated during new()
    dynamic_state: DynamicState,
    /// Generated during new()
    sampler: Arc<Sampler>,

    /// hold the image, texture, and Fonts the same was as we do CompuState
    image_buffers: Vec<Arc<CanvasImage>>,
    texture_buffers: Vec<Arc<CanvasTexture>>,
    font_buffers: Vec<Arc<CanvasFont>>,

    /// Compiled Graphics pipelines have a handle which self describe their position in this vector
    shader_buffers: Vec<Arc<Box<dyn CompiledShader>>>,

    /// Looks like we gotta hold onto the queue for managing textures
    queue: Arc<Queue>,
    device: Arc<Device>,
    render_pass: Arc<dyn RenderPassAbstract + Send + Sync>,
}


impl CanvasState {
    /// This method is called once during initialization, then again whenever the window is resized
    pub fn window_size_dependent_setup(&mut self, images: &[Arc<SwapchainImage<Window>>])
                                       -> Vec<Arc<dyn FramebufferAbstract + Send + Sync>> {
        let dimensions = images[0].dimensions();

        self.dynamic_state.viewports =
            Some(vec![Viewport {
                origin: [0.0, 0.0],
                dimensions: [dimensions.width() as f32, dimensions.height() as f32],
                depth_range: 0.0..1.0,
            }]);

        let dimensions = [dimensions.width(), dimensions.height()];
        let depth_buffer = AttachmentImage::transient(self.device.clone(), dimensions, Format::D32Sfloat_S8Uint).unwrap();

        images.iter().map(|image| {
            Arc::new(
                Framebuffer::start(self.render_pass.clone())
                    .add(image.clone()).unwrap()
                    .add(depth_buffer.clone()).unwrap()
                    .build().unwrap()
            ) as Arc<dyn FramebufferAbstract + Send + Sync>
        }).collect::<Vec<_>>()
    }

    /// Creates a Canvas State. Which at this point is pretty empty
    pub fn new(queue: Arc<Queue>,
               device: Arc<Device>,
               physical: PhysicalDevice,
               capabilities: Capabilities) -> CanvasState {
        let format = capabilities.supported_formats[0].0;

        let render_pass = Arc::new(vulkano::single_pass_renderpass!(
            device.clone(),

            // Attachments are outgoing like f_color
            attachments: {
                // `color` is a custom name we give to the first and only attachment.
                color: {
                    // `load: Clear` means that we ask the GPU to clear the content of this
                    // attachment at the start of the drawing.
                    load: Clear,
                    // `store: Store` means that we ask the GPU to store the output of the draw
                    // in the actual image. We could also ask it to discard the result.
                    store: Store,
                    // `format: <ty>` indicates the type of the format of the image. This has to
                    // be one of the types of the `vulkano::format` module (or alternatively one
                    // of your structs that implements the `FormatDesc` trait). Here we use the
                    // same format as the swapchain.
                    format: format,
                    samples: 1,
                },

                depth: {
                    load: Clear,
                    store: DontCare,
                    format: Format::D32Sfloat_S8Uint,
                    samples: 1,
                }
            },
            pass: {
                // We use the attachment named `color` as the one and only color attachment.
                color: [color],
                // No depth-stencil attachment is indicated with empty brackets.
                depth_stencil: {depth}
            }
        ).unwrap());


        CanvasState {

            // TODO: Might need to move this
            dynamic_state: DynamicState {
                line_width: None,
                viewports: None,
                scissors: None,
                compare_mask: Some(DynamicStencilValue {
                    face: StencilFaceFlags::StencilFrontAndBack,
                    value: 0xFF,
                }),
                write_mask: Some(DynamicStencilValue {
                    face: StencilFaceFlags::StencilFrontAndBack,
                    value: 0xFF,
                }),
                reference: Some(DynamicStencilValue {
                    face: StencilFaceFlags::StencilFrontAndBack,
                    value: 0xFF,
                }),
            },
            sampler: Sampler::new(device.clone(),
                                  Filter::Linear, Filter::Linear,
                                  MipmapMode::Nearest,
                                  SamplerAddressMode::Repeat, SamplerAddressMode::Repeat,
                                  SamplerAddressMode::Repeat, 0.0, 1.0, 0.0, 0.0).unwrap(),
            image_buffers: vec![],
            texture_buffers: vec![],
            shader_buffers: vec![],
            font_buffers: vec![],

            queue: queue.clone(),
            device: device.clone(),
            render_pass: render_pass.clone(),
        }
    }

    /// Using the dimensions and suggested usage, load a CanvasImage and return it's handle
    pub fn create_image(&mut self, dimensions: (u32, u32), usage: ImageUsage) -> Arc<CanvasImageHandle> {
        let handle = Arc::new(CanvasImageHandle { handle: self.image_buffers.len() as u32 });

        let image = CanvasImage {
            handle: handle.clone(),
            buffer: AttachmentImage::with_usage(
                self.device.clone(),
                [dimensions.0, dimensions.1],
                Format::R8G8B8A8Uint,
                usage).unwrap(),
            size: dimensions,
        };

        self.image_buffers.push(Arc::new(image));

        handle
    }

    /// Return the image buffer from an input image handle
    pub fn get_image(&self, image_handle: Arc<CanvasImageHandle>) -> Arc<AttachmentImage> {
        self.image_buffers.get((*image_handle).clone().get_handle() as usize).unwrap()
            .clone().buffer.clone()
    }

    /// Load a texture buffer from an input filename
    fn get_texture_from_file(&self, image_filename: String) -> Arc<ImmutableImage<Format>> {
        let project_root =
            std::env::current_dir()
                .expect("failed to get root directory");

        let mut compute_path = project_root.clone();
        compute_path.push(PathBuf::from("resources/images/"));
        compute_path.push(PathBuf::from(image_filename.clone()));

        let img = image::open(compute_path).expect("Couldn't find image");

        let xy = img.dimensions();

        let data_length = xy.0 * xy.1 * 4;
        let pixel_count = img.raw_pixels().len();

        let mut image_buffer = Vec::new();

        if pixel_count != data_length as usize {
            println!("Creating alpha channel for {}", image_filename.clone());
            for i in img.raw_pixels().iter() {
                if (image_buffer.len() + 1) % 4 == 0 {
                    image_buffer.push(255);
                }
                image_buffer.push(*i);
            }
            image_buffer.push(255);
        } else {
            image_buffer = img.raw_pixels();
        }

        let (texture, tex_future) = ImmutableImage::from_iter(
            image_buffer.iter().cloned(),
            Dimensions::Dim2d { width: xy.0, height: xy.1 },
            Format::R8G8B8A8Srgb,
            self.queue.clone(),
        ).unwrap();

        texture
    }

    /// Load a texture using it's filename from a file. Returns the handle of the loaded texture
    pub fn load_texture(&mut self, filename: String) -> Option<Arc<CanvasTextureHandle>> {
        let handle = Arc::new(CanvasTextureHandle {
            handle: self.texture_buffers.len() as u32
        });

        let texture = Arc::new(CanvasTexture {
            handle: handle.clone(),
            buffer: self.get_texture_from_file(filename.clone()),
            name: filename.clone(),
            size: (0, 0),
        });

        self.texture_buffers.push(texture);

        Some(handle)
    }

    /// Load and Compile a shader with the filename at resources/shaders
    /// Takes physical and capabilities as we don't store that in Canvas
    pub fn load_shader<T: 'static, V>(&mut self,
                                      filename: String,
                                      capabilities: Capabilities) -> Option<Arc<CompiledShaderHandle>>
        where T: CompiledShader, V: Vertex {
        let handle = Arc::new(CompiledShaderHandle {
            handle: self.shader_buffers.len() as u32
        });

        let shader: Box<dyn CompiledShader> = Box::new(T::new::<V>(
            filename.clone(),
            self.device.clone(),
            handle.clone(),
            self.render_pass.clone(),
        ));

        self.shader_buffers.push(Arc::new(shader));

        Some(handle)
    }

    ///
    pub fn load_font(&mut self, name: String) -> Arc<CanvasFontHandle> {
        let handle = Arc::new(CanvasFontHandle { handle: self.font_buffers.len() as u32 });

        self.font_buffers.push(Arc::new({
            let font = Font::from_bytes({
                let mut f = File::open("resources/fonts/sansation.ttf")
                    .expect("Font file not found");
                let mut font_data = Vec::new();
                f.read_to_end(&mut font_data).expect("Dont know");
                font_data
            }).unwrap();

            let mut current_x = 0;
            let mut current_y = 0;

            let mut accumulator = Vec::new();


            for i in (0..255) {
                let glyph = font.glyph('d');

                let s = glyph.scaled(Scale { x: 1.0, y: 1.0 });

                let shape = s.shape().unwrap();

                for contour in shape {
                    for segment in contour.segments {
                        match segment {
                            Segment::Line(l) => {
                                accumulator.push(TextVertex3D {
                                    position: [l.p[0].x as f32, l.p[0].y as f32, 0.0],
                                });
                            }
                            Segment::Curve(c) => {
                                accumulator.push(TextVertex3D {
                                    position: [c.p[0].x as f32, c.p[0].y as f32, 0.0],
                                });
                            }
                        }
                    }
                }
            }

            CanvasFont {
                handle: handle.clone(),
                font: font.clone(),
                name: name,
                buffer: ImmutableBuffer::from_iter(
                    accumulator.iter().cloned(),
                    BufferUsage::vertex_buffer(), self.queue.clone()).unwrap().0,
            }
        }));

        handle
    }

    /// Using the texture name, iterates through the stored textures and matches by the name
    pub fn get_texture_handle(&self, texture_name: String)
                              -> Option<Arc<CanvasTextureHandle>> {
        for i in self.texture_buffers.clone() {
            if i.name == texture_name {
                return Some(i.handle.clone());
            }
        }
        None
    }

    /// Using the shader name, iterates through the stored shaders and matches by the name
    pub fn get_shader_handle(&self, shader_name: String)
                             -> Option<Arc<CompiledShaderHandle>> {
        for shader in self.shader_buffers.clone() {
            if shader.get_name() == shader_name {
                return Some(shader.get_handle().clone());
            }
        }
        None
    }

    /// Using the font name, iterates through the stored fonts and matches by the name
    pub fn get_font_handle(&self, font_name: String) -> Option<Arc<CanvasFontHandle>> {
        for font in self.font_buffers.clone() {
            if font.name == font_name {
                return Some(font.handle.clone());
            }
        }
        None
    }

    /// Using the texture handle, grab the stored texture and return the buffer
    pub fn get_texture(&self, texture_handle: Arc<CanvasTextureHandle>)
                       -> Arc<ImmutableImage<Format>> {
        let handle = texture_handle.get_handle() as usize;

        if let Some(i) = self.texture_buffers.get(handle) {
            return i.clone().buffer.clone();
        } else {
            panic!("{} : Texture not loaded", handle);
        }
    }

    /// Builds the descriptor set for solid colors using the input kernel (needs to support solid colors)
    fn get_solid_color_descriptor_set(&self, kernel: Arc<GenericShader>) -> Box<dyn DescriptorSet + Send + Sync> {
        let o: Box<dyn DescriptorSet + Send + Sync> = Box::new(
            PersistentDescriptorSet::start(
                kernel.clone().get_pipeline().clone().descriptor_set_layout(0).unwrap().clone(),
            ).build().unwrap());
        o
    }

    /// Consume and allocate the canvas frame data to the GPU
    pub fn allocate(&mut self, canvas_frame: CanvasFrame) -> CanvasFrameAllocation {

        let mut colored_vertex_buffer: Vec<ColorVertex3D> = Vec::default();
        let mut textured_vertex_buffer: HashMap<Arc<CanvasTextureHandle>, Vec<TextureVertex3D>> = HashMap::new();
        let mut image_vertex_buffer: HashMap<Arc<CanvasImageHandle>, Vec<ImageVertex3D>> = HashMap::new();
        let mut text_instances: HashMap<Arc<CanvasFontHandle>, Vec<TextVertex3D>> = HashMap::new();
        let mut text_vertex_buffer: Vec<ColorVertex3D> = Vec::new();

        // separate the mux of vertex containers back out
        for value in canvas_frame.map {
            match value {
                VertexType::TextureType(vertices, handle) => {
                    textured_vertex_buffer.entry(handle).or_insert(vertices.clone()).extend(vertices);
                }
                VertexType::ImageType(vertices, handle) => {
                    image_vertex_buffer.entry(handle).or_insert(vertices.clone()).extend(vertices);
                }
                VertexType::ColorType(vertices) => {
                    colored_vertex_buffer.extend(vertices);
                }
                VertexType::ThreeDType(vertices) => {

                }
                VertexType::TextType(vertices) => {
                    text_vertex_buffer.extend(vertices);
                }
            };
        };

        let mut allocated_colored_buffer: Vec<Arc<(dyn BufferAccess + Send + Sync)>> = Vec::new();
        if !colored_vertex_buffer.is_empty() {
            allocated_colored_buffer.push(ImmutableBuffer::from_iter(
                colored_vertex_buffer.iter().cloned(),
                BufferUsage::vertex_buffer(),
                self.queue.clone(),
            ).unwrap().0);
        }

        let mut allocated_text_buffer: Vec<Arc<(dyn BufferAccess + Send + Sync)>> = Vec::new();
        if !text_vertex_buffer.is_empty() {
            allocated_text_buffer.push(ImmutableBuffer::from_iter(
                text_vertex_buffer.iter().cloned(),
                BufferUsage::vertex_buffer(),
                self.queue.clone(),
            ).unwrap().0);
        }

        CanvasFrameAllocation {
            colored_vertex_buffer: allocated_colored_buffer,
            textured_vertex_buffer: textured_vertex_buffer.into_iter().map(|(k, v)| {
                (k,
                 ImmutableBuffer::from_iter(
                     v.iter().cloned(),
                     BufferUsage::vertex_buffer(),
                     self.queue.clone(),
                 ).unwrap().0 as Arc<(dyn BufferAccess + Send + Sync)>)
            }).collect(),
            image_vertex_buffer: image_vertex_buffer.into_iter().map(|(k, v)| {
                (k,
                 ImmutableBuffer::from_iter(
                     v.iter().cloned(),
                     BufferUsage::vertex_buffer(),
                     self.queue.clone(),
                 ).unwrap().0 as Arc<(dyn BufferAccess + Send + Sync)>)
            }).collect(),
            text_instances: Default::default(),
            text_vertex_buffer: allocated_text_buffer,
        }
    }

    /// Pushes the draw commands to the command buffer. Requires the framebuffers and
    /// image number to be passed in as they are taken care of by the vkprocessor
    pub fn draw_commands(&mut self,
                         mut command_buffer: &mut AutoCommandBufferBuilder,
                         framebuffers: Vec<Arc<dyn FramebufferAbstract + Send + Sync>>,
                         image_num: usize,
                         allocated_buffers: CanvasFrameAllocation) {

        // Specify the color to clear the framebuffer with i.e. blue
        let clear_values = vec!(
            ClearValue::Float([0.0, 0.0, 0.0, 0.0]),
            ClearValue::DepthStencil((1.0, 0x00)),
        );

        self.dynamic_state = DynamicState {
            line_width: None,
            viewports: self.dynamic_state.viewports.clone(),
            scissors: None,
            compare_mask: None,
            write_mask: None,
            reference: None,
        };

        command_buffer = command_buffer.begin_render_pass(
            framebuffers[image_num].clone(), false, clear_values.clone(),
        ).unwrap();

        // Solid colors
        let mut shader = self.shader_buffers.get(
            self.get_shader_handle(String::from("color-passthrough"))
                .unwrap().clone().get_handle() as usize
        ).unwrap();

        // This looks a little weird as colored_vertex_buffer is a vec of GPU allocated vecs.
        // But we can pass in multiple vertex buffers

        if !allocated_buffers.colored_vertex_buffer.is_empty() {
            command_buffer = command_buffer.draw(
                shader.get_pipeline().clone(),
                &self.dynamic_state.clone(),
                allocated_buffers.colored_vertex_buffer.clone(),
                (), (),
            ).unwrap();
        }

        // Images
        let mut shader = self.shader_buffers.get(
            self.get_shader_handle(String::from("simple_image"))
                .unwrap().clone().get_handle() as usize
        ).unwrap();


        if !allocated_buffers.image_vertex_buffer.is_empty() {
            for (image_handle, vertex_buffer) in allocated_buffers.image_vertex_buffer.clone() {
                let handle = image_handle.clone().get_handle() as usize;
                let descriptor_set = self.image_buffers.get(handle).clone().unwrap().clone()
                    .get_descriptor_set(shader.get_pipeline().clone());

                command_buffer = command_buffer.draw(
                    shader.get_pipeline().clone(),
                    // Multiple vertex buffers must have their definition in the pipeline!
                    &self.dynamic_state.clone(), vec![vertex_buffer],
                    vec![descriptor_set], (),
                ).unwrap();
            }
        }

        // Textures
        let mut shader = self.shader_buffers.get(
            self.get_shader_handle(String::from("simple_texture"))
                .unwrap().clone().get_handle() as usize
        ).unwrap();

        if !allocated_buffers.textured_vertex_buffer.is_empty() {
            for (texture_handle, vertex_buffer) in allocated_buffers.textured_vertex_buffer.clone() {
                let handle = texture_handle.clone().get_handle() as usize;
                let descriptor_set = self.texture_buffers.get(handle).clone().unwrap().clone()
                    .get_descriptor_set(shader.get_pipeline(), self.sampler.clone());

                command_buffer = command_buffer.draw(
                    shader.get_pipeline().clone(),
                    // Multiple vertex buffers must have their definition in the pipeline!
                    &self.dynamic_state.clone(), vec![vertex_buffer],
                    vec![descriptor_set], (),
                ).unwrap();
            }
        }

        // Text
        // let mut shader = self.shader_buffers.get(
        //     self.get_shader_handle(String::from("simple_text"))
        //         .unwrap().clone().get_handle() as usize
        // ).unwrap();
//
//        self.dynamic_state = DynamicState {
//            line_width: None,
//            viewports: self.dynamic_state.viewports.clone(),
//            scissors: None,
//            compare_mask: Some(DynamicStencilValue {
//                face: StencilFaceFlags::StencilFrontAndBack,
//                value: 0x00,
//            }),
//            write_mask: Some(DynamicStencilValue {
//                face: StencilFaceFlags::StencilFrontAndBack,
//                value: 0xFF,
//            }),
//            reference: Some(DynamicStencilValue {
//                face: StencilFaceFlags::StencilFrontAndBack,
//                value: 0x00,
//            }),
//        };
//
//        if !allocated_buffers.text_vertex_buffer.is_empty() {
//            command_buffer = command_buffer.draw(
//                shader.get_pipeline().clone(),
//                &self.dynamic_state.clone(),
//                allocated_buffers.text_vertex_buffer.clone(),
//                (), (),
//            ).unwrap();
//        }
//
//        self.dynamic_state = DynamicState {
//            line_width: None,
//            viewports: self.dynamic_state.viewports.clone(),
//            scissors: None,
//            compare_mask: Some(DynamicStencilValue {
//                face: StencilFaceFlags::StencilFrontAndBack,
//                value: 0xFF,
//            }),
//            write_mask: Some(DynamicStencilValue {
//                face: StencilFaceFlags::StencilFrontAndBack,
//                value: 0x00,
//            }),
//            reference: Some(DynamicStencilValue {
//                face: StencilFaceFlags::StencilFrontAndBack,
//                value: 0x00,
//            }),
//        };

        // if !allocated_buffers.text_vertex_buffer.is_empty() {
        //     command_buffer = command_buffer.draw(
        //         shader.get_pipeline().clone(),
        //         &self.dynamic_state.clone(),
        //         allocated_buffers.text_vertex_buffer.clone(),
        //         (), (),
        //     ).unwrap();
        // }

        command_buffer
            .end_render_pass()
            .unwrap();
    }
}