//! Defines node - building block for framegraph.

//!


pub mod present;
pub mod render;

use {
    crate::{
        command::{Capability, Families, Family, FamilyId, Fence, Queue, Submission, Submittable},
        factory::{Factory, UploadError},
        frame::Frames,
        graph::GraphContext,
        wsi::SwapchainError,
        BufferId, ImageId, NodeId,
    },
    rendy_core::hal::{queue::QueueFamilyId, Backend},
};

/// Buffer access node will perform.

/// Node must not perform any access to the buffer not specified in `access`.

/// All access must be between logically first and last `stages`.

#[derive(Clone, Copy, Debug)]
pub struct BufferAccess {
    /// Access flags.

    pub access: rendy_core::hal::buffer::Access,

    /// Intended usage flags for buffer.

    /// TODO: Could derive from access?

    pub usage: rendy_core::hal::buffer::Usage,

    /// Pipeline stages at which buffer is accessd.

    pub stages: rendy_core::hal::pso::PipelineStage,
}

/// Buffer pipeline barrier.

#[derive(Clone, Debug)]
pub struct BufferBarrier {
    /// State transition for the buffer.

    pub states: std::ops::Range<rendy_core::hal::buffer::State>,

    /// Stages at which buffer is accessd.

    pub stages: std::ops::Range<rendy_core::hal::pso::PipelineStage>,

    /// Transfer between families.

    pub families: Option<std::ops::Range<QueueFamilyId>>,
}

/// Buffer shared between nodes.

///

/// If Node doesn't actually use the buffer it can merge acquire and release barriers into one.

/// TODO: Make merge function.

#[derive(Clone, Debug)]
pub struct NodeBuffer {
    /// Id of the buffer.

    pub id: BufferId,

    /// Region of the buffer that is the transient resource.

    pub range: std::ops::Range<u64>,

    /// Acquire barrier.

    /// Node implementation must insert it before first command that uses the buffer.

    /// Barrier must be inserted even if this node doesn't use the buffer.

    pub acquire: Option<BufferBarrier>,

    /// Release barrier.

    /// Node implementation must insert it after last command that uses the buffer.

    /// Barrier must be inserted even if this node doesn't use the buffer.

    pub release: Option<BufferBarrier>,
}

/// Image access node wants to perform.

#[derive(Clone, Copy, Debug)]
pub struct ImageAccess {
    /// Access flags.

    pub access: rendy_core::hal::image::Access,

    /// Intended usage flags for image.

    /// TODO: Could derive from access?

    pub usage: rendy_core::hal::image::Usage,

    /// Preferred layout for access.

    /// Actual layout will be reported int `NodeImage`.

    /// Actual layout is guaranteed to support same operations.

    /// TODO: Could derive from access?

    pub layout: rendy_core::hal::image::Layout,

    /// Pipeline stages at which image is accessd.

    pub stages: rendy_core::hal::pso::PipelineStage,
}

/// Image pipeline barrier.

/// Node implementation must insert it before first command that uses the image.

/// Barrier must be inserted even if this node doesn't use the image.

#[derive(Clone, Debug)]
pub struct ImageBarrier {
    /// State transition for the image.

    pub states: std::ops::Range<rendy_core::hal::image::State>,

    /// Stages at which image is accessd.

    pub stages: std::ops::Range<rendy_core::hal::pso::PipelineStage>,

    /// Transfer between families.

    pub families: Option<std::ops::Range<QueueFamilyId>>,
}

/// Image shared between nodes.

#[derive(Clone, Debug)]
pub struct NodeImage {
    /// Id of the image.

    pub id: ImageId,

    /// Region of the image that is the transient resource.

    pub range: rendy_core::hal::image::SubresourceRange,

    /// Image state for node.

    pub layout: rendy_core::hal::image::Layout,

    /// Specify that node should clear image to this value.

    pub clear: Option<rendy_core::hal::command::ClearValue>,

    /// Acquire barrier.

    /// Node implementation must insert it before first command that uses the image.

    /// Barrier must be inserted even if this node doesn't use the image.

    pub acquire: Option<ImageBarrier>,

    /// Release barrier.

    /// Node implementation must insert it after last command that uses the image.

    /// Barrier must be inserted even if this node doesn't use the image.

    pub release: Option<ImageBarrier>,
}

/// NodeSubmittable

pub trait NodeSubmittable<'a, B: Backend> {
    /// Submittable type returned from `Node`.

    type Submittable: Submittable<B> + 'a;

    /// Iterator over submittables returned from `Node`.

    type Submittables: IntoIterator<Item = Self::Submittable>;
}

/// The node is building block of the framegraph.

/// Node defines set of resources and operations to perform over them.

/// Read-only data for operations comes from auxiliary data source `T`.

///

/// # Parameters

///

/// `B` - backend type.

/// `T` - auxiliary data type.

///

pub trait Node<B: Backend, T: ?Sized>:
    for<'a> NodeSubmittable<'a, B> + std::fmt::Debug + Sized + Sync + Send + 'static
{
    /// Capability required by node.

    /// Graph will execute this node on command queue that supports this capability level.

    type Capability: Capability;

    /// Record commands required by node.

    /// Returned submits are guaranteed to be submitted within specified frame.

    fn run<'a>(
        &'a mut self,
        ctx: &GraphContext<B>,
        factory: &Factory<B>,
        aux: &T,
        frames: &'a Frames<B>,
    ) -> <Self as NodeSubmittable<'a, B>>::Submittables;

    /// Dispose of the node.

    ///

    /// # Safety

    ///

    /// Must be called after waiting for device idle.

    unsafe fn dispose(self, factory: &mut Factory<B>, aux: &T);
}

/// Description of the node.

/// Implementation of the builder type provide framegraph with static information about node

/// that is used for building the node.

pub trait NodeDesc<B: Backend, T: ?Sized>: std::fmt::Debug + Sized + 'static {
    /// Node this builder builds.

    type Node: Node<B, T>;

    /// Make node builder.

    fn builder(self) -> DescBuilder<B, T, Self> {
        DescBuilder::new(self)
    }

    /// Get set or buffer resources the node uses.

    fn buffers(&self) -> Vec<BufferAccess> {
        Vec::new()
    }

    /// Get set or image resources the node uses.

    fn images(&self) -> Vec<ImageAccess> {
        Vec::new()
    }

    /// Build the node.

    ///

    /// # Parameters

    ///

    /// `factory`    - factory instance.

    /// `aux`       - auxiliary data.

    /// `family`    - id of the family this node will be executed on.

    /// `resources` - set of transient resources managed by graph.

    ///               with barriers required for interface resources.

    ///

    fn build<'a>(
        self,
        ctx: &GraphContext<B>,
        factory: &mut Factory<B>,
        family: &mut Family<B>,
        queue: usize,
        aux: &T,
        buffers: Vec<NodeBuffer>,
        images: Vec<NodeImage>,
    ) -> Result<Self::Node, NodeBuildError>;
}

/// Trait-object safe `Node`.

pub trait DynNode<B: Backend, T: ?Sized>: std::fmt::Debug + Sync + Send {
    /// Record commands required by node.

    /// Recorded buffers go into `submits`.

    unsafe fn run<'a>(
        &mut self,
        ctx: &GraphContext<B>,
        factory: &Factory<B>,
        queue: &mut Queue<B>,
        aux: &T,
        frames: &Frames<B>,
        waits: &[(&'a B::Semaphore, rendy_core::hal::pso::PipelineStage)],
        signals: &[&'a B::Semaphore],
        fence: Option<&mut Fence<B>>,
    );

    /// Dispose of the node.

    ///

    /// # Safety

    ///

    /// Must be called after waiting for device idle.

    unsafe fn dispose(self: Box<Self>, factory: &mut Factory<B>, aux: &T);
}

impl<B, T, N> DynNode<B, T> for (N,)
where
    B: Backend,
    T: ?Sized,
    N: Node<B, T>,
{
    unsafe fn run<'a>(
        &mut self,
        ctx: &GraphContext<B>,
        factory: &Factory<B>,
        queue: &mut Queue<B>,
        aux: &T,
        frames: &Frames<B>,
        waits: &[(&'a B::Semaphore, rendy_core::hal::pso::PipelineStage)],
        signals: &[&'a B::Semaphore],
        fence: Option<&mut Fence<B>>,
    ) {
        let submittables = Node::run(&mut self.0, ctx, factory, aux, frames);
        queue.submit(
            Some(
                Submission::new()
                    .submits(submittables)
                    .wait(waits.iter().cloned())
                    .signal(signals.iter().cloned()),
            ),
            fence,
        )
    }

    unsafe fn dispose(self: Box<Self>, factory: &mut Factory<B>, aux: &T) {
        N::dispose(self.0, factory, aux);
    }
}

/// Error building a node of the graph.

#[derive(Debug)]
pub enum NodeBuildError {
    /// Filed to uplaod the data.

    Upload(UploadError),
    /// Mismatched queue family.

    QueueFamily(FamilyId),
    /// Failed to create an imate view.

    View(rendy_core::hal::image::ViewError),
    /// Failed to create a pipeline.

    Pipeline(rendy_core::hal::pso::CreationError),
    /// Failed to create a swap chain.

    Swapchain(SwapchainError),
    /// Ran out of memory when creating something.

    OutOfMemory(rendy_core::hal::device::OutOfMemory),
}

/// Dynamic node builder that emits `DynNode`.

pub trait NodeBuilder<B: Backend, T: ?Sized>: std::fmt::Debug {
    /// Pick family for this node to be executed onto.

    fn family(&self, factory: &mut Factory<B>, families: &Families<B>) -> Option<FamilyId>;

    /// Get buffer accessed by the node.

    fn buffers(&self) -> Vec<(BufferId, BufferAccess)>;

    /// Get images accessed by the node.

    fn images(&self) -> Vec<(ImageId, ImageAccess)>;

    /// Indices of nodes this one dependes on.

    fn dependencies(&self) -> Vec<NodeId>;

    /// Build node.

    fn build<'a>(
        self: Box<Self>,
        ctx: &GraphContext<B>,
        factory: &mut Factory<B>,
        family: &mut Family<B>,
        queue: usize,
        aux: &T,
        buffers: Vec<NodeBuffer>,
        images: Vec<NodeImage>,
    ) -> Result<Box<dyn DynNode<B, T>>, NodeBuildError>;
}

/// Builder for the node.

pub struct DescBuilder<B: Backend, T: ?Sized, N> {
    desc: N,
    buffers: Vec<BufferId>,
    images: Vec<ImageId>,
    dependencies: Vec<NodeId>,
    marker: std::marker::PhantomData<fn(B, &T)>,
}

impl<B, T, N> std::fmt::Debug for DescBuilder<B, T, N>
where
    B: Backend,
    T: ?Sized,
    N: std::fmt::Debug,
{
    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        fmt.debug_struct("DescBuilder")
            .field("desc", &self.desc)
            .field("buffers", &self.buffers)
            .field("images", &self.images)
            .field("dependencies", &self.dependencies)
            .finish()
    }
}

impl<B, T, N> DescBuilder<B, T, N>
where
    B: Backend,
    T: ?Sized,
{
    /// Create new builder out of desc

    pub fn new(desc: N) -> Self {
        DescBuilder {
            desc,
            buffers: Vec::new(),
            images: Vec::new(),
            dependencies: Vec::new(),
            marker: std::marker::PhantomData,
        }
    }
    /// Add buffer to the node.

    /// This method must be called for each buffer node uses.

    pub fn add_buffer(&mut self, buffer: BufferId) -> &mut Self {
        self.buffers.push(buffer);
        self
    }

    /// Add buffer to the node.

    /// This method must be called for each buffer node uses.

    pub fn with_buffer(mut self, buffer: BufferId) -> Self {
        self.add_buffer(buffer);
        self
    }

    /// Add image to the node.

    /// This method must be called for each image node uses.

    pub fn add_image(&mut self, image: ImageId) -> &mut Self {
        self.images.push(image);
        self
    }

    /// Add image to the node.

    /// This method must be called for each image node uses.

    pub fn with_image(mut self, image: ImageId) -> Self {
        self.add_image(image);
        self
    }

    /// Add dependency.

    /// Node will be placed after its dependencies.

    pub fn add_dependency(&mut self, dependency: NodeId) -> &mut Self {
        self.dependencies.push(dependency);
        self
    }

    /// Add dependency.

    /// Node will be placed after its dependencies.

    pub fn with_dependency(mut self, dependency: NodeId) -> Self {
        self.add_dependency(dependency);
        self
    }
}

impl<B, T, N> NodeBuilder<B, T> for DescBuilder<B, T, N>
where
    B: Backend,
    T: ?Sized,
    N: NodeDesc<B, T>,
{
    fn family(&self, _factory: &mut Factory<B>, families: &Families<B>) -> Option<FamilyId> {
        families.with_capability::<<N::Node as Node<B, T>>::Capability>()
    }

    fn buffers(&self) -> Vec<(BufferId, BufferAccess)> {
        let desc_buffers = self.desc.buffers();
        assert_eq!(self.buffers.len(), desc_buffers.len());

        self.buffers.iter().cloned().zip(desc_buffers).collect()
    }

    fn images(&self) -> Vec<(ImageId, ImageAccess)> {
        let desc_images = self.desc.images();
        assert_eq!(self.images.len(), desc_images.len());

        self.images.iter().cloned().zip(desc_images).collect()
    }

    fn dependencies(&self) -> Vec<NodeId> {
        self.dependencies.clone()
    }

    fn build<'a>(
        self: Box<Self>,
        ctx: &GraphContext<B>,
        factory: &mut Factory<B>,
        family: &mut Family<B>,
        queue: usize,
        aux: &T,
        buffers: Vec<NodeBuffer>,
        images: Vec<NodeImage>,
    ) -> Result<Box<dyn DynNode<B, T>>, NodeBuildError> {
        Ok(Box::new((self.desc.build(
            ctx, factory, family, queue, aux, buffers, images,
        )?,)))
    }
}

/// Convert graph barriers into gfx barriers.

pub fn gfx_acquire_barriers<'a, 'b, B: Backend>(
    ctx: &'a GraphContext<B>,
    buffers: impl IntoIterator<Item = &'b NodeBuffer>,
    images: impl IntoIterator<Item = &'b NodeImage>,
) -> (
    std::ops::Range<rendy_core::hal::pso::PipelineStage>,
    Vec<rendy_core::hal::memory::Barrier<'a, B>>,
) {
    let mut bstart = rendy_core::hal::pso::PipelineStage::empty();
    let mut bend = rendy_core::hal::pso::PipelineStage::empty();

    let mut istart = rendy_core::hal::pso::PipelineStage::empty();
    let mut iend = rendy_core::hal::pso::PipelineStage::empty();

    let barriers: Vec<rendy_core::hal::memory::Barrier<'_, B>> = buffers
        .into_iter()
        .filter_map(|buffer| {
            buffer.acquire.as_ref().map(|acquire| {
                bstart |= acquire.stages.start;
                bend |= acquire.stages.end;

                rendy_core::hal::memory::Barrier::Buffer {
                    states: acquire.states.clone(),
                    families: acquire.families.clone(),
                    target: ctx
                        .get_buffer(buffer.id)
                        .expect("Buffer does not exist")
                        .raw(),
                    range: Some(buffer.range.start)..Some(buffer.range.end),
                }
            })
        })
        .chain(images.into_iter().filter_map(|image| {
            image.acquire.as_ref().map(|acquire| {
                istart |= acquire.stages.start;
                iend |= acquire.stages.end;

                rendy_core::hal::memory::Barrier::Image {
                    states: acquire.states.clone(),
                    families: acquire.families.clone(),
                    target: ctx.get_image(image.id).expect("Image does not exist").raw(),
                    range: image.range.clone(),
                }
            })
        }))
        .collect();

    (bstart | istart..bend | iend, barriers)
}

/// Convert graph barriers into gfx barriers.

pub fn gfx_release_barriers<'a, B: Backend>(
    ctx: &'a GraphContext<B>,
    buffers: impl IntoIterator<Item = &'a NodeBuffer>,
    images: impl IntoIterator<Item = &'a NodeImage>,
) -> (
    std::ops::Range<rendy_core::hal::pso::PipelineStage>,
    Vec<rendy_core::hal::memory::Barrier<'a, B>>,
) {
    let mut bstart = rendy_core::hal::pso::PipelineStage::empty();
    let mut bend = rendy_core::hal::pso::PipelineStage::empty();

    let mut istart = rendy_core::hal::pso::PipelineStage::empty();
    let mut iend = rendy_core::hal::pso::PipelineStage::empty();

    let barriers: Vec<rendy_core::hal::memory::Barrier<'_, B>> = buffers
        .into_iter()
        .filter_map(|buffer| {
            buffer.release.as_ref().map(|release| {
                bstart |= release.stages.start;
                bend |= release.stages.end;

                rendy_core::hal::memory::Barrier::Buffer {
                    states: release.states.clone(),
                    families: release.families.clone(),
                    target: ctx
                        .get_buffer(buffer.id)
                        .expect("Buffer does not exist")
                        .raw(),
                    range: Some(buffer.range.start)..Some(buffer.range.end),
                }
            })
        })
        .chain(images.into_iter().filter_map(|image| {
            image.release.as_ref().map(|release| {
                istart |= release.stages.start;
                iend |= release.stages.end;

                rendy_core::hal::memory::Barrier::Image {
                    states: release.states.clone(),
                    families: release.families.clone(),
                    target: ctx.get_image(image.id).expect("Image does not exist").raw(),
                    range: image.range.clone(),
                }
            })
        }))
        .collect();

    (bstart | istart..bend | iend, barriers)
}