use crate::prelude::*;
use anyhow::Result;
use std::sync::{Arc, Mutex};

pub enum ClearValue {
    Color([f32; 4]),
    Depth(f32, u32),
}

pub struct CustomTarget {
    pub usage: VkImageUsageFlagBits,
    pub format: VkFormat,
    pub clear_on_load: bool,
    pub store_on_save: bool,
    pub attach_sampler: bool,
    pub use_as_input: bool,
    pub clear_value: ClearValue,
}

impl CustomTarget {
    pub fn depth() -> CustomTarget {
        CustomTarget {
            usage: VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT.into(),
            format: VK_FORMAT_D16_UNORM,
            clear_on_load: true,
            store_on_save: false,
            attach_sampler: false,
            use_as_input: false,
            clear_value: ClearValue::Depth(1.0, 0),
        }
    }

    fn to_attachment_info(
        &self,
        device: &Arc<Device>,
        queue: &Arc<Mutex<Queue>>,
        width: u32,
        height: u32,
        sample_count: VkSampleCountFlags,
    ) -> Result<AttachmentInfo> {
        let clear_operation = SubPassBuilder::clear_op(self.clear_on_load);
        let store_operation = SubPassBuilder::store_op(self.store_on_save);

        // set clear values
        let clear_value = match self.clear_value {
            ClearValue::Color(color) => VkClearValue::color(VkClearColorValue::float32(color)),
            ClearValue::Depth(depth, stencil) => {
                VkClearValue::depth_stencil(VkClearDepthStencilValue { depth, stencil })
            }
        };

        // check for color attachment flag
        let (format, aspect, description, usage, layout) =
            if (self.usage & VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT) != 0 {
                // set color attachment
                let description = VkAttachmentDescription::new(
                    0,
                    self.format,
                    sample_count,
                    clear_operation,
                    store_operation,
                    VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                    VK_ATTACHMENT_STORE_OP_DONT_CARE,
                    VK_IMAGE_LAYOUT_UNDEFINED,
                    if self.attach_sampler || self.use_as_input {
                        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
                    } else {
                        VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL
                    },
                );

                (
                    self.format,
                    VK_IMAGE_ASPECT_COLOR_BIT,
                    description,
                    AttachmentInfoUsage::Output,
                    VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
                )
            // check for depth attachment flag
            } else if (self.usage & VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT) != 0 {
                // set depth attachment
                let description = VkAttachmentDescription::new(
                    0,
                    self.format,
                    sample_count,
                    clear_operation,
                    store_operation,
                    VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                    VK_ATTACHMENT_STORE_OP_DONT_CARE,
                    VK_IMAGE_LAYOUT_UNDEFINED,
                    VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
                );

                // take format and aspect mask
                (
                    self.format,
                    VK_IMAGE_ASPECT_DEPTH_BIT,
                    description,
                    AttachmentInfoUsage::Depth,
                    VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
                )
            } else {
                // TODO: add more as required
                unimplemented!();
            };

        let mut image_builder = Image::empty(width, height, self.usage, sample_count)
            .format(format)
            .aspect_mask(aspect);

        if self.attach_sampler {
            image_builder = image_builder.attach_sampler(Sampler::nearest_sampler().build(device)?);
        }

        let image = image_builder.build(device, queue)?;

        match aspect {
            VK_IMAGE_ASPECT_DEPTH_BIT => {
                Image::convert_layout(&image, VK_IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL)?
            }
            VK_IMAGE_ASPECT_COLOR_BIT => {
                Image::convert_layout(&image, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL)?
            }
            _ => unimplemented!(),
        }

        Ok(AttachmentInfo {
            images: vec![image],
            clear_value,
            layout,
            description,
            usage,
        })
    }
}

pub enum ResolveTarget {
    CustomTarget(CustomTarget),
    PreparedTargets(Vec<Arc<Image>>),
}

impl From<CustomTarget> for ResolveTarget {
    fn from(custom_target: CustomTarget) -> Self {
        Self::CustomTarget(custom_target)
    }
}

impl From<Vec<Arc<Image>>> for ResolveTarget {
    fn from(prepared_targets: Vec<Arc<Image>>) -> Self {
        Self::PreparedTargets(prepared_targets)
    }
}

impl<'a> From<&'a Vec<Arc<Image>>> for ResolveTarget {
    fn from(prepared_targets: &'a Vec<Arc<Image>>) -> Self {
        Self::PreparedTargets(prepared_targets.clone())
    }
}

pub struct InputAttachmentInfo {
    pub sub_pass_index: usize,
    pub input_indices: Vec<usize>,
}

pub struct SubPassBuilder<'a> {
    width: u32,
    height: u32,
    sample_count: VkSampleCountFlags,

    target_infos: Vec<CustomTarget>,

    input_info: Option<InputAttachmentInfo>,

    // (images, index, clear_color, clear_on_load)
    prepared_targets: Option<(&'a [Arc<Image>], usize, [f32; 4], bool)>,
    resolve_targets: Vec<ResolveTarget>,

    output_usage: VkAccessFlagBits,
}

impl<'a> SubPassBuilder<'a> {
    pub fn set_sample_count(mut self, sample_count: VkSampleCountFlags) -> Self {
        self.sample_count = sample_count;

        self
    }

    pub fn add_target_info(mut self, target: CustomTarget) -> Self {
        self.target_infos.push(target);

        self
    }

    pub fn set_input_attachment_info(mut self, input_info: InputAttachmentInfo) -> Self {
        self.input_info = Some(input_info);

        self
    }

    pub fn set_output_usage(mut self, output_usage: impl Into<VkAccessFlagBits>) -> Self {
        self.output_usage = output_usage.into();

        self
    }

    pub fn set_prepared_targets(
        mut self,
        prepared_targets: &'a [Arc<Image>],
        target_index: usize,
        clear_color: impl Into<[f32; 4]>,
        clear_on_load: bool,
    ) -> Self {
        self.prepared_targets = Some((
            prepared_targets,
            target_index,
            clear_color.into(),
            clear_on_load,
        ));

        self
    }

    pub fn add_resolve_targets(mut self, resolve_target: impl Into<ResolveTarget>) -> Self {
        self.resolve_targets.push(resolve_target.into());

        self
    }

    pub fn build(self, device: &Arc<Device>, queue: &Arc<Mutex<Queue>>) -> Result<SubPass> {
        let attachments = self.create_images(device, queue)?;

        Ok(SubPass {
            extent: VkExtent2D {
                width: self.width,
                height: self.height,
            },

            input_info: self.input_info,
            attachments,

            output_usage: self.output_usage,
        })
    }

    #[inline]
    fn create_images(
        &self,
        device: &Arc<Device>,
        queue: &Arc<Mutex<Queue>>,
    ) -> Result<Vec<AttachmentInfo>> {
        // check for correct sample count
        let checked_sample_count = device.max_supported_sample_count(self.sample_count);

        // throw an error if we don't use muultisampling and have an resolve target
        if checked_sample_count == VK_SAMPLE_COUNT_1_BIT && !self.resolve_targets.is_empty() {
            panic!("Sample count 1 and using resolve target is not supported");
        }

        let mut attachment_infos = Vec::new();

        for target_info in self.target_infos.iter() {
            attachment_infos.push(target_info.to_attachment_info(
                device,
                queue,
                self.width,
                self.height,
                self.sample_count,
            )?);
        }

        // insert prepared images
        if let Some((prepared_images, index, clear_color, clear_on_load)) = self.prepared_targets {
            let clear_operation = Self::clear_op(clear_on_load);

            attachment_infos.insert(
                index,
                AttachmentInfo {
                    images: prepared_images.iter().map(|image| image.clone()).collect(),
                    clear_value: VkClearValue::color(VkClearColorValue::float32(clear_color)),
                    layout: VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
                    description: VkAttachmentDescription::new(
                        0,
                        prepared_images[0].vk_format(),
                        VK_SAMPLE_COUNT_1_BIT,
                        clear_operation,
                        VK_ATTACHMENT_STORE_OP_STORE,
                        VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                        VK_ATTACHMENT_STORE_OP_DONT_CARE,
                        prepared_images[0].image_layout(),
                        prepared_images[0].image_layout(),
                    ),
                    usage: AttachmentInfoUsage::Output,
                },
            );
        }

        // add resolve target if possible
        for resolve_target in self.resolve_targets.iter() {
            match resolve_target {
                ResolveTarget::CustomTarget(custom_target) => {
                    let mut attachment_info = custom_target.to_attachment_info(
                        device,
                        queue,
                        self.width,
                        self.height,
                        VK_SAMPLE_COUNT_1_BIT,
                    )?;

                    attachment_info.usage = AttachmentInfoUsage::Resolve;

                    attachment_infos.push(attachment_info);
                }
                ResolveTarget::PreparedTargets(prepared_targets) => {
                    attachment_infos.push(AttachmentInfo {
                        images: prepared_targets.iter().map(|image| image.clone()).collect(),
                        clear_value: VkClearValue::color(VkClearColorValue::float32([
                            0.0, 0.0, 0.0, 1.0,
                        ])),
                        layout: VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
                        description: VkAttachmentDescription::new(
                            0,
                            prepared_targets[0].vk_format(),
                            VK_SAMPLE_COUNT_1_BIT,
                            VK_ATTACHMENT_LOAD_OP_CLEAR,
                            VK_ATTACHMENT_STORE_OP_STORE,
                            VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                            VK_ATTACHMENT_STORE_OP_DONT_CARE,
                            VK_IMAGE_LAYOUT_UNDEFINED,
                            prepared_targets[0].image_layout(),
                        ),
                        usage: AttachmentInfoUsage::Resolve,
                    });
                }
            }
        }

        Ok(attachment_infos)
    }

    #[inline]
    fn clear_op(clear_on_load: bool) -> VkAttachmentLoadOp {
        if clear_on_load {
            VK_ATTACHMENT_LOAD_OP_CLEAR
        } else {
            VK_ATTACHMENT_LOAD_OP_LOAD
        }
    }

    #[inline]
    fn store_op(store_on_save: bool) -> VkAttachmentStoreOp {
        if store_on_save {
            VK_ATTACHMENT_STORE_OP_STORE
        } else {
            VK_ATTACHMENT_STORE_OP_DONT_CARE
        }
    }
}

#[derive(Eq, PartialEq, Hash, Clone, Copy)]
pub enum AttachmentInfoUsage {
    Depth,
    Resolve,
    Output,
}

pub struct AttachmentInfo {
    images: Vec<Arc<Image>>,
    pub(crate) clear_value: VkClearValue,
    pub(crate) layout: VkImageLayout,
    pub(crate) description: VkAttachmentDescription,
    pub(crate) usage: AttachmentInfoUsage,
}

impl AttachmentInfo {
    pub fn image(&self, mut index: usize) -> &Arc<Image> {
        debug_assert!(!self.images.is_empty());

        if index >= self.images.len() {
            index = self.images.len() - 1;
        }

        &self.images[index]
    }
}

pub struct SubPass {
    extent: VkExtent2D,

    input_info: Option<InputAttachmentInfo>,
    attachments: Vec<AttachmentInfo>,

    output_usage: VkAccessFlagBits,
}

impl SubPass {
    pub fn builder<'a>(width: u32, height: u32) -> SubPassBuilder<'a> {
        SubPassBuilder {
            width,
            height,
            sample_count: VK_SAMPLE_COUNT_1_BIT,

            input_info: None,

            target_infos: Vec::new(),

            prepared_targets: None,
            resolve_targets: Vec::new(),

            output_usage: VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT
                | VK_ACCESS_COLOR_ATTACHMENT_READ_BIT,
        }
    }

    pub(crate) fn inputs(&self) -> Option<&InputAttachmentInfo> {
        self.input_info.as_ref()
    }

    pub(crate) fn output_usage(&self) -> VkAccessFlagBits {
        self.output_usage
    }

    pub fn extent(&self) -> VkExtent2D {
        self.extent
    }

    pub fn attachments(&self) -> &[AttachmentInfo] {
        &self.attachments
    }

    pub fn max_images_per_attachment(&self) -> usize {
        let mut max_images = 0;

        for attachment in self.attachments.iter() {
            max_images = max_images.max(attachment.images.len());
        }

        max_images
    }
}