//! Simple frame graph implementation

use std::mem::{discriminant, Discriminant};
use std::ops::Range;
use std::vec::Drain;

use gfx_hal::command::CommandBuffer;
use gfx_hal::pool::CommandPoolCreateFlags;
use gfx_hal::pso::PipelineStage;
use gfx_hal::{Backend, Device, Graphics, Submission};

use bitvec::BitVec;
use failure::Fallible;
use fnv::FnvHashMap;
use smallvec::SmallVec;

use crate::encoder::Encoder;
use crate::factory::Factory;
use crate::node::{AnyNode, AnyNodePlan, Node, NodeId, NodePlan};

type NodeDeps = SmallVec<[NodeId; 8]>;

#[derive(derivative::Derivative)]
#[derivative(Default(bound = ""))]
/// Dynamic group of nodes of the same type
struct NodeGroup<B: Backend> {
    /// The range of nodes we can rebuild during node group updating
    #[derivative(Default(value = "0..0"))]
    visit_range: Range<usize>,

    nodes: Vec<AnyNode<B>>,
    plans: Vec<AnyNodePlan<B>>,

    // Every node is associated with a single semaphore, which will be signaled after node execution
    wait_semaphores: Vec<NodeDeps>,
    signal_semaphores: Vec<NodeId>,
}

impl<B: Backend> NodeGroup<B> {
    fn new() -> NodeGroup<B> {
        Default::default()
    }

    /// Called every frame to avoid memory leaks and reuse nodes instead of rebuilding them
    fn reset(&mut self, f: &mut Factory<B>) {
        if self.visit_range.start != self.visit_range.end {
            for node in self.nodes.drain(self.visit_range.start..) {
                node.destroy(f);
            }

            self.plans.truncate(self.visit_range.start);
        }

        self.visit_range = 0..self.plans.len()
    }

    /// Get next node id to rebuild
    fn next_plan(&mut self) -> Option<usize> {
        // TODO: choose the best node to rebuild, instead of the first one
        if self.visit_range.start < self.visit_range.end {
            self.visit_range.start += 1;
            Some(self.visit_range.start - 1)
        } else {
            self.visit_range.start += 1;
            None
        }
    }

    /// Add a new node to the group, rebuilding the existing one or instantiating a new one.
    fn add_node(&mut self, f: &mut Factory<B>, id: NodeId, deps: NodeDeps, plan: AnyNodePlan<B>) {
        if let Some(idx) = self.next_plan() {
            self.rebuild_node(f, idx, id, deps, plan);
        } else {
            self.push_node(f, id, deps, plan);
        };
    }

    fn rebuild_node(
        &mut self,
        f: &mut Factory<B>,
        idx: usize,
        id: NodeId,
        deps: NodeDeps,
        plan: AnyNodePlan<B>,
    ) {
        plan.rebuild(f, &mut self.nodes[idx]);
        self.wait_semaphores[idx] = deps;
        self.signal_semaphores[idx] = id;
    }

    fn push_node(&mut self, f: &mut Factory<B>, id: NodeId, deps: NodeDeps, plan: AnyNodePlan<B>) {
        self.nodes.push(plan.build(f));
        self.plans.push(plan);
        self.wait_semaphores.push(deps);
        self.signal_semaphores.push(id);
    }
}

/// Graph containing a "recipe" with dependencies for rendering the next frame
pub struct FrameGraph<B: Backend> {
    nodes: FnvHashMap<Discriminant<AnyNodePlan<B>>, NodeGroup<B>>,
    num_semaphores: usize,
    factory: Factory<B>,
    final_semaphores: Vec<NodeId>,
}

impl<B: Backend> FrameGraph<B> {
    /// Create a new frame graph encapsulating the specified `Factory`.
    pub fn new(factory: Factory<B>) -> FrameGraph<B> {
        FrameGraph {
            nodes: Default::default(),
            num_semaphores: 0,
            factory,
            final_semaphores: vec![],
        }
    }

    /// Create a new frame graph plan. Does the same thing as `FrameGraphPlan::new` does.
    pub fn plan() -> FrameGraphPlan<B> {
        Default::default()
    }

    fn update(
        &mut self,
        deps: Drain<NodeDeps>,
        plans: Drain<AnyNodePlan<B>>,
        final_nodes: &BitVec,
    ) {
        self.num_semaphores = plans.len();

        for (id, (deps, plan)) in deps.zip(plans).enumerate() {
            if final_nodes[id] {
                self.final_semaphores.push(NodeId(id as _));
            }

            self.nodes
                .entry(discriminant(&plan))
                .or_insert(NodeGroup::new())
                .add_node(&mut self.factory, NodeId(id as _), deps, plan);
        }

        for group in self.nodes.values_mut() {
            group.reset(&mut self.factory);
        }
    }

    /// Execute the frame graph.
    ///
    /// Records commands from every node and submit command buffers to the queue
    /// (currently only a single queue is used), putting semaphores where needed.
    ///
    /// The order of submitted command buffers is not defined, however this shouldn't matter
    /// when the nodes have correct dependencies.
    pub fn run(&mut self) -> Fallible<()> {
        let mut semaphores = Vec::with_capacity(self.num_semaphores);

        for _ in 0..self.num_semaphores {
            semaphores.push(self.factory.device.create_semaphore()?);
        }

        let mut cpool = unsafe {
            self.factory.device.create_command_pool_typed(
                &self.factory.queue_group,
                CommandPoolCreateFlags::empty(),
            )?
        };

        for group in self.nodes.values_mut() {
            for (id, node) in group.nodes.iter_mut().enumerate() {
                let mut cbuf = cpool.acquire_command_buffer();
                let encoder = Encoder { buf: &mut cbuf };

                node.run(&mut self.factory, encoder);

                let wait = group.wait_semaphores[id]
                    .iter()
                    .map(|&NodeId(v)| (&semaphores[v as usize], PipelineStage::BOTTOM_OF_PIPE));

                let signal = &semaphores[group.signal_semaphores[id].0 as usize];

                let submission = Submission {
                    command_buffers: &[cbuf],
                    wait_semaphores: wait,
                    signal_semaphores: Some(signal),
                };

                let queue = &mut self.factory.queue_group.queues[0];
                unsafe { queue.submit(submission, None) };
            }
        }

        let wait = self
            .final_semaphores
            .iter()
            .map(|&NodeId(v)| (&semaphores[v as usize], PipelineStage::BOTTOM_OF_PIPE));

        let submission = Submission {
            command_buffers: Option::<&CommandBuffer<B, Graphics>>::None,
            wait_semaphores: wait,
            signal_semaphores: Option::<&B::Semaphore>::None,
        };

        let fence = self.factory.device.create_fence(false)?;

        let queue = &mut self.factory.queue_group.queues[0];
        unsafe { queue.submit(submission, Some(&fence)) };

        unsafe { self.factory.device.wait_for_fence(&fence, std::u64::MAX) }?;

        unsafe { self.factory.device.destroy_fence(fence) };

        for semaphore in semaphores.drain(..) {
            unsafe { self.factory.device.destroy_semaphore(semaphore) };
        }

        unsafe { self.factory.device.destroy_command_pool(cpool.into_raw()) };

        Ok(())
    }
}

impl<B: Backend> Drop for FrameGraph<B> {
    fn drop(&mut self) {
        for (_, mut group) in self.nodes.drain() {
            for node in group.nodes.drain(..) {
                node.destroy(&mut self.factory);
            }
        }
    }
}

/// Plan for building or making changes to the frame graph.
#[derive(Clone, derivative::Derivative)]
#[derivative(Default(bound = ""))]
pub struct FrameGraphPlan<B: Backend> {
    plans: Vec<AnyNodePlan<B>>,
    deps: Vec<NodeDeps>,
    final_nodes: BitVec,
}

impl<B: Backend> FrameGraphPlan<B> {
    /// Create a new frame graph plan. Does the same thing as `FrameGraph::plan` does.
    pub fn new() -> FrameGraphPlan<B> {
        Default::default()
    }

    /// Add a new node to the plan with specified dependencies.
    ///
    /// Returns ID of the newly created node.
    pub fn add_node<P: Into<AnyNodePlan<B>>>(&mut self, deps: &[NodeId], plan: P) -> NodeId {
        for dep in deps {
            self.final_nodes.set(dep.0 as _, false);
        }

        self.final_nodes.push(true);
        self.plans.push(plan.into());
        self.deps.push(deps.into());
        NodeId((self.plans.len() - 1) as _)
    }

    /// Commit changes to the framegraph allocating or just updating resources for planned nodes.
    ///
    /// After that, plan is cleared and can (and should) be used again next frame to reduce
    /// memory allocation count. If you want to commit the same plan twice, you can `.clone()` it.
    pub fn commit(&mut self, fg: &mut FrameGraph<B>) {
        fg.update(self.deps.drain(..), self.plans.drain(..), &self.final_nodes);
        self.final_nodes.clear();
    }
}