use std::sync::Arc;

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

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

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

impl Command for PassNoop {
    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 PassNoop {
    fn allowed_in_render_pass(&self) -> bool {
        true
    }

    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 pass_buffer_read_scope() -> SubresourceScope {
    SubresourceScope {
        stage: vk::PipelineStageFlags2::VERTEX_SHADER,
        access: vk::AccessFlags2::SHADER_READ,
        subresource: SubresourceScopeInner::Buffer(BufferRange {
            offset: 0,
            size: 256,
        }),
    }
}

fn pass_texture_read_scope() -> SubresourceScope {
    SubresourceScope {
        stage: vk::PipelineStageFlags2::FRAGMENT_SHADER,
        access: vk::AccessFlags2::SHADER_SAMPLED_READ,
        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::SHADER_READ_ONLY_OPTIMAL,
        },
    }
}

fn supports_render_pass_test(test: &TestVulkan) -> bool {
    let features = test.device().enabled_features().core13;
    if features.synchronization2 != vk::TRUE {
        eprintln!("Skipping graph render pass: synchronization2 is unsupported");
        return false;
    }
    if features.dynamic_rendering != vk::TRUE {
        eprintln!("Skipping graph render pass: dynamic_rendering is unsupported");
        return false;
    }
    true
}

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_scoped_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);
    let alias_barrier = compiled
        .barriers()
        .iter()
        .find(|barrier| barrier.before_command == CommandIndex(1))
        .expect("incoming buffer's first use has an alias barrier");
    assert!(alias_barrier.memory_barriers.is_empty());
    let incoming = alias_barrier
        .resources
        .iter()
        .find(|resource| resource.id == second.id())
        .expect("alias barrier is scoped to the incoming buffer");
    assert_eq!(incoming.src_stage, vk::PipelineStageFlags2::COPY);
    assert_eq!(incoming.src_access, vk::AccessFlags2::TRANSFER_WRITE);
    assert_eq!(incoming.dst_stage, vk::PipelineStageFlags2::COPY);
    assert_eq!(incoming.dst_access, vk::AccessFlags2::TRANSFER_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();
}

#[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();
}

#[test]
fn created_render_pass_records_submits_and_validates() {
    let Some(test) = TestVulkan::acquire() else {
        return;
    };
    if !supports_render_pass_test(&test) {
        return;
    }
    let device = test.device();
    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 render-pass test buffer"),
        )
    };
    let producer_source = make_buffer();
    let producer_destination = make_buffer();
    let unrelated_destination = make_buffer();

    let mut graph = RenderGraph::new();
    let target = graph.create_texture(GraphTextureDesc {
        format: vk::Format::R8G8B8A8_UNORM,
        extent: vk::Extent3D {
            width: 16,
            height: 16,
            depth: 1,
        },
        ..Default::default()
    });
    let producer_source = graph
        .import_buffer(producer_source, ResourceState::default())
        .expect("import producer source");
    let producer_destination = graph
        .import_buffer(producer_destination, ResourceState::default())
        .expect("import producer destination");
    let unrelated_destination = graph
        .import_buffer(unrelated_destination, ResourceState::default())
        .expect("import unrelated destination");
    let buffer_range = BufferRange {
        offset: 0,
        size: 256,
    };
    let image_range = vk::ImageSubresourceRange {
        aspect_mask: vk::ImageAspectFlags::COLOR,
        base_mip_level: 0,
        level_count: 1,
        base_array_layer: 0,
        layer_count: 1,
    };

    graph.push(ClearTexture {
        texture: target,
        range: image_range,
        layout: vk::ImageLayout::TRANSFER_DST_OPTIMAL,
        color: [0.1, 0.2, 0.3, 1.0],
    });
    let producer = graph.push(CopyBuffer {
        src: producer_source,
        src_range: buffer_range,
        dst: producer_destination,
        dst_range: buffer_range,
    });
    let begin = graph.push(BeginRendering {
        color: vec![ColorAttachment::created(
            target,
            ImageViewDesc::color_2d(),
            vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL,
            vk::AttachmentLoadOp::LOAD,
            vk::AttachmentStoreOp::STORE,
            vk::ClearValue::default(),
        )],
        depth: None,
    });
    let interior = graph.push(PassNoop(vec![(
        producer_destination.id(),
        pass_buffer_read_scope(),
    )]));
    let end = graph.push(EndRendering);
    let post_end = graph.push(ScopedNoop(vec![(target.id(), pass_texture_read_scope())]));
    let unrelated = graph.push(CopyBuffer {
        src: producer_destination,
        src_range: buffer_range,
        dst: unrelated_destination,
        dst_range: buffer_range,
    });

    let mut compiled = graph.compile().expect("compile render-pass graph");
    let execution_order = compiled.execution_order();
    let position = |command| {
        execution_order
            .iter()
            .position(|&item| item == command)
            .unwrap()
    };
    assert!(position(begin) < position(interior));
    assert!(position(interior) < position(end));
    assert!(position(unrelated) > position(end));
    let begin_barrier = compiled
        .barriers()
        .iter()
        .find(|barrier| barrier.before_command == begin)
        .expect("incoming barriers are hoisted before BeginRendering");
    assert!(begin_barrier.resources.iter().any(|resource| {
        resource.id == target.id()
            && resource.old_layout == Some(vk::ImageLayout::TRANSFER_DST_OPTIMAL)
            && resource.new_layout == Some(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL)
    }));
    assert!(begin_barrier.resources.iter().any(|resource| {
        resource.id == producer_destination.id()
            && resource.origin == renderer_rendergraph::BarrierOrigin::Command(producer)
    }));
    assert!(compiled.barriers().iter().any(|barrier| {
        barrier.before_command == post_end
            && barrier.resources.iter().any(|resource| {
                resource.id == target.id()
                    && resource.old_layout == Some(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL)
                    && resource.new_layout == Some(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL)
            })
    }));

    compiled
        .resolve(device)
        .expect("resolve created attachment view");
    assert!(
        compiled
            .created_view(target.id(), &ImageViewDesc::color_2d())
            .is_some()
    );
    submit_graph(&test, &mut compiled);
    test.assert_no_validation_errors();
    drop(compiled);
    device
        .wait_idle()
        .expect("wait before render target teardown");
    test.assert_no_validation_errors();
}