use std::sync::Arc;

use ash::vk;
use renderer_device::{
    device::{MemoryLocation, QueueFlags},
    testing::TestVulkan,
};
use renderer_rendergraph::{
    RenderGraph,
    commands::Command,
    commands::CopyBuffer,
    recorder::{CommandIndex, CommandRecorder, RecordError, RecordedCommand, SideEffectMap},
    resources::{
        BufferRange, GraphBufferDesc, GraphTextureDesc, ResourceId, ResourceState,
        SubresourceScope, SubresourceScopeInner,
    },
};
use renderer_resources::buffers::{Buffer, BufferDesc};

struct ScopedNoop(Vec<(ResourceId, SubresourceScope)>);

impl Command for ScopedNoop {
    fn side_effects(&self, command_index: CommandIndex, map: &mut SideEffectMap) {
        for &(resource, scope) in &self.0 {
            map.insert(command_index, resource, scope);
        }
    }
}

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

fn buffer_scope() -> SubresourceScope {
    SubresourceScope {
        stage: vk::PipelineStageFlags2::COPY,
        access: vk::AccessFlags2::TRANSFER_WRITE,
        subresource: SubresourceScopeInner::Buffer(BufferRange {
            offset: 0,
            size: 256,
        }),
    }
}

fn image_scope() -> SubresourceScope {
    SubresourceScope {
        stage: vk::PipelineStageFlags2::COMPUTE_SHADER,
        access: vk::AccessFlags2::SHADER_STORAGE_WRITE,
        subresource: SubresourceScopeInner::Texture {
            range: vk::ImageSubresourceRange {
                aspect_mask: vk::ImageAspectFlags::COLOR,
                base_mip_level: 0,
                level_count: 1,
                base_array_layer: 0,
                layer_count: 1,
            },
            layout: vk::ImageLayout::GENERAL,
        },
    }
}

fn submit_graph(test: &TestVulkan, graph: &mut renderer_rendergraph::CompiledGraph) {
    let device = test.device();
    let pool_info = vk::CommandPoolCreateInfo::default()
        .queue_family_index(device.queues.graphics().family.index.index());
    let pool = unsafe { device.dev().create_command_pool(&pool_info, None) }
        .expect("create alias test command pool");
    let allocation_info = vk::CommandBufferAllocateInfo::default()
        .command_pool(pool)
        .level(vk::CommandBufferLevel::PRIMARY)
        .command_buffer_count(1);
    let command_buffer = unsafe { device.dev().allocate_command_buffers(&allocation_info) }
        .expect("allocate alias test command buffer")[0];
    unsafe {
        device
            .dev()
            .begin_command_buffer(command_buffer, &vk::CommandBufferBeginInfo::default())
            .expect("begin alias test command buffer");
    }
    let mut recorder = CommandRecorder::new(device.dev(), command_buffer);
    graph
        .record_commands(&mut recorder)
        .expect("record alias graph");
    unsafe {
        device
            .dev()
            .end_command_buffer(command_buffer)
            .expect("end alias test command buffer");
        let command_buffers = [command_buffer];
        let submit = vk::SubmitInfo::default().command_buffers(&command_buffers);
        device
            .dev()
            .queue_submit(device.queues.graphics().raw, &[submit], vk::Fence::null())
            .expect("submit alias graph");
    }
    device.wait_idle().expect("wait for alias graph idle");
    unsafe { device.dev().destroy_command_pool(pool, None) };
}

#[test]
fn imported_buffer_copy_records_under_validation() {
    let Some(test) = TestVulkan::acquire() else {
        return;
    };
    let device = test.device();
    if device.enabled_features().core13.synchronization2 != vk::TRUE {
        eprintln!("Skipping graph recording: synchronization2 is unsupported");
        return;
    }
    let make_buffer = || {
        Arc::new(
            Buffer::new(
                device.clone(),
                BufferDesc {
                    size: 256,
                    usage: vk::BufferUsageFlags::TRANSFER_SRC | vk::BufferUsageFlags::TRANSFER_DST,
                    queue_families: QueueFlags::empty(),
                    mem_location: MemoryLocation::CpuToGpu,
                    ..Default::default()
                },
            )
            .expect("create graph test buffer"),
        )
    };
    let source = make_buffer();
    let destination = make_buffer();
    let mut graph = RenderGraph::new();
    let source_id = graph
        .import_buffer(source, ResourceState::default())
        .expect("import source buffer");
    let destination_id = graph
        .import_buffer(destination, ResourceState::default())
        .expect("import destination buffer");
    graph.push(CopyBuffer {
        src: source_id,
        src_range: BufferRange {
            offset: 0,
            size: 256,
        },
        dst: destination_id,
        dst_range: BufferRange {
            offset: 0,
            size: 256,
        },
    });
    let mut graph = graph.compile().expect("compile imported buffer graph");

    let pool_info = vk::CommandPoolCreateInfo::default()
        .queue_family_index(device.queues.graphics().family.index.index());
    let pool = unsafe { device.dev().create_command_pool(&pool_info, None) }
        .expect("create graph command pool");
    let allocation_info = vk::CommandBufferAllocateInfo::default()
        .command_pool(pool)
        .level(vk::CommandBufferLevel::PRIMARY)
        .command_buffer_count(1);
    let command_buffer = unsafe { device.dev().allocate_command_buffers(&allocation_info) }
        .expect("allocate graph command buffer")[0];
    unsafe {
        device
            .dev()
            .begin_command_buffer(command_buffer, &vk::CommandBufferBeginInfo::default())
            .expect("begin graph command buffer");
    }
    let mut recorder = CommandRecorder::new(device.dev(), command_buffer);
    graph
        .record_commands(&mut recorder)
        .expect("record imported buffer graph");
    unsafe {
        device
            .dev()
            .end_command_buffer(command_buffer)
            .expect("end graph command buffer");
        device.dev().destroy_command_pool(pool, None);
    }
    device.wait_idle().expect("wait for graph device idle");
    test.assert_no_validation_errors();
}

#[test]
fn resolved_buffers_alias_and_submit_with_memory_barrier() {
    let Some(test) = TestVulkan::acquire() else {
        return;
    };
    let device = test.device();
    if device.enabled_features().core13.synchronization2 != vk::TRUE {
        eprintln!("Skipping graph aliasing: synchronization2 is unsupported");
        return;
    }
    let bridge = Arc::new(
        Buffer::new(
            device.clone(),
            BufferDesc {
                size: 256,
                usage: vk::BufferUsageFlags::TRANSFER_DST,
                mem_location: MemoryLocation::CpuToGpu,
                ..Default::default()
            },
        )
        .expect("create dependency bridge buffer"),
    );
    let mut graph = RenderGraph::new();
    let first = graph.create_buffer(GraphBufferDesc {
        size: 256,
        ..Default::default()
    });
    let second = graph.create_buffer(GraphBufferDesc {
        size: 256,
        ..Default::default()
    });
    let bridge = graph
        .import_buffer(bridge, ResourceState::default())
        .expect("import dependency bridge");
    graph.push(ScopedNoop(vec![
        (first.id(), buffer_scope()),
        (bridge.id(), buffer_scope()),
    ]));
    graph.push(ScopedNoop(vec![
        (bridge.id(), buffer_scope()),
        (second.id(), buffer_scope()),
    ]));

    let mut compiled = graph.compile().expect("compile buffer alias graph");
    compiled
        .resolve(device)
        .expect("resolve buffer alias graph");
    compiled
        .resolve(device)
        .expect("resolving an already-resolved graph is idempotent");
    let first_placement = compiled.resource_lifetimes()[&first.id()].placement;
    let second_placement = compiled.resource_lifetimes()[&second.id()].placement;
    assert_eq!(first_placement.pool, second_placement.pool);
    assert_eq!(first_placement.offset, second_placement.offset);
    assert!(first_placement.aliased && second_placement.aliased);
    assert!(compiled.barriers().iter().any(|barrier| {
        barrier.before_command == CommandIndex(1) && !barrier.memory_barriers.is_empty()
    }));

    submit_graph(&test, &mut compiled);
    test.assert_no_validation_errors();
    drop(compiled);
    device
        .wait_idle()
        .expect("wait before shared allocation teardown");
    test.assert_no_validation_errors();
}

#[test]
fn resolved_images_alias_and_submit_with_undefined_transition() {
    let Some(test) = TestVulkan::acquire() else {
        return;
    };
    let device = test.device();
    if device.enabled_features().core13.synchronization2 != vk::TRUE {
        eprintln!("Skipping graph aliasing: synchronization2 is unsupported");
        return;
    }
    let bridge = Arc::new(
        Buffer::new(
            device.clone(),
            BufferDesc {
                size: 256,
                usage: vk::BufferUsageFlags::TRANSFER_DST,
                mem_location: MemoryLocation::CpuToGpu,
                ..Default::default()
            },
        )
        .expect("create dependency bridge buffer"),
    );
    let mut graph = RenderGraph::new();
    let first = graph.create_texture(GraphTextureDesc {
        format: vk::Format::R8G8B8A8_UNORM,
        extent: vk::Extent3D {
            width: 16,
            height: 16,
            depth: 1,
        },
        extra_usage: vk::ImageUsageFlags::STORAGE,
        ..Default::default()
    });
    let second = graph.create_texture(GraphTextureDesc {
        format: vk::Format::R8G8B8A8_UNORM,
        extent: vk::Extent3D {
            width: 16,
            height: 16,
            depth: 1,
        },
        extra_usage: vk::ImageUsageFlags::STORAGE,
        ..Default::default()
    });
    let bridge = graph
        .import_buffer(bridge, ResourceState::default())
        .expect("import dependency bridge");
    graph.push(ScopedNoop(vec![
        (first.id(), image_scope()),
        (bridge.id(), buffer_scope()),
    ]));
    graph.push(ScopedNoop(vec![
        (bridge.id(), buffer_scope()),
        (second.id(), image_scope()),
    ]));

    let mut compiled = graph.compile().expect("compile image alias graph");
    compiled.resolve(device).expect("resolve image alias graph");
    compiled
        .resolve(device)
        .expect("resolving an already-resolved graph is idempotent");
    let first_placement = compiled.resource_lifetimes()[&first.id()].placement;
    let second_placement = compiled.resource_lifetimes()[&second.id()].placement;
    assert_eq!(first_placement.pool, second_placement.pool);
    assert_eq!(first_placement.offset, second_placement.offset);
    assert!(first_placement.aliased && second_placement.aliased);
    assert!(compiled.barriers().iter().any(|barrier| {
        barrier.before_command == CommandIndex(1)
            && barrier.resources.iter().any(|resource| {
                resource.id == second.id()
                    && resource.old_layout == Some(vk::ImageLayout::UNDEFINED)
                    && resource.src_access == vk::AccessFlags2::SHADER_STORAGE_WRITE
            })
    }));

    submit_graph(&test, &mut compiled);
    test.assert_no_validation_errors();
    drop(compiled);
    device
        .wait_idle()
        .expect("wait before shared allocation teardown");
    test.assert_no_validation_errors();
}