use std::{collections::HashMap, iter};

use ash::{
    khr::{surface, swapchain},
    vk::{self, DescriptorSet},
};

use crate::{
    Game, InWorld, LoadingWorld, RenderAvailable, RenderCtx, game::Camera, render::MemType,
};

use super::{
    MAX_HEIGHT, MAX_WIDTH, SwapchainCtx,
    swapchain::{
        make_depth_img, make_depth_view, make_framebufs, make_swap_images, make_swap_views,
        setup_swapchain,
    },
};

/// Anything can be rendered in world
pub trait WorldComponent {
    fn descriptors(&self) -> HashMap<vk::DescriptorType, u32>;

    fn desc_layout(&self, ctx: &RenderCtx) -> vk::DescriptorSetLayout;

    /// Skybox is passed in for reflections
    fn write_desc_set(
        &self,
        ctx: &RenderCtx,
        desc_set: vk::DescriptorSet,
        skybox: (vk::ImageView, vk::Sampler),
    );

    fn pipeline(&self) -> (vk::Pipeline, vk::PipelineLayout);

    fn push_constants(
        &self,
        camera: &Camera,
        res: [u32; 2],
        base_color: [f32; 4],
    ) -> Vec<(vk::ShaderStageFlags, u32, Vec<u8>)>;
}

pub struct Component {
    pub inner: Box<dyn WorldComponent>,
    pub desc_set: DescriptorSet,
}

impl Game<RenderAvailable> {
    pub fn start_load_world(self) -> Game<LoadingWorld> {
        let Self {
            state:
                RenderAvailable {
                    sdl_context,
                    ctx,
                    window,
                },
        } = self;

        let depth_mem = ctx.mem_alloc(MemType::HostInvisible, (MAX_WIDTH * MAX_HEIGHT * 8) as u64);

        let pass = setup_render_pass(&ctx);

        let surface_loader = surface::Instance::new(&ctx.entry, &ctx.instance);
        let swapchain_loader = swapchain::Device::new(&ctx.instance, &ctx.dev);

        let (swapchain, extents) = setup_swapchain(
            &ctx,
            &surface_loader,
            &swapchain_loader,
            window.width,
            window.height,
        );

        let depth_img = make_depth_img(&ctx, window.width, window.height, depth_mem);
        let swap_images = make_swap_images(swapchain, &swapchain_loader);
        let swap_views = make_swap_views(&ctx, &swap_images, vk::Format::B8G8R8A8_UNORM);
        let depth_view = make_depth_view(&ctx, depth_img, vk::Format::D32_SFLOAT);
        let framebufs = make_framebufs(
            &ctx,
            &swap_views,
            &depth_view,
            &pass,
            window.width,
            window.height,
        );

        let skybox = ctx.create_skybox(pass);

        let requested_descriptors = skybox.descriptors();

        Game {
            state: LoadingWorld {
                sdl_context,
                ctx,
                window,

                requested_descriptors,

                swapchain: SwapchainCtx {
                    surface_loader,
                    swapchain_loader,
                    swap_images,
                    swap_views,
                    depth_mem,
                    depth_image: depth_img,
                    depth_view,
                    framebufs,
                    swapchain,
                },
                pass,

                skybox,
                components: Vec::new(),
            },
        }
    }
}

impl Game<LoadingWorld> {
    pub fn add_component<C: WorldComponent + 'static>(&mut self, component: C) {
        let Self {
            state:
                LoadingWorld {
                    requested_descriptors,
                    components,
                    ..
                },
        } = self;
        for (desc, n) in component.descriptors() {
            match requested_descriptors.remove(&desc) {
                Some(m) => requested_descriptors.insert(desc, n + m),
                None => requested_descriptors.insert(desc, n),
            };
        }
        components.push(Box::new(component));
    }

    pub fn start_world(self) -> Game<InWorld> {
        let Self {
            state:
                LoadingWorld {
                    sdl_context,
                    window,
                    ctx,

                    requested_descriptors,

                    swapchain,
                    pass,

                    skybox,
                    components,
                },
        } = self;
        let desc_pool = setup_desc_pool(&ctx, requested_descriptors);

        let mut desc_sets = make_desc_sets(
            &ctx,
            &desc_pool,
            &components
                .iter()
                .map(|c| c.desc_layout(&ctx))
                .chain(iter::once(skybox.desc_layout(&ctx)))
                .collect::<Vec<_>>(),
        );

        let skybox_descriptor = desc_sets.pop().unwrap();

        let components = components
            .into_iter()
            .zip(desc_sets.into_iter())
            .map(|(c, d)| Component {
                inner: c,
                desc_set: d,
            })
            .collect::<Vec<_>>();

        for comp in &components {
            let Component { inner, desc_set } = comp;
            inner.write_desc_set(&ctx, *desc_set, (skybox.imageview, skybox.sampler));
        }

        skybox.write_desc_set(&ctx, skybox_descriptor, (skybox.imageview, skybox.sampler));

        let sem_avail = ctx.make_sem();
        let sem_finish = ctx.make_sem();
        let fence_flight = ctx.make_fence();

        Game {
            state: InWorld {
                sdl_context,
                window,
                ctx,

                skybox: (skybox, skybox_descriptor),
                components,

                sem_avail,
                sem_finish,
                fence_flight,

                swapchain,
                pass,

                camera: Camera::default(),
            },
        }
    }
}

fn setup_render_pass(ctx: &RenderCtx) -> vk::RenderPass {
    let color_attach = vk::AttachmentDescription::default()
        .format(vk::Format::B8G8R8A8_UNORM)
        .samples(vk::SampleCountFlags::TYPE_1)
        .load_op(vk::AttachmentLoadOp::CLEAR)
        .store_op(vk::AttachmentStoreOp::STORE)
        .stencil_load_op(vk::AttachmentLoadOp::DONT_CARE)
        .stencil_store_op(vk::AttachmentStoreOp::DONT_CARE)
        .initial_layout(vk::ImageLayout::UNDEFINED)
        .final_layout(vk::ImageLayout::PRESENT_SRC_KHR);
    let depth_attach = vk::AttachmentDescription::default()
        .format(vk::Format::D32_SFLOAT)
        .samples(vk::SampleCountFlags::TYPE_1)
        .load_op(vk::AttachmentLoadOp::CLEAR)
        .store_op(vk::AttachmentStoreOp::DONT_CARE)
        .stencil_load_op(vk::AttachmentLoadOp::DONT_CARE)
        .stencil_store_op(vk::AttachmentStoreOp::DONT_CARE)
        .initial_layout(vk::ImageLayout::UNDEFINED)
        .final_layout(vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL);

    let color_ref = vk::AttachmentReference::default()
        .attachment(0)
        .layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL);

    let depth_ref = vk::AttachmentReference::default()
        .attachment(1)
        .layout(vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL);

    let color_refs = [color_ref];
    let depth_refs = [depth_ref];

    let subpass = vk::SubpassDescription::default()
        .pipeline_bind_point(vk::PipelineBindPoint::GRAPHICS)
        .color_attachments(color_refs.as_ref())
        .depth_stencil_attachment(&depth_ref);

    let subpass_dep = vk::SubpassDependency::default()
        .src_subpass(vk::SUBPASS_EXTERNAL)
        .dst_subpass(0)
        .src_stage_mask(
            vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT
                | vk::PipelineStageFlags::EARLY_FRAGMENT_TESTS,
        )
        .dst_stage_mask(
            vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT
                | vk::PipelineStageFlags::EARLY_FRAGMENT_TESTS,
        )
        .src_access_mask(vk::AccessFlags::empty())
        .dst_access_mask(
            vk::AccessFlags::COLOR_ATTACHMENT_WRITE
                | vk::AccessFlags::DEPTH_STENCIL_ATTACHMENT_WRITE,
        );

    let attach_desc = [color_attach, depth_attach];

    let subpasses = [subpass];

    let deps = [subpass_dep];

    let pass_info = vk::RenderPassCreateInfo::default()
        .attachments(&attach_desc)
        .subpasses(subpasses.as_ref())
        .dependencies(deps.as_ref());

    let pass = unsafe {
        ctx.dev
            .create_render_pass(&pass_info, None)
            .expect("Failed to create render pass")
    };
    pass
}

fn setup_desc_pool(
    ctx: &RenderCtx,
    requested_descriptors: HashMap<vk::DescriptorType, u32>,
) -> vk::DescriptorPool {
    let pool_sizes = requested_descriptors
        .into_iter()
        .map(|(desc, count)| {
            vk::DescriptorPoolSize::default()
                .ty(desc)
                .descriptor_count(count)
        })
        .collect::<Vec<_>>();

    let pool_info = vk::DescriptorPoolCreateInfo::default()
        .max_sets(1)
        .pool_sizes(&pool_sizes);
    let pool = unsafe {
        ctx.dev
            .create_descriptor_pool(&pool_info, None)
            .expect("Failed to create descriptor pool")
    };
    pool
}

fn make_desc_sets(
    ctx: &RenderCtx,
    desc_pool: &vk::DescriptorPool,
    layouts: &[vk::DescriptorSetLayout],
) -> Vec<vk::DescriptorSet> {
    let alloc_info = vk::DescriptorSetAllocateInfo::default()
        .descriptor_pool(*desc_pool)
        .set_layouts(&layouts);
    let sets = unsafe {
        ctx.dev
            .allocate_descriptor_sets(&alloc_info)
            .expect("Failed to allocate descriptor sets")
    };
    sets
}