mod oct_tree;
mod empty_volume;
mod light;
mod memorizable;
pub mod generators;

use anyhow::Ok;
use light::{DirectionalLight, LightSource, PointLight};
use vulkanalia::prelude::v1_0::*;
use anyhow::Result;

use cgmath::{vec2, vec3, Vector3};

use std::cell::RefCell;
use std::rc::Rc;

use crate::scene::memorizable::Memorizable;
use crate::app_data::AppData;
use crate::buffer;
use crate::primitives::rec_cuboid::Cuboid;
use crate::vertex;
use crate::primitives::cube::Cube;
use crate::primitives::drawable::Drawable;
use crate::scene::oct_tree::{OctTree, OctTreeIter, CHUNK_SIZE};
use crate::scene::empty_volume::EmptyVolume;

extern crate rand;
use rand::Rng;

#[repr(C)]
#[derive(Clone, Debug, Default)]
pub struct Scene {
    pub vertices: Vec<vertex::Vertex>,
    pub sized_vertices: Vec<vertex::SizedVertex>,
    pub rt_vertices: Vec<vertex::RTVertex>,
    pub indices_cube: Vec<u32>,
    pub indices_cuboid: Vec<u32>,
    pub indices_rt: Vec<u32>,

    pub vertex_buffer_cube: vk::Buffer,
    pub vertex_buffer_memory_cube: vk::DeviceMemory,

    pub index_buffer_cube: vk::Buffer,
    pub index_buffer_memory_cube: vk::DeviceMemory,

    pub vertex_buffer_cuboid: vk::Buffer,
    pub vertex_buffer_memory_cuboid: vk::DeviceMemory,

    pub index_buffer_cuboid: vk::Buffer,
    pub index_buffer_memory_cuboid: vk::DeviceMemory,

    pub vertex_buffer_quad: vk::Buffer,
    pub vertex_buffer_memory_quad: vk::DeviceMemory,

    pub index_buffer_quad: vk::Buffer,
    pub index_buffer_memory_quad: vk::DeviceMemory,

    pub rt_memory: Vec<u32>,

    pub oct_trees: Vec<Vec<Vec<Rc<RefCell<OctTree<Cube>>>>>>,

    pub point_lights: Vec<Rc<RefCell<PointLight>>>,
    pub directional_lights: Vec<Rc<RefCell<DirectionalLight>>>,

    pub memorizables: Vec<Rc<RefCell<dyn Memorizable>>>,
}

impl Scene {
    pub unsafe fn prepare_data(&mut self, instance: &vulkanalia::Instance, device: &vulkanalia::Device, data: &mut AppData) -> Result<()> {
        // todo store the chunks somewhere (or only use them as intermediary for neighbourhood calculation idc)

        let mut empty_volumes: Vec<Rc<RefCell<EmptyVolume>>> = vec![];

        let mut neighbor_trees: Vec<Vec<Vec<Rc<OctTree<Rc<RefCell<EmptyVolume>>>>>>> = vec![];

        
        
        let mut z_index = 0;
        for oct_tree_plane_xy in &self.oct_trees {
            neighbor_trees.push(vec![]);
            let mut y_index = 0;
            for oct_tree_line_y in oct_tree_plane_xy {
                neighbor_trees[z_index].push(vec![]);
                let mut x_index = 0;
                for oct_tree in oct_tree_line_y {
                    let mut new_volumes: Vec<Rc<RefCell<EmptyVolume>>>;
                    let new_neighbors;
                    (new_volumes, new_neighbors) = EmptyVolume::from_oct_tree(oct_tree, Vector3 { x: (x_index * CHUNK_SIZE) as f32 * oct_tree.borrow().scale, y: (y_index * CHUNK_SIZE) as f32 * oct_tree.borrow().scale, z: (z_index * CHUNK_SIZE) as f32 * oct_tree.borrow().scale });
                    empty_volumes.append(&mut new_volumes);

                    neighbor_trees[z_index][y_index].push(Rc::new(new_neighbors));

                    x_index += 1;
                }
                y_index += 1;
            }
            z_index += 1;
        }
        
        let mut z_index = 0;
        for oct_tree_plane_xy in &self.oct_trees {
            let mut y_index = 0;
            for oct_tree_line_x in oct_tree_plane_xy {
                let mut x_index = 0;
                for oct_tree in oct_tree_line_x {
                    
                    if oct_tree_line_x.len() > x_index + 1 {
                        EmptyVolume::combine_results(oct_tree, &neighbor_trees[z_index][y_index][x_index], &oct_tree_line_x[x_index + 1], &neighbor_trees[z_index][y_index][x_index + 1], vertex::Facing::Right);
                        EmptyVolume::combine_results(&oct_tree_line_x[x_index + 1], &neighbor_trees[z_index][y_index][x_index + 1], oct_tree, &neighbor_trees[z_index][y_index][x_index], vertex::Facing::Left);
                    }
                    if oct_tree_plane_xy.len() > y_index + 1 {
                        EmptyVolume::combine_results(oct_tree, &neighbor_trees[z_index][y_index][x_index], &oct_tree_plane_xy[y_index + 1][x_index], &neighbor_trees[z_index][y_index + 1][x_index], vertex::Facing::Back);
                        EmptyVolume::combine_results(&oct_tree_plane_xy[y_index + 1][x_index], &neighbor_trees[z_index][y_index + 1][x_index], oct_tree, &neighbor_trees[z_index][y_index][x_index], vertex::Facing::Front);
                    }
                    if self.oct_trees.len() > z_index + 1 {
                        EmptyVolume::combine_results(oct_tree, &neighbor_trees[z_index][y_index][x_index], &self.oct_trees[z_index + 1][y_index][x_index], &neighbor_trees[z_index + 1][y_index][x_index], vertex::Facing::Top);
                        EmptyVolume::combine_results(&self.oct_trees[z_index + 1][y_index][x_index], &neighbor_trees[z_index + 1][y_index][x_index], oct_tree, &neighbor_trees[z_index][y_index][x_index], vertex::Facing::Bottom);
                    }
                    x_index += 1;
                }
                y_index += 1;
            }
            z_index += 1;
        }
        println!("number of empty volumes is {}", empty_volumes.len());

        for light in &self.point_lights {
            self.memorizables.push(light.clone());
        }

        for light in &self.directional_lights {
            self.memorizables.push(light.clone());
        }

        for volume in &empty_volumes {
            self.memorizables.push(volume.clone());
        }

        self.update_memory(data, false);

        for volume in &empty_volumes {
            let quads = volume.borrow().to_quads();
            for quad in quads {
                quad.draw(&data.topology, self.rt_vertices.len(), self);
            }
        }
        
        if self.vertices.len() != 0 {
            (self.vertex_buffer_cube, self.vertex_buffer_memory_cube) = buffer::create_vertex_buffer(instance, device, &data, &self.vertices)?;
            (self.index_buffer_cube, self.index_buffer_memory_cube) = buffer::create_index_buffer(&instance, &device, &data, &self.indices_cube)?;    
        }
        
        if self.sized_vertices.len() != 0 {
            (self.vertex_buffer_cuboid, self.vertex_buffer_memory_cuboid) = buffer::create_vertex_buffer(instance, device, &data, &self.sized_vertices)?;
            (self.index_buffer_cuboid, self.index_buffer_memory_cuboid) = buffer::create_index_buffer(&instance, &device, &data, &self.indices_cuboid)?;
        }

        if self.rt_vertices.len() != 0 {
            println!("number of quad vertices is {}", self.rt_vertices.len());
            (self.vertex_buffer_quad, self.vertex_buffer_memory_quad) = buffer::create_vertex_buffer(instance, device, &data, &self.rt_vertices)?;
            (self.index_buffer_quad, self.index_buffer_memory_quad) = buffer::create_index_buffer(&instance, &device, &data, &self.indices_rt)?;
        }

        Ok(())

    }

    pub fn is_dirty(&self) -> bool {
        for memorizable in &self.memorizables {
            if memorizable.borrow().is_dirty() {
                return true
            }
        }
        return false
    }

    pub fn update_memory(&mut self, data: &mut AppData, reuse_memory: bool) {
        // reuse_memory controls whether a fresh data vector is created or the existing one is used if it is the right size
        let mut memory_index = 6; 
        // 0 - location for the maximum number of lights referenced per chunk (also will be the invalid memory allocation for pointing to a nonexistant neighbor)
        // 1 - location for the max iterations per light
        // 2 - diffuse raster samples (2*n + 1) * (2*n + 1) so as to always have at least the central fragment covered
        // 3 - diffuse raster size
        // 4 - max recursive rays
        // 5 - diffuse rays per hit
        for memorizable in &self.memorizables {
            memorizable.borrow_mut().set_memory_start(memory_index);
            memory_index += memorizable.borrow_mut().get_buffer_mem_size(data) as usize;
        }

        //println!("Memory size is {} kB, max indes is {}", memory_index * 32 / 8 /1024 + 1, memory_index);
        let mut volume_vec;
        let needs_overwrite;
        if !reuse_memory || memory_index != self.rt_memory.len() {
            volume_vec = vec![data.num_lights_per_volume; memory_index];
            needs_overwrite = true;
        } else {
            needs_overwrite = false;
            volume_vec = self.rt_memory.clone();
        }
        volume_vec[1] = data.max_iterations_per_light;
        volume_vec[2] = data.diffuse_raster_steps;
        volume_vec[3] = u32::from_ne_bytes(data.diffuse_raster_size.to_ne_bytes());
        volume_vec[4] = data.max_recursive_rays;
        volume_vec[5] = data.diffuse_rays_per_hit;
        
        for memorizable in &self.memorizables {
            if needs_overwrite || memorizable.borrow().is_dirty() {
                volume_vec = memorizable.borrow_mut().insert_into_memory(volume_vec, data, &self);
            }
        }

        self.rt_memory = volume_vec;
        data.scene_rt_memory_size = (self.rt_memory.len() * 4) as u64; // size of the needed buffer size in bytes
    }

    pub unsafe fn destroy(&mut self, device: &vulkanalia::Device) {
        device.destroy_buffer(self.index_buffer_cube, None);
        device.free_memory(self.index_buffer_memory_cube, None);

        device.destroy_buffer(self.vertex_buffer_cube, None);
        device.free_memory(self.vertex_buffer_memory_cube, None);

        device.destroy_buffer(self.index_buffer_cuboid, None);
        device.free_memory(self.index_buffer_memory_cuboid, None);

        device.destroy_buffer(self.vertex_buffer_cuboid, None);
        device.free_memory(self.vertex_buffer_memory_cuboid, None);

        device.destroy_buffer(self.index_buffer_quad, None);
        device.free_memory(self.index_buffer_memory_quad, None);

        device.destroy_buffer(self.vertex_buffer_quad, None);
        device.free_memory(self.vertex_buffer_memory_quad, None);   
    }

    fn get_light_iter(&self) -> LightsIter {
        LightsIter::new(self)
    }
}


pub struct LightsIter<'a> {
    light_index: usize,
    scene: &'a Scene,
}

impl<'a> LightsIter<'a> {
    fn new(scene: &'a Scene) -> Self {
        LightsIter {light_index: 0, scene: scene}
    }
}

impl<'a> Iterator for LightsIter<'a> {
    type Item = Rc<RefCell<dyn LightSource>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.light_index >= self.scene.point_lights.len() {
            if self.light_index - self.scene.point_lights.len() >= self.scene.directional_lights.len() {
                None
            } else {
                let result = self.scene.directional_lights[self.light_index - self.scene.point_lights.len()].clone();
                self.light_index += 1;
                Some(result)
            }
        } else {
            let result = self.scene.point_lights[self.light_index].clone();
            self.light_index += 1;
            Some(result)
        }
    }    
}