Trac3r-rust/src/canvas.rs

use crate::vertex_2d::{ColoredVertex2D, Vertex2D};
use vulkano::command_buffer::{AutoCommandBufferBuilder, DynamicState};
use std::collections::HashMap;
use vulkano::buffer::{BufferAccess, BufferUsage, ImmutableBuffer, CpuAccessibleBuffer};
use std::sync::Arc;
use vulkano::format::{ClearValue, Format};
use vulkano::framebuffer::{FramebufferAbstract, Framebuffer};
use vulkano::device::{Device, Queue};
use vulkano::instance::PhysicalDevice;
use vulkano::image::immutable::ImmutableImage;
use crate::util::shader_kernels::ShaderKernels;
use vulkano::image::{Dimensions, ImageAccess, ImageDimensions, SwapchainImage, ImageUsage, AttachmentImage};
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 crate::util::compute_image::ComputeImage;
use std::iter::FromIterator;
use vulkano::swapchain::Capabilities;
use winit::Window;
use vulkano::pipeline::viewport::Viewport;

// Canvas is the accumulator of Sprites for drawing

// Needs to know:
//   textured?
//   colored?
//   vertices

/*

If it is textured. It needs to be rendered with the texture shader which requires a separate
    graphics pipeline. Might as well have a new render pass as well.


So framebuffer is tied to the swapchains images as well as the renderpass

it appears that renderpass is tied to the individual shader


*/

// I want to be able to draw 2d sprites.

// These sprites might be textured or a single color

// All of the single colors will be grouped into one batch using colored vertices.
// The rest will be grouped by their texture and run individually


pub trait Vertex {
    fn position(&self) -> (f32, f32) {
        (0.0, 0.0)
    }

    fn color(&self) -> Option<(f32, f32, f32, f32)> {
        Some((0., 0., 0., 0.))
    }
}


impl Vertex for ColoredVertex2D {
    fn position(&self) -> (f32, f32) {
        (0.0, 0.0)
    }

    fn color(&self) -> Option<(f32, f32, f32, f32)> {
        Some((0., 0., 0., 0.))
    }
}

pub trait Drawable {
    fn get_vertices(&self) -> Vec<(f32, f32)>;
    fn get_color(&self) -> (f32, f32, f32, f32);
    fn get_texture_handle(&self) -> Option<Arc<u32>>;
    fn get_image_handle(&self) -> Option<Arc<u32>>;
}

// Need three types of shaders. Solid, Textured, Image
#[derive(PartialEq, Eq, Hash, Clone)]
pub enum ShaderType {
    SOLID = 0,
    TEXTURED = 1,
    IMAGE = 2,
}

pub struct CanvasFrame {
    colored_drawables: Vec<ColoredVertex2D>,
    textured_drawables: HashMap<Arc<u32>, Vec<Vec<Vertex2D>>>,
    image_drawables: HashMap<Arc<u32>, Vec<Vec<Vertex2D>>>,
}

impl CanvasFrame {

    pub fn new() -> CanvasFrame {
        CanvasFrame {
            colored_drawables: vec![],
            textured_drawables: Default::default(),
            image_drawables: Default::default()
        }
    }

    // Accumulates the drawables vertices and colors
    pub fn draw(&mut self, drawable: &dyn Drawable) {

        match drawable.get_texture_handle() {
            Some(handle) => {
                self.textured_drawables
                    .entry(handle.clone())
                    .or_insert(Vec::new())
                    .push(drawable.get_vertices().iter().map(|n|
                        Vertex2D {
                            position: [n.0, n.1],
                        }
                    ).collect::<Vec<Vertex2D>>());
            }
            None => {

                match drawable.get_image_handle() {

                    Some(handle) => {
                        self.image_drawables
                            .entry(handle.clone())
                            .or_insert(Vec::new())
                            .push(drawable.get_vertices().iter().map(|n|
                                Vertex2D {
                                    position: [n.0, n.1],
                                }
                            ).collect());
                    }
                    None => {
                        let colors = drawable.get_color();

                        self.colored_drawables.extend(
                            drawable.get_vertices().iter().map(|n|
                                ColoredVertex2D {
                                    position: [n.0, n.1],
                                    color: [colors.0, colors.1, colors.2, colors.3],
                                }
                            )
                        );
                    }
                }
            }
        }
    }
}


#[derive(Clone)]
pub struct Canvas {
    shader_kernels: HashMap<ShaderType, ShaderKernels>,

    texture_store: HashMap<String, Arc<ImmutableImage<Format>>>,

    dynamic_state: DynamicState,
    sampler: Arc<Sampler>,

    // hold the image, texture, and shader buffers the same was as we do CompuState
    image_buffers: Vec<std::sync::Arc<vulkano::image::attachment::AttachmentImage>>,
    image_buffer_handles: Vec<Arc<u32>>,

    texture_buffers: Vec<Arc<ImmutableImage<Format>>>,
    texture_buffer_handles: Vec<Arc<u32>>,

    shader_buffers: HashMap<ShaderType, ShaderKernels>,
    shader_buffer_handles: Vec<Arc<u32>>,

    // Hold onto the vertices we get from the Compu and Canvas Frames
    // When the run comes around, push the vertices to the GPU
    colored_drawables: Vec<ColoredVertex2D>,
    colored_vertex_buffer: Vec<Arc<(dyn BufferAccess + std::marker::Send + std::marker::Sync)>>,

    textured_drawables: HashMap<Arc<u32>, Vec<Vec<Vertex2D>>>,
    textured_vertex_buffer: HashMap<Arc<u32>, Arc<(dyn BufferAccess + std::marker::Send + std::marker::Sync)>>,

    image_drawables: HashMap<Arc<u32>, Vec<Vec<Vertex2D>>>,
    image_vertex_buffer: HashMap<Arc<u32>, Arc<(dyn BufferAccess + std::marker::Send + std::marker::Sync)>>,

    // Looks like we gotta hold onto the queue for managing textures
    queue: Arc<Queue>,
    device: Arc<Device>,
}


impl Canvas {

    // 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,
            }]);

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

    // needs to take in the texture list
    pub fn new(queue: Arc<Queue>,
               device: Arc<Device>,
               physical: PhysicalDevice,
               capabilities: Capabilities) -> Canvas {

        let solid_color_kernel = String::from("color-passthrough");
        let texture_kernel = String::from("simple_texture");

        let shader_kernels : HashMap<ShaderType, ShaderKernels> = HashMap::from_iter(vec![
            (ShaderType::SOLID, ShaderKernels::new(solid_color_kernel, capabilities.clone(), queue.clone(), physical.clone(), device.clone())),
            (ShaderType::TEXTURED, ShaderKernels::new(texture_kernel, capabilities.clone(), queue.clone(), physical.clone(), device.clone()))
        ]);


        Canvas {
            shader_kernels: Default::default(),
            texture_store: Default::default(),

            dynamic_state: DynamicState { line_width: None, viewports: None, scissors: None },
            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![],
            image_buffer_handles: vec![],
            texture_buffers: vec![],
            texture_buffer_handles: vec![],
            shader_buffers: Default::default(),
            shader_buffer_handles: vec![],

            colored_drawables: vec![],
            colored_vertex_buffer: vec![],
            textured_drawables: Default::default(),
            textured_vertex_buffer: Default::default(),
            image_drawables: Default::default(),
            image_vertex_buffer: Default::default(),

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

    pub fn create_image(&mut self, dimensions: (u32, u32), usage: ImageUsage) -> Arc<u32> {
        self.image_buffers.push(
            AttachmentImage::with_usage(
                self.device.clone(),
                [dimensions.0, dimensions.1],
                Format::R8G8B8A8Uint,
                usage).unwrap());
        let id = Arc::new(self.image_buffers.len() as u32);
        self.image_buffer_handles.push(id.clone());
        id
    }

    pub fn get_image(&self, image_handle: Arc<u32>) -> std::sync::Arc<vulkano::image::attachment::AttachmentImage> {
        self.image_buffers.get((*image_handle).clone() as usize).unwrap().clone()
    }

    // TODO Handle file not found gracefully
    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));

        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 apha channel...");
            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
    }

    pub fn load_texture_from_filename(&mut self, filename: String) -> Arc<ImmutableImage<Format>> {
        if self.texture_store.contains_key(&filename.clone()) {
            println!("{} Already exists, not going to replace it.", filename.clone());
            self.texture_store.get(&filename.clone()).unwrap().clone()
        } else {
            let texture = self.get_texture_from_file(filename.clone());
            self.texture_store.insert(filename, texture.clone());
            texture
        }
    }

    pub fn load_texture(&mut self, filename: String) -> Option<Arc<u32>> {
        let texture_buffer = self.get_texture_from_file(filename.clone());
        self.texture_buffers.push(texture_buffer.clone());
        let id = Arc::new(self.texture_buffers.len() as u32);
        self.texture_buffer_handles.push(id.clone());
        Some(id)
    }


    // After done using this, need to call allocated vertex buffers
    pub fn draw(&mut self, canvas_frame: CanvasFrame) {
        self.textured_drawables = canvas_frame.textured_drawables;
        self.colored_drawables = canvas_frame.colored_drawables;
        self.image_drawables = canvas_frame.image_drawables;

        self.allocate_vertex_buffers(self.device.clone());
    }

    fn get_texture(&self, texture_id: String) -> Arc<ImmutableImage<Format>> {
        if let Some(i) = self.texture_store.get(&texture_id) {
            return i.clone();
        } else {
            panic!("{} : Texture not loaded", texture_id);
        }
    }

    pub fn allocate_vertex_buffers(&mut self, device: Arc<Device>) {
        self.colored_vertex_buffer.clear();
        self.textured_vertex_buffer.clear();
        self.image_vertex_buffer.clear();


        //TODO should probably use cpu accessible buffer instead of recreating immutes each frame
//        CpuAccessibleBuffer::from_iter(
//            device.clone(),
//            BufferUsage::vertex_buffer(),
//            self.colored_drawables.iter().cloned(),
//        ).unwrap().0;

        self.colored_vertex_buffer.push(
            ImmutableBuffer::from_iter(
                self.colored_drawables.iter().cloned(),
                BufferUsage::vertex_buffer(),
                self.queue.clone(),
            ).unwrap().0
        );

        for (k, v) in self.textured_drawables.drain() {
            self.textured_vertex_buffer.insert(
                k.clone(),
                ImmutableBuffer::from_iter(
                    v.iter().cloned(),
                    BufferUsage::vertex_buffer(),
                    self.queue.clone(),
                ).unwrap().0,
            );
        }
    }

    fn get_solid_color_descriptor_set(&self) -> Box<dyn DescriptorSet + Send + Sync> {
        println!("{}", self.shader_kernels.get(&ShaderType::SOLID).unwrap().clone().get_pipeline().clone().num_sets());

        let o: Box<dyn DescriptorSet + Send + Sync> = Box::new(
            PersistentDescriptorSet::start(
                self.shader_kernels.get(&ShaderType::SOLID).unwrap().clone().get_pipeline().clone(), 0,
            ).build().unwrap());
        o
    }

    fn get_textured_descriptor_set(&self, texture_id: String)
                                   -> Box<dyn DescriptorSet + Send + Sync> {
        let o: Box<dyn DescriptorSet + Send + Sync> = Box::new(
            PersistentDescriptorSet::start(
                self.shader_kernels.get(&ShaderType::TEXTURED).unwrap().clone().get_pipeline().clone(), 0,
            )
                .add_sampled_image(self.get_texture(texture_id), self.sampler.clone()).unwrap()
                .build().unwrap());
        o
    }

    // This is the image which is written to by the write compute buffer
    // I suppose I could just have a general image set maker instead of compue... they are
    // somewhat similar
    fn get_compute_swap_descriptor_set(&mut self,
                                       device: Arc<Device>,
                                       compute_image: &ComputeImage) -> Box<dyn DescriptorSet + Send + Sync> {
        let 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();

        let o: Box<dyn DescriptorSet + Send + Sync> = Box::new(
            PersistentDescriptorSet::start(
                self.shader_kernels.get(&ShaderType::IMAGE).clone().unwrap().clone().get_pipeline(), 0,
            )
                .add_image(compute_image.clone().get_swap_buffer().clone()).unwrap()
                .build().unwrap());
        o
    }

    pub fn draw_commands(&self,
                         mut command_buffer: AutoCommandBufferBuilder,
                         framebuffers: Vec<Arc<dyn FramebufferAbstract + Send + Sync>>,
                         image_num: usize) -> AutoCommandBufferBuilder {

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

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


        for (shader_type, kernel) in self.shader_kernels.clone().iter() {
            match shader_type {
                ShaderType::SOLID => {
                    command_buffer = command_buffer.draw(
                        kernel.clone().get_pipeline().clone(),
                        &self.dynamic_state.clone(),
                        self.colored_vertex_buffer.clone(),
                        (), (),
                    ).unwrap();
                }
                ShaderType::TEXTURED => {
//                    command_buffer = command_buffer.draw(
//                        kernel.clone().get_pipeline().clone(),
//                        &dynamic_state.clone(), self.textured_vertex_buffer.clone(),
//                        vec![self.get_textured_descriptor_set(String::from("funky-bird.jpg"))], ()
//                    ).unwrap();
                }
                ShaderType::IMAGE => {

                }
            }
        }

        command_buffer
            .end_render_pass()
            .unwrap()
    }


}