mod cube;

use std::{path::PathBuf, sync::Arc, time::Duration};

use anyhow::Result;
use cube::{Cube, CubeCorners};
use ecs::*;
use engine::prelude::{cgmath::*, shader_type::*, *};

pub struct SkyBoxImages {
    left: PathBuf,
    right: PathBuf,
    front: PathBuf,
    back: PathBuf,
    top: PathBuf,
    bottom: PathBuf,
}

impl<T: ExactSizeIterator<Item = PathBuf>> From<T> for SkyBoxImages {
    fn from(mut paths: T) -> Self {
        debug_assert_eq!(paths.len(), 6);

        Self {
            left: paths.next().unwrap(),
            right: paths.next().unwrap(),
            front: paths.next().unwrap(),
            back: paths.next().unwrap(),
            top: paths.next().unwrap(),
            bottom: paths.next().unwrap(),
        }
    }
}

pub struct SkyBox {
    enabled: bool,

    _cube_map: Arc<Image>,
    cube_buffer: Arc<Buffer<PositionOnly>>,

    vertex_shader: Arc<ShaderModule<Vertex>>,
    fragment_shader: Arc<ShaderModule<Fragment>>,

    render_target: TargetMode<RenderTarget>,
    pipeline: TargetMode<Arc<Pipeline>>,
    descriptor_set: Arc<DescriptorSet>,
}

impl SkyBox {
    pub fn new(world: &mut WorldBuilder, images: impl Into<SkyBoxImages>) -> Result<()> {
        let sample_count = world
            .resources
            .get::<EngineSettings>()
            .graphics_info()?
            .sample_count;

        let context = world.resources.get_mut_unchecked::<Context>();
        context.render_core_mut().add_render_routine::<Self>(10);

        let images = images.into();
        let cube_map = Image::cube_map([
            (
                images.right.try_into()?,
                vec![ImageModifier::FlipV, ImageModifier::Rotate90],
            ),
            (
                images.left.try_into()?,
                vec![ImageModifier::Rotate90, ImageModifier::FlipV],
            ),
            (images.front.try_into()?, vec![ImageModifier::FlipV]),
            (images.back.try_into()?, vec![ImageModifier::FlipH]),
            (images.top.try_into()?, vec![ImageModifier::FlipV]),
            (images.bottom.try_into()?, vec![ImageModifier::FlipH]),
        ])?
        .format(VK_FORMAT_R8G8B8A8_UNORM)
        .max_mip_map_levels()
        .attach_pretty_sampler(context.device())?
        .build(context.device(), context.queue())?;

        let descriptor_set_layout = DescriptorSetLayout::builder()
            .add_layout_binding(
                0,
                VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                VK_SHADER_STAGE_VERTEX_BIT,
                0,
            )
            .add_layout_binding(
                1,
                VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                VK_SHADER_STAGE_FRAGMENT_BIT,
                0,
            )
            .build(context.device().clone())?;

        let render_target = Self::create_render_target(context, &context.images(), sample_count)?;

        let pipeline_layout = PipelineLayout::builder()
            .add_descriptor_set_layout(&descriptor_set_layout)
            .build(context.device().clone())?;

        let vertex_shader = ShaderModule::from_slice(
            context.device().clone(),
            include_bytes!("../shader/skybox.vert.spv"),
        )?;

        let fragment_shader = ShaderModule::from_slice(
            context.device().clone(),
            include_bytes!("../shader/skybox.frag.spv"),
        )?;

        let pipeline = Self::create_pipeline(
            context,
            sample_count,
            &render_target,
            &pipeline_layout,
            &vertex_shader,
            &fragment_shader,
        )?;

        let descriptor_pool = DescriptorPool::builder()
            .set_layout(descriptor_set_layout)
            .build(context.device().clone())?;
        let descriptor_set = descriptor_pool.prepare_set().allocate()?;

        let scene = world.resources.get::<Scene>();
        let view = scene.view();

        descriptor_set.update(&[
            DescriptorWrite::uniform_buffers(0, &[view.buffer()]),
            DescriptorWrite::combined_samplers(1, &[&cube_map]),
        ])?;

        let cube: Cube = CubeCorners::new(vec3(-1.0, -1.0, -1.0), vec3(2.0, 2.0, 2.0)).into();
        let cube_mesh = cube.triangulate();

        let command_buffer = CommandBuffer::new_primary()
            .build(context.device().clone(), context.queue().clone())?;

        let cube_buffer = SingleSubmit::builder(&command_buffer, context.queue(), |recorder| {
            Buffer::builder()
                .set_memory_usage(MemoryUsage::CpuToGpu)
                .set_usage(VK_BUFFER_USAGE_TRANSFER_SRC_BIT)
                .set_data(&cube_mesh)
                .build(context.device().clone())?
                .into_device_local(
                    recorder,
                    VK_ACCESS_VERTEX_ATTRIBUTE_READ_BIT,
                    VK_PIPELINE_STAGE_VERTEX_INPUT_BIT,
                    VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
                )
        })
        .wait_for_timeout(Duration::from_secs(2))
        .submit()?;

        let me = Self {
            enabled: true,

            _cube_map: cube_map,
            cube_buffer,

            vertex_shader,
            fragment_shader,

            render_target,
            pipeline,
            descriptor_set,
        };
        world.resources.insert(me);

        Ok(())
    }

    pub fn enable(&mut self) {
        self.enabled = true;
    }

    pub fn disable(&mut self) {
        self.enabled = false;
    }

    fn create_render_target(
        context: &Context,
        images: &TargetMode<Vec<Arc<Image>>>,
        sample_count: VkSampleCountFlags,
    ) -> Result<TargetMode<RenderTarget>> {
        images.execute(|images| {
            let first = images.first().unwrap();

            let width = first.width();
            let height = first.height();
            let format = first.vk_format();

            if sample_count == VK_SAMPLE_COUNT_1_BIT {
                RenderTarget::builder()
                    .add_sub_pass(
                        SubPass::builder(width, height)
                            // render directly into swapchain target if there is no multi sampling
                            .set_prepared_targets(images, 0, None)
                            .use_queue(context.queue().clone())
                            .build(context.device())?,
                    )
                    .build(context.device())
            } else {
                RenderTarget::builder()
                    .add_sub_pass(
                        SubPass::builder(width, height)
                            // render into multi sampled images
                            .add_target_info(CustomTarget {
                                usage: VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT.into(),
                                format,
                                clear_on_load: true,
                                store_on_save: true,
                                attach_sampler: false,
                                use_as_input: false,
                                clear_value: ClearValue::Color([0.0, 0.0, 0.0, 0.0]),
                            })
                            .set_sample_count(sample_count)
                            // resolve multi sampling into swapchain image
                            .add_resolve_targets((images, false))
                            .use_queue(context.queue().clone())
                            .build(context.device())?,
                    )
                    .build(context.device())
            }
        })
    }

    fn create_pipeline(
        context: &Context,
        sample_count: VkSampleCountFlags,
        render_target: &TargetMode<RenderTarget>,
        pipeline_layout: &Arc<PipelineLayout>,
        vertex_shader: &Arc<ShaderModule<Vertex>>,
        fragment_shader: &Arc<ShaderModule<Fragment>>,
    ) -> Result<TargetMode<Arc<Pipeline>>> {
        render_target.execute(|render_target| {
            Pipeline::new_graphics()
                .set_vertex_shader::<PositionOnly>(vertex_shader.clone())
                .set_fragment_shader(fragment_shader.clone())
                .input_assembly(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST, false)
                .default_multisample(sample_count)
                .default_color_blend(vec![VkPipelineColorBlendAttachmentState::default()])
                .default_rasterization(VK_CULL_MODE_NONE, VK_FRONT_FACE_CLOCKWISE)
                .whole_area(render_target.width(), render_target.height())
                .build(
                    context.device().clone(),
                    pipeline_layout,
                    render_target.render_pass(),
                    0,
                )
        })
    }
}

impl TScene for SkyBox {
    fn process(
        &mut self,
        buffer_recorder: &mut CommandBufferRecorder<'_>,
        _images: &TargetMode<Vec<Arc<Image>>>,
        indices: &TargetMode<usize>,
        _world: &mut World,
    ) -> Result<()> {
        if !self.enabled {
            return Ok(());
        }

        self.render_target
            .chain(indices)
            .chain(&self.pipeline)
            .unfold()
            .execute(|(render_target, index, pipeline)| {
                render_target.begin(buffer_recorder, VK_SUBPASS_CONTENTS_INLINE, ***index);

                buffer_recorder.bind_pipeline(pipeline)?;
                buffer_recorder.bind_descriptor_sets_minimal(&[&self.descriptor_set]);

                buffer_recorder.bind_vertex_buffer(&self.cube_buffer);
                buffer_recorder.draw_complete_single_instance(self.cube_buffer.size() as u32);

                render_target.end(buffer_recorder);

                Ok(())
            })?;

        Ok(())
    }

    fn resize(
        &mut self,
        _window_width: f32,
        _window_height: f32,
        images: &TargetMode<Vec<Arc<Image>>>,
        world: &mut World,
    ) -> Result<()> {
        let sample_count = world
            .resources
            .get::<EngineSettings>()
            .graphics_info()?
            .sample_count;

        let context = world.resources.get::<Context>();

        let pipeline_layout = match &self.pipeline {
            TargetMode::Mono(p) => p.pipeline_layout().clone(),
            TargetMode::Stereo(p, _) => p.pipeline_layout().clone(),
        };

        self.render_target = Self::create_render_target(context, images, sample_count)?;
        self.pipeline = Self::create_pipeline(
            context,
            sample_count,
            &self.render_target,
            &pipeline_layout,
            &self.vertex_shader,
            &self.fragment_shader,
        )?;

        Ok(())
    }
}