//! These graph shapes compare command-level batching with render-pass span
//! contraction and barrier hoisting. Each has two disjoint buffer access sets
//! joined by a shared-buffer dependency across the passes.

use std::{hint::black_box, time::Duration};

use ash::vk;
use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
use renderer_rendergraph::{
    Buffer, BufferRange, GraphBufferDesc, RenderGraph,
    commands::{BeginRendering, Command, EndRendering},
    recorder::{CommandIndex, CommandRecorder, RecordError, RecordedCommand, SideEffectMap},
    resources::{SubresourceScope, SubresourceScopeInner},
};

struct Use {
    buffer: Buffer,
    write: bool,
}

impl Command for Use {
    fn side_effects(&self, command_index: CommandIndex, map: &mut SideEffectMap) {
        map.insert(
            command_index,
            self.buffer.id(),
            SubresourceScope {
                stage: vk::PipelineStageFlags2::COMPUTE_SHADER,
                access: if self.write {
                    vk::AccessFlags2::SHADER_WRITE
                } else {
                    vk::AccessFlags2::SHADER_READ
                },
                subresource: SubresourceScopeInner::Buffer(BufferRange::full()),
            },
        );
    }
}

impl RecordedCommand for Use {
    fn allowed_in_render_pass(&self) -> bool {
        true
    }

    fn apply(&mut self, _recorder: &mut CommandRecorder<'_>) -> Result<(), RecordError> {
        Ok(())
    }
}

fn build(n_groups: usize, per: usize, pass: bool) -> RenderGraph {
    let mut graph = RenderGraph::new();
    let a: Vec<Buffer> = (0..n_groups)
        .map(|_| {
            graph.create_buffer(GraphBufferDesc {
                size: 4096,
                ..Default::default()
            })
        })
        .collect();
    let b: Vec<Buffer> = (0..n_groups)
        .map(|_| {
            graph.create_buffer(GraphBufferDesc {
                size: 4096,
                ..Default::default()
            })
        })
        .collect();
    let shared = graph.create_buffer(GraphBufferDesc {
        size: 4096,
        ..Default::default()
    });

    if pass {
        graph.push(BeginRendering {
            color: vec![],
            depth: None,
        });
    }
    for _ in 0..per {
        for &buffer in &a {
            graph.push(Use {
                buffer,
                write: true,
            });
        }
    }
    graph.push(Use {
        buffer: shared,
        write: true,
    });
    if pass {
        graph.push(EndRendering);
        graph.push(BeginRendering {
            color: vec![],
            depth: None,
        });
    }
    for _ in 0..per {
        for &buffer in &b {
            graph.push(Use {
                buffer,
                write: true,
            });
        }
    }
    graph.push(Use {
        buffer: shared,
        write: false,
    });
    if pass {
        graph.push(EndRendering);
    }

    graph
}

fn compile_benchmarks(criterion: &mut Criterion) {
    let mut group = criterion.benchmark_group("compile");
    group
        .measurement_time(Duration::from_secs(4))
        .sample_size(20);

    for (pass_name, pass) in [("no_pass", false), ("pass", true)] {
        for per in [20, 50, 100] {
            group.bench_function(format!("two_pass/{pass_name}/per_{per}"), |bencher| {
                bencher.iter_batched(
                    || build(50, per, pass),
                    |graph| {
                        let compiled = graph.compile().unwrap();
                        black_box(compiled.barriers().len())
                    },
                    BatchSize::LargeInput,
                );
            });
        }
    }

    group.finish();
}

criterion_group!(benches, compile_benchmarks);
criterion_main!(benches);