use ash::vk;

use crate::{
    CommandIndex, GraphError, SideEffectMap,
    graph::{Kahn, RenderPassSpan},
    recorder::BoxedCommand,
    resources::{ResourceDesc, SubresourceScope, SubresourceScopeInner},
};

struct SpanBuilder {
    spans: Vec<RenderPassSpan>,
    current: Option<CommandIndex>,
}

impl SpanBuilder {
    fn new() -> Self {
        Self {
            spans: Vec::new(),
            current: None,
        }
    }

    fn in_span(&self) -> bool {
        self.current.is_some()
    }

    fn current_begin(&self) -> Option<CommandIndex> {
        self.current
    }

    fn begin(&mut self, index: CommandIndex) -> Result<(), GraphError> {
        if self.current.is_some() {
            Err(GraphError::NestedRenderPass)
        } else {
            self.current = Some(index);
            Ok(())
        }
    }

    fn end(&mut self, end: CommandIndex) -> Result<(), GraphError> {
        if let Some(begin) = self.current.take() {
            self.spans.push(RenderPassSpan { begin, end });
            Ok(())
        } else {
            Err(GraphError::UnbalancedRenderPass)
        }
    }

    fn finish(self) -> Result<Vec<RenderPassSpan>, GraphError> {
        if self.current.is_some() {
            Err(GraphError::UnbalancedRenderPass)
        } else {
            Ok(self.spans)
        }
    }
}

pub(crate) struct Batch {
    pub(crate) start: usize,
    pub(crate) count: usize,
}

pub(crate) struct CommandOrderings {
    pub(crate) execution_order: Vec<CommandIndex>,
    pub(crate) batches: Vec<Batch>,
    pub(crate) spans: Vec<RenderPassSpan>,
    pub(crate) node_of_command: Vec<usize>,
    pub(crate) node_representatives: Vec<CommandIndex>,
    pub(crate) dependency_edges: Vec<Vec<usize>>,
}

pub(crate) fn build_command_orderings(
    resources: &[ResourceDesc],
    commands: &[BoxedCommand],
    side_effects: &mut SideEffectMap,
) -> Result<CommandOrderings, GraphError> {
    let per_command_resource_access = side_effects.command_accesses();
    let mut kahn = Kahn::new(commands.len());
    let mut access_by_resource: Vec<Vec<(CommandIndex, Vec<SubresourceScope>)>> =
        vec![Vec::new(); resources.len()];

    let mut span_builder = SpanBuilder::new();
    let mut node_of_command = vec![usize::MAX; commands.len()];
    let mut node_representatives = Vec::new();

    for (index, (command, accesses)) in commands
        .iter()
        .zip(per_command_resource_access)
        .enumerate()
        .map(|(i, c)| (CommandIndex::from(i), c))
    {
        let rep = span_builder.current_begin().unwrap_or(index);
        kahn.add_node(rep.index());

        // Node ids follow first encounter of each representative in command order.
        let node = if node_of_command[rep.index()] == usize::MAX {
            let node = node_representatives.len();
            node_representatives.push(rep);
            node
        } else {
            node_of_command[rep.index()]
        };
        node_of_command[index.index()] = node;

        for access in accesses {
            let Some(resource_access) = access_by_resource.get_mut(access.resource.index()) else {
                // `validate_resources` reports invalid ids after dependency analysis.
                continue;
            };

            let priors = resource_access
                .iter()
                .take_while(|&&(c, _)| c < rep)
                .filter(|(_, scopes)| scopes.iter().any(|scope| scope.conflicts(&access.scope)));

            for &(prior, _) in priors {
                kahn.add_edge(prior.index(), rep.index());
            }

            if let Some(&mut (last, ref mut scopes)) = resource_access.last_mut()
                && last == rep
            {
                scopes.push(access.scope);
            } else {
                resource_access.push((rep, vec![access.scope]))
            }
        }

        match command.kind() {
            crate::CommandKind::Begin => {
                if let Some(desc) = command.render_pass_desc() {
                    crate::graph::validate_render_pass_attachments(resources, desc)?;
                }
                span_builder.begin(index)?;
            }
            crate::CommandKind::Graphics => {
                if !command.allowed_in_render_pass() {
                    return Err(GraphError::CommandNotAllowedInRenderPass(index));
                }
                if !span_builder.in_span() {
                    return Err(GraphError::DrawOutsideRenderPass(index));
                }
            }
            crate::CommandKind::Transfer
            | crate::CommandKind::Compute
            | crate::CommandKind::Other => {
                if span_builder.in_span() && !command.allowed_in_render_pass() {
                    return Err(GraphError::CommandNotAllowedInRenderPass(index));
                }
            }
            crate::CommandKind::End => {
                span_builder.end(index)?;
            }
        }
    }

    // spans are ordered by beginning command
    let spans = span_builder.finish()?;

    for span in &spans {
        // Snapshot accesses intentionally predate lifetime-only attachment injection:
        // these accesses extend lifetimes but must not add hazard edges.
        if let Some(desc) = commands[span.begin.index()].render_pass_desc() {
            for attachment in desc.color {
                side_effects.insert_lifetime_only(
                    span.end,
                    attachment.texture.id(),
                    SubresourceScope {
                        stage: vk::PipelineStageFlags2::COLOR_ATTACHMENT_OUTPUT,
                        access: vk::AccessFlags2::COLOR_ATTACHMENT_WRITE,
                        subresource: SubresourceScopeInner::Texture {
                            range: attachment.range,
                            layout: attachment.layout,
                        },
                    },
                );
            }

            if let Some(attachment) = desc.depth {
                side_effects.insert_lifetime_only(
                    span.end,
                    attachment.texture.id(),
                    SubresourceScope {
                        stage: vk::PipelineStageFlags2::EARLY_FRAGMENT_TESTS
                            | vk::PipelineStageFlags2::LATE_FRAGMENT_TESTS,
                        access: vk::AccessFlags2::DEPTH_STENCIL_ATTACHMENT_WRITE,
                        subresource: SubresourceScopeInner::Texture {
                            range: attachment.range,
                            layout: attachment.layout,
                        },
                    },
                );
            }
        }
    }

    let mut execution_order: Vec<CommandIndex> = Vec::with_capacity(commands.len());
    let mut batches = Vec::<Batch>::new();

    let mut span_iter = spans.iter().peekable();

    while !kahn.is_complete() {
        let first = execution_order.len();

        for cmd in kahn.iter_batch().map(CommandIndex::from) {
            use core::cmp::Ordering;
            match span_iter.peek().map(|s| s.begin.cmp(&cmd)) {
                Some(Ordering::Equal) => {
                    let span = span_iter.next().unwrap();
                    execution_order
                        .extend((span.begin.index()..=span.end.index()).map(CommandIndex::from));
                }
                _ => execution_order.push(cmd),
            }
        }

        let count = execution_order.len() - first;

        if count == 0 {
            return Err(GraphError::DependencyCycle);
        }

        for command in &execution_order[first..] {
            kahn.complete(command.index());
        }

        batches.push(Batch {
            start: first,
            count,
        });
    }

    let mut dependency_edges = vec![Vec::new(); node_representatives.len()];
    for (node, representative) in node_representatives.iter().enumerate() {
        let representative = representative.index();
        let node_edges = &mut dependency_edges[node];
        for &edge in &kahn.edges[representative] {
            node_edges.push(node_of_command[edge]);
        }
    }

    Ok(CommandOrderings {
        execution_order,
        batches,
        spans,
        node_of_command,
        node_representatives,
        dependency_edges,
    })
}