use std::sync::{Arc};
use std::time::SystemTime;

use cgmath::{Matrix4, SquareMatrix};
use image::{ImageBuffer, ImageFormat, Rgb, Rgba};
use image::buffer::ConvertBuffer;
use vulkano::buffer::{BufferUsage, CpuAccessibleBuffer};
use vulkano::command_buffer::{AutoCommandBuffer, AutoCommandBufferBuilder, DynamicState};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSet};
use vulkano::descriptor::DescriptorSet;
use vulkano::device::{Device, DeviceExtensions, Queue};
use vulkano::format::{ClearValue, Format};
use vulkano::framebuffer::{Framebuffer, FramebufferAbstract, RenderPassAbstract, Subpass};
use vulkano::image::{AttachmentImage, Dimensions, ImageUsage, ImmutableImage, SwapchainImage};
use vulkano::instance::{ApplicationInfo, Instance, InstanceExtensions, PhysicalDevice, Version};
use vulkano::instance::debug::{DebugCallback, MessageSeverity, MessageType};
use vulkano::memory::pool::{PotentialDedicatedAllocation, StdMemoryPoolAlloc};
use vulkano::pipeline::{GraphicsPipeline, GraphicsPipelineAbstract};
use vulkano::pipeline::viewport::Viewport;
use vulkano::sampler::{Filter, MipmapMode, Sampler, SamplerAddressMode};
use vulkano::swapchain::{AcquireError, ColorSpace, FullscreenExclusive, PresentMode, Surface, SurfaceTransform, Swapchain, SwapchainCreationError};
use vulkano::swapchain;
use vulkano::sync::{FlushError, GpuFuture};
use vulkano::sync;
use vulkano_win::VkSurfaceBuild;
use winit::event::{Event, WindowEvent};
use winit::event_loop::{EventLoop, ControlFlow};
use winit::window::{Window, WindowBuilder};

use line_vs::ty::LinePushConstants;
use vs::ty::PushConstants;

use crate::mesh::CPUMesh;

const VALIDATION_LAYERS: &[&str] = &[
    "VK_LAYER_LUNARG_standard_validation"
];

#[derive(Default, Debug, Clone)]
pub struct Vertex {
    pub position: [f32; 3],
    pub uv: [f32; 2],
    pub normal: [f32; 3],
}
vulkano::impl_vertex!(Vertex, position, uv, normal);

#[derive(Default, Debug, Clone)]
pub struct LinePoint {
    pub position: [f32; 3],
}
vulkano::impl_vertex!(LinePoint, position);

pub trait Game {
    /// Returns true if event should be ignored by the vulkan handler
    fn on_window_event(self: &mut Self, event: &Event<()>);

    fn update(self: &mut Self, renderer: &mut VulkanRenderer) -> vs::ty::UniformBufferObject;
}

pub struct Mesh {
    vertex_buffer: Arc<CpuAccessibleBuffer<[Vertex]>>,
    index_buffer: Arc<CpuAccessibleBuffer<[u32]>>,
}

pub struct GameObject {
    pub mesh_index: usize,
    pub texture_index: usize,
    pub model_matrix: Matrix4<f32>,
}

pub(crate) type GameObjectHandle = usize;
pub(crate) type MeshHandle = usize;

pub struct GameData {
    pub start_time: SystemTime,
    pub line_vertices: Vec<LinePoint>,
    pub push_constants: PushConstants,
    pub line_push_constants: LinePushConstants,
    pub recreate_pipeline: bool,
    pub dimensions: [u32; 2],
    pub shutdown: bool,
    pub game_objects: Vec<GameObject>,
    pub meshes: Vec<Mesh>,
    pub textures: Vec<Arc<ImmutableImage<Format>>>,
}

type Yeet = dyn DescriptorSet + Send + Sync;

pub struct VulkanRenderer {
    pub game_data: GameData,
    pub device: Arc<Device>,
    pub framebuffers: Vec<Arc<dyn FramebufferAbstract + Send + Sync>>,
    pub sampler: Arc<Sampler>,
    pub dynamic_state: DynamicState,
    pub pipeline: Arc<dyn GraphicsPipelineAbstract + Send + Sync>,
    pub line_pipeline: Arc<dyn GraphicsPipelineAbstract + Send + Sync>,
    pub line_vertex_buffer: Arc<CpuAccessibleBuffer<[LinePoint], PotentialDedicatedAllocation<StdMemoryPoolAlloc>>>,
    pub surface: Arc<Surface<Window>>,
    pub swapchain: Arc<Swapchain<Window>>,
    pub render_pass: Arc<dyn RenderPassAbstract + Send + Sync>,
    pub queue: Arc<Queue>,
    pub recreate_swapchain: bool,
    pub debug_callback: Option<DebugCallback>,
    pub previous_frame_end: Option<Box<dyn GpuFuture>>,
    pub uniform_buffers: Vec<Arc<CpuAccessibleBuffer<vs::ty::UniformBufferObject>>>,
    pub descriptor_sets: Vec<Arc<Yeet>>,
}

impl VulkanRenderer {
    pub fn init(line_vertices: Vec<LinePoint>, enable_validation_layers: bool) -> (VulkanRenderer, EventLoop<()>) {
        let mut data = GameData {
            push_constants: PushConstants {
                model: Matrix4::identity().into(),
                view: Matrix4::identity().into(),
                projection: Matrix4::identity().into(),
            },
            line_push_constants: LinePushConstants {
                model: Matrix4::identity().into(),
                view: Matrix4::identity().into(),
                projection: Matrix4::identity().into(),
            },
            start_time: SystemTime::now(),
            recreate_pipeline: false,
            shutdown: false,
            line_vertices,
            dimensions: [0, 0],
            meshes: vec![],
            game_objects: vec![],
            textures: vec![],
        };

        if enable_validation_layers {
            println!("Enabling validation layers...");
        }

        let instance = {
            let extensions = InstanceExtensions {
                ext_debug_utils: true,
                ..vulkano_win::required_extensions()
            };

            let app_info = ApplicationInfo {
                application_name: Some("Asuro's Editor".into()),
                application_version: Some(Version { major: 0, minor: 1, patch: 0 }),
                engine_name: Some("Asuro's Rust Engine".into()),
                engine_version: Some(Version { major: 0, minor: 1, patch: 0 })
            };

            if enable_validation_layers {
                let available_layers = vulkano::instance::layers_list().unwrap().map(|layer| String::from(layer.name())).collect::<Vec<String>>();

                VALIDATION_LAYERS.iter().for_each(|wanted_layer_name| {
                    if !available_layers.iter().any(|available_layer_name| available_layer_name == wanted_layer_name) {
                        panic!("Validation layer not found: {:?}. Available layers: {:?}", wanted_layer_name, &available_layers.join(", "));
                    }
                });

                Instance::new(Some(&app_info), &extensions, VALIDATION_LAYERS.iter().cloned()).expect("failed to create Vulkan instance")
            } else {
                Instance::new(Some(&app_info), &extensions, None).expect("failed to create Vulkan instance")
            }
        };

        // lifetime of this is important, even tho it isn't used!
        let mut debug_callback = None;
        if enable_validation_layers {
            let msg_severity = MessageSeverity {
                verbose: false,
                information: true,
                warning: true,
                error: true
            };

            let msg_types = MessageType {
                general: true,
                performance: true,
                validation: true
            };

            debug_callback = DebugCallback::new(&instance, msg_severity, msg_types, |msg| {
                let type_str = match (msg.severity.error, msg.severity.warning, msg.severity.information, msg.severity.verbose) {
                    (true, _, _, _) => "!!",
                    (_, true, _, _) => "!",
                    (_, _, _, true) => "i",
                    _ => " "
                };

                let layer_str = msg.layer_prefix;

                println!("[{}][{}]: {}", type_str, layer_str, msg.description);
            }).ok();
        }

        let physical = PhysicalDevice::enumerate(&instance).next().unwrap();
        println!("Using device: {} (type: {:?})", physical.name(), physical.ty());

        let events_loop = EventLoop::new();
        let surface = WindowBuilder::new().build_vk_surface(&events_loop, instance.clone()).unwrap();
        let window = surface.window();

        // In a real-life application, we would probably use at least a graphics queue and a transfers
        // queue to handle data transfers in parallel. In this example we only use one queue.
        let queue_family = physical.queue_families().find(|&q| {
            q.supports_graphics() && surface.is_supported(q).unwrap_or(false)
        }).unwrap();

        let device_ext = DeviceExtensions { khr_swapchain: true, ..DeviceExtensions::none() };
        let (device, mut queues) = Device::new(physical, physical.supported_features(), &device_ext,
                                               [(queue_family, 0.5)].iter().cloned()).unwrap();
        let queue = queues.next().unwrap();

        let (swapchain, images) = {
            let caps = surface.capabilities(physical).unwrap();
            let usage = caps.supported_usage_flags;
            let alpha = caps.supported_composite_alpha.iter().next().unwrap();
            let format = caps.supported_formats[0].0;
            let inner_size = window.inner_size();
            data.dimensions = [inner_size.width, inner_size.height];

            Swapchain::new(device.clone(), surface.clone(), caps.min_image_count, format,
                           data.dimensions, 1, usage, &queue, SurfaceTransform::Identity, alpha,
                           PresentMode::Fifo, FullscreenExclusive::Default, true, ColorSpace::SrgbNonLinear).unwrap()
        };

        // what is host_cached
        let line_vertex_buffer = CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::vertex_buffer(), false, data.line_vertices.iter().cloned()).unwrap();

        let render_pass = Arc::new(vulkano::single_pass_renderpass!(
            device.clone(),
            attachments: {
                color: {
                    load: Clear,
                    store: Store,
                    format: swapchain.format(),
                    samples: 1,
                },
                depth: {
                    load: Clear,
                    store: DontCare,
                    format: Format::D16Unorm,
                    samples: 1,
                    initial_layout: ImageLayout::Undefined,
                    final_layout: ImageLayout::DepthStencilAttachmentOptimal,
                }
            },
            pass: {
                color: [color],
                depth_stencil: {depth}
            }
        ).unwrap());

        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 pipeline: Arc<dyn GraphicsPipelineAbstract + Send + Sync> =
            create_pipeline::<Vertex>(device.clone(), render_pass.clone(), false).unwrap();
        let line_pipeline: Arc<dyn GraphicsPipelineAbstract + Send + Sync> =
            create_pipeline::<LinePoint>(device.clone(), render_pass.clone(), true).unwrap();

        let default_tex = {
            let image = image::load_from_memory_with_format(include_bytes!("../models/missing-texture.jpg"),
                                                            ImageFormat::Jpeg).unwrap().to_rgba();
            let image_data = image.into_raw().clone();

            let (image_view, future) = ImmutableImage::from_iter(
                image_data.iter().cloned(),
                Dimensions::Dim2d { width: 128, height: 128 },
                Format::R8G8B8A8Unorm,
                queue.clone(),
            ).unwrap();

            future.flush().unwrap();

            image_view
        };

        // Dynamic viewports allow us to recreate just the viewport when the window is resized
        // Otherwise we would have to recreate the whole pipeline.
        let mut dynamic_state = DynamicState { line_width: None, viewports: None, scissors: None, compare_mask: None, write_mask: None, reference: None };

        // The render pass we created above only describes the layout of our framebuffers. Before we
        // can draw we also need to create the actual framebuffers.
        let framebuffers = window_size_dependent_setup(device.clone(), &images, render_pass.clone(), &mut dynamic_state);

        let mut uniform_buffers = Vec::new();
        let uniform_buffer = vs::ty::UniformBufferObject { view: Matrix4::identity().into(), projection: Matrix4::identity().into(), time: 0.0 };

        for _ in 0..swapchain.num_images() {
            uniform_buffers.push(CpuAccessibleBuffer::from_data(
                device.clone(),
                BufferUsage::uniform_buffer_transfer_destination(),
                false,
                uniform_buffer,
            ).unwrap());
        }

        let descriptor_set_layout = pipeline.descriptor_set_layout(0).unwrap().clone();

        let descriptor_sets = uniform_buffers.iter().map(|uniform_buffer| {
            let builder = PersistentDescriptorSet::start(descriptor_set_layout.clone());
            let result: Arc<Yeet> = Arc::new(builder
                .add_buffer(uniform_buffer.clone()).unwrap()
                .add_sampled_image(default_tex.clone(), sampler.clone()).unwrap()
                .build().unwrap());
            result
        }).collect();

        data.textures.push(default_tex);

        // In the loop below we are going to submit commands to the GPU. Submitting a command produces
        // an object that implements the `GpuFuture` trait, which holds the resources for as long as
        // they are in use by the GPU.
        //
        // Destroying the `GpuFuture` blocks until the GPU is finished executing it. In order to avoid
        // that, we store the submission of the previous frame here.
        let previous_frame_end = Some(Box::new(sync::now(device.clone())) as Box<dyn GpuFuture>);

        (VulkanRenderer { game_data: data, device, framebuffers, sampler,
            dynamic_state, pipeline, line_pipeline, uniform_buffers, descriptor_sets,
            surface, swapchain, render_pass, queue, line_vertex_buffer,
            recreate_swapchain: false, debug_callback, previous_frame_end }, events_loop)
    }

    fn create_command_buffer(self: &mut Self, fb_index: usize, ubo: vs::ty::UniformBufferObject) -> Arc<AutoCommandBuffer> {
        let mut builder = AutoCommandBufferBuilder::primary_simultaneous_use(self.device.clone(), self.queue.family()).unwrap();
        builder.update_buffer(self.uniform_buffers[fb_index].clone(), ubo).unwrap();
        builder.begin_render_pass(self.framebuffers[fb_index].clone(), false, vec![ClearValue::Float([0.0, 0.0, 0.0, 1.0]), ClearValue::Depth(1.0)]).unwrap();

        for i in 0..self.game_data.game_objects.len() {
            let game_object = &self.game_data.game_objects[i];
            let mesh = &self.game_data.meshes[game_object.mesh_index];
            self.game_data.push_constants.model = game_object.model_matrix.into();

            builder.draw_indexed(
                self.pipeline.clone(),
                &self.dynamic_state,
                vec![mesh.vertex_buffer.clone()],
                mesh.index_buffer.clone(),
                self.descriptor_sets[fb_index].clone(),
                self.game_data.push_constants.clone()).unwrap();
        }

        builder.draw(self.line_pipeline.clone(), &self.dynamic_state, vec![self.line_vertex_buffer.clone()], (), self.game_data.line_push_constants.clone()).unwrap()
            .end_render_pass().unwrap();

        Arc::new(builder.build().unwrap())
    }

    pub fn render_loop(self: &mut Self, new_ubo: vs::ty::UniformBufferObject) {
        // It is important to call this function from time to time, otherwise resources will keep
        // accumulating and you will eventually reach an out of memory error.
        // Calling this function polls various fences in order to determine what the GPU has
        // already processed, and frees the resources that are no longer needed.
        self.previous_frame_end.as_mut().unwrap().cleanup_finished();

        if self.recreate_swapchain {
            let window = self.surface.window();
            let inner_size = window.inner_size();
            self.game_data.dimensions = [inner_size.width, inner_size.height];

            let (new_swapchain, new_images) = match self.swapchain.recreate_with_dimensions(self.game_data.dimensions) {
                Ok(r) => r,
                // This error tends to happen when the user is manually resizing the window.
                // Simply restarting the loop is the easiest way to fix this issue.
                Err(SwapchainCreationError::UnsupportedDimensions) => {
                    println!("Swapchain rejected: UnsupportedDimensions");
                    return;
                }
                Err(err) => panic!("{:?}", err),
            };

            self.swapchain = new_swapchain;
            // Because framebuffers contains an Arc on the old swapchain, we need to
            // recreate framebuffers as well.
            self.framebuffers = window_size_dependent_setup(self.device.clone(), &new_images, self.render_pass.clone(), &mut self.dynamic_state);

            self.recreate_swapchain = false;
        }

        if self.game_data.recreate_pipeline {
            if let Some(pipeline_ok) = create_pipeline::<Vertex>(self.device.clone(), self.render_pass.clone(), false) {
                self.pipeline = pipeline_ok;
                println!("Updated pipeline.");
            } else {
                println!("Failed to update pipeline.");
            }
            self.game_data.recreate_pipeline = false;
        }

        // Before we can draw on the output, we have to *acquire* an image from the swapchain. If
        // no image is available (which happens if you submit draw commands too quickly), then the
        // function will block.
        // This operation returns the index of the image that we are allowed to draw upon.
        //
        // This function can block if no image is available. The parameter is an optional timeout
        // after which the function call will return an error.
        let (fb_index, _, acquire_future) = match swapchain::acquire_next_image(self.swapchain.clone(), None) {
            Ok(r) => r,
            Err(AcquireError::OutOfDate) => {
                self.recreate_swapchain = true;
                return;
            },
            Err(err) => panic!("{:?}", err)
        };

        let command_buffer = self.create_command_buffer(fb_index, new_ubo).clone();

        let future = self.previous_frame_end.take().unwrap()
            .join(acquire_future)
            .then_execute(self.queue.clone(), command_buffer).unwrap()
            .then_swapchain_present(self.queue.clone(), self.swapchain.clone(), fb_index)
            .then_signal_fence_and_flush();

        match future {
            Ok(future) => {
                // we're joining on the previous future but the CPU is running faster than the GPU so
                // eventually it stutters, and jumps ahead to the newer frames.
                //
                // See vulkano issue 1135: https://github.com/vulkano-rs/vulkano/issues/1135
                // This makes sure the CPU stays in sync with the GPU in situations when the CPU is
                // running "too fast"
                #[cfg(target_os = "macos")]
                future.wait(None).unwrap();

                self.previous_frame_end = Some(Box::new(future) as Box<_>);
            },
            Err(FlushError::OutOfDate) => {
                println!("Swapchain out of date!");
                self.recreate_swapchain = true;
                self.previous_frame_end = Some(Box::new(sync::now(self.device.clone())) as Box<_>);
            }
            Err(e) => {
                println!("{:?}", e);
                self.previous_frame_end = Some(Box::new(sync::now(self.device.clone())) as Box<_>);
            }
        };
    }

    pub fn upload_mesh(self: &mut Self, mesh: CPUMesh) -> usize {
        let vertex_buffer = CpuAccessibleBuffer::from_iter(self.device.clone(), BufferUsage::vertex_buffer(), false, mesh.vertices.into_iter()).unwrap();
        let index_buffer  = CpuAccessibleBuffer::from_iter(self.device.clone(), BufferUsage::index_buffer(),  false, mesh.indices.into_iter()).unwrap();
        self.game_data.meshes.push(Mesh { vertex_buffer, index_buffer });
        self.game_data.meshes.len() - 1
    }

    pub fn upload_texture(self: &mut Self, texture_data: &gltf::image::Data) {
        let buffer: ImageBuffer<Rgb<u8>, Vec<u8>> = image::ImageBuffer::from_raw(texture_data.width, texture_data.height, texture_data.pixels.clone()).unwrap();
        let new_buffer: ImageBuffer<Rgba<u8>, Vec<u8>> = buffer.convert();

        let (image_view, future) = ImmutableImage::from_iter(
            new_buffer.iter().cloned(),
            Dimensions::Dim2d { width: texture_data.width, height: texture_data.height },
            Format::R8G8B8A8Unorm,
            self.queue.clone(),
        ).unwrap();

        future.flush().unwrap();

        self.game_data.textures.push(image_view);
    }

    pub fn add_game_object(self: &mut Self, game_object: GameObject) -> usize {
        self.game_data.game_objects.push(game_object);
        self.game_data.game_objects.len() - 1
    }
}

pub fn start_event_loop(mut renderer: VulkanRenderer, mut game: Box<dyn Game>, event_loop: EventLoop<()>) {
    let mut recreate_swapchain = false;
    let ubo = game.update(&mut renderer);

    event_loop.run(move |event, _, control_flow| {
        game.on_window_event(&event);
        match event {
            Event::WindowEvent { event: WindowEvent::CloseRequested, .. } => {
                *control_flow = ControlFlow::Exit;
            },
            Event::WindowEvent { event: WindowEvent::Resized(_), .. } => {
                recreate_swapchain = true;
            },
            Event::RedrawEventsCleared => {
                game.update(&mut renderer);
                renderer.render_loop(ubo);
            },
            _ => {}
        }
    });
}

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

    let viewport = Viewport {
        origin: [0.0, 0.0],
        dimensions: [dimensions[0] as f32, dimensions[1] as f32],
        depth_range: 0.0 .. 1.0,
    };
    dynamic_state.viewports = Some(vec!(viewport));

    let depth_image = AttachmentImage::with_usage(device.clone(), dimensions, Format::D16Unorm, ImageUsage { depth_stencil_attachment: true, ..ImageUsage::none() }).unwrap();

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

pub mod vs {
    vulkano_shaders::shader!{
        ty: "vertex",
        path: "shaders/triangle.vert"
    }
}

pub mod line_vs {
    vulkano_shaders::shader!{
        ty: "vertex",
        path: "shaders/line.vert"
    }
}

pub mod fs {
    vulkano_shaders::shader! {
        ty: "fragment",
        path: "shaders/triangle.frag"
    }
}

pub mod line_fs {
    vulkano_shaders::shader! {
        ty: "fragment",
        path: "shaders/line.frag"
    }
}

fn create_pipeline<V: vulkano::pipeline::vertex::Vertex>(device: Arc<Device>, render_pass: Arc<dyn RenderPassAbstract + Send + Sync>, is_line: bool) -> Option<Arc<dyn GraphicsPipelineAbstract + Send + Sync>> {
    let sub_pass = Subpass::from(render_pass.clone(), 0).unwrap();
    let pipeline;

    if is_line {
        let vertex_shader = line_vs::Shader::load(device.clone()).unwrap();
        let fragment_shader = line_fs::Shader::load(device.clone()).unwrap();

        pipeline = Arc::new(GraphicsPipeline::start()
            .vertex_input_single_buffer::<V>()
            .vertex_shader(vertex_shader.main_entry_point(), ())
            .line_list()
            .viewports_dynamic_scissors_irrelevant(1)
            .depth_stencil_simple_depth()
            .fragment_shader(fragment_shader.main_entry_point(), ())
            .render_pass(sub_pass.clone())
            .build(device.clone())
            .unwrap());
    } else {
        let vertex_shader = vs::Shader::load(device.clone()).unwrap();
        let fragment_shader = fs::Shader::load(device.clone()).unwrap();

        pipeline = Arc::new(GraphicsPipeline::start()
            .vertex_input_single_buffer::<V>()
            .vertex_shader(vertex_shader.main_entry_point(), ())
            .triangle_list()
            .viewports_dynamic_scissors_irrelevant(1)
            .depth_stencil_simple_depth()
            .fragment_shader(fragment_shader.main_entry_point(), ())
            .blend_alpha_blending()
            .cull_mode_back()
            .render_pass(sub_pass.clone())
            .build(device.clone())
            .unwrap());
    }

    return Some(pipeline);
}

impl GameObject {
    pub fn new(mesh: MeshHandle, texture_index: usize) -> GameObject {
        GameObject { mesh_index: mesh, texture_index, model_matrix: Matrix4::identity() }
    }
}