//! 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. The resolve benchmarks
//! include device-dependent placement, alias-barrier rebuilding, and pass hoisting
//! in the timed region.

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

use ash::vk;
use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
use renderer_device::device::{MemoryLocation, QueueFlags};
use renderer_device::testing::TestVulkan;
use renderer_rendergraph::{
    Buffer, BufferRange, GraphBufferDesc, RenderGraph,
    commands::{BeginRendering, Command, EndRendering},
    recorder::{CommandIndex, CommandRecorder, RecordError, RecordedCommand, SideEffectMap},
    resources::{ResourceState, SubresourceScope, SubresourceScopeInner},
};
use renderer_resources::{buffers::Buffer as DeviceBuffer, buffers::BufferDesc};

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(())
    }
}

struct UseWithBridge {
    buffer: Buffer,
    bridge: Buffer,
}

impl Command for UseWithBridge {
    fn side_effects(&self, command_index: CommandIndex, map: &mut SideEffectMap) {
        for buffer in [self.buffer, self.bridge] {
            map.insert(
                command_index,
                buffer.id(),
                SubresourceScope {
                    stage: vk::PipelineStageFlags2::COMPUTE_SHADER,
                    access: vk::AccessFlags2::SHADER_WRITE,
                    subresource: SubresourceScopeInner::Buffer(BufferRange::full()),
                },
            );
        }
    }
}

impl RecordedCommand for UseWithBridge {
    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 build_overlapping(n_groups: usize, per: usize) -> RenderGraph {
    let mut graph = RenderGraph::new();
    let buffers: 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()
    });

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

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

    graph
}

fn build_aliasing(
    n_groups: usize,
    per: usize,
    bridge: &Arc<DeviceBuffer>,
) -> 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 bridge = graph
        .import_buffer(bridge.clone(), ResourceState::default())
        .unwrap();

    // Keep the first set outside the pass; render-pass resources interfere with
    // each other by design and otherwise prevent placement reuse.
    for _ in 0..per {
        for &buffer in &a {
            graph.push(Use {
                buffer,
                write: true,
            });
        }
    }
    for &buffer in &a {
        graph.push(UseWithBridge { buffer, bridge });
    }

    graph.push(BeginRendering {
        color: vec![],
        depth: None,
    });
    for &buffer in &b {
        graph.push(UseWithBridge { buffer, bridge });
    }
    for _ in 0..per {
        for &buffer in &b {
            graph.push(Use {
                buffer,
                write: true,
            });
        }
    }
    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();
}

fn resolve_benchmarks(criterion: &mut Criterion, device: &renderer_device::Device) {
    let bridge = Arc::new(
        DeviceBuffer::new(
            device.clone(),
            BufferDesc {
                size: 4096,
                usage: vk::BufferUsageFlags::STORAGE_BUFFER,
                queue_families: QueueFlags::empty(),
                mem_location: MemoryLocation::GpuOnly,
                ..Default::default()
            },
        )
        .expect("create alias benchmark bridge buffer"),
    );

    let mut group = criterion.benchmark_group("compile_resolve");
    group
        .measurement_time(Duration::from_secs(4))
        .sample_size(20);

    for (name, overlapping) in [("alias", false), ("no_alias", true)] {
        for per in [20, 50, 100] {
            let graph = if overlapping {
                build_overlapping(50, per)
            } else {
                build_aliasing(50, per, &bridge)
            };
            let mut probe = graph.compile().unwrap();
            probe.resolve(device).unwrap();
            assert_eq!(
                probe
                    .resource_lifetimes()
                    .values()
                    .any(|lifetime| lifetime.placement.aliased),
                !overlapping,
                "unexpected placement aliasing in {name}/per_{per}"
            );
            let alias_barriers = probe
                .barriers()
                .iter()
                .flat_map(|barrier| &barrier.resources)
                .filter(|resource| {
                    probe
                        .resource_lifetimes()
                        .get(&resource.id)
                        .is_some_and(|lifetime| lifetime.placement.aliased)
                        && resource.origin != renderer_rendergraph::BarrierOrigin::Initial
                })
                .count();
            assert_eq!(
                alias_barriers > 0,
                !overlapping,
                "unexpected alias barriers in {name}/per_{per}"
            );

            group.bench_function(format!("{name}/per_{per}"), |bencher| {
                bencher.iter_batched(
                    || {
                        if overlapping {
                            build_overlapping(50, per)
                        } else {
                            build_aliasing(50, per, &bridge)
                        }
                    },
                    |graph| {
                        let mut compiled = graph.compile().unwrap();
                        compiled.resolve(device).unwrap();
                        black_box(compiled.barriers().len())
                    },
                    BatchSize::LargeInput,
                );
            });
        }
    }

    group.finish();
}

fn benchmarks(criterion: &mut Criterion) {
    let test = TestVulkan::acquire();

    compile_benchmarks(criterion);

    if let Some(test) = &test {
        resolve_benchmarks(criterion, test.device());
    } else {
        eprintln!("Skipping compile_resolve benchmarks: no Vulkan device is available");
    }
}

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