//! Launch-overhead probe — null dispatch and 1 MiB host→device copy.
//!
//! Two median latencies, both reported in nanoseconds:
//!
//! - **Null launch.** Pipeline with an empty `@workgroup_size(1)` shader
//!   and no descriptors. Measures submit → queue-idle round trip for the
//!   smallest unit of GPU work the host can ask for.
//! - **1 MiB memcpy.** `cmd_copy_buffer` from a host-visible staging
//!   buffer into a device-local buffer, both 1 MiB. Measures the cost
//!   of staging the smallest practical chunk of work onto the device.
//!
//! Both numbers are medians over a sample set so OS-scheduling jitter
//! does not skew the result.

use std::time::Instant;

use ash::vk;
use naga::back::spv;
use naga::front::wgsl;
use naga::valid;

use super::vulkan_ctx::VulkanContext;

const NULL_LAUNCH_WGSL: &str = r"
@compute @workgroup_size(1)
fn main() {}
";

const SAMPLE_COUNT: usize = 100;
const WARMUP_COUNT: usize = 8;
const MEMCPY_BYTES: vk::DeviceSize = 1024 * 1024;

#[derive(Clone, Copy, Debug)]
pub struct LaunchOverheadResult {
    pub null_launch_ns: f64,
    pub memcpy_1mib_ns: f64,
}

#[derive(Debug)]
pub enum LaunchError {
    ShaderParse(String),
    ShaderValidate(String),
    ShaderEmit(String),
    Vulkan(vk::Result),
    NoDeviceLocalMemory,
    NoHostVisibleMemory,
}

impl core::fmt::Display for LaunchError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::ShaderParse(s) => write!(f, "WGSL parse failed: {s}"),
            Self::ShaderValidate(s) => write!(f, "WGSL validation failed: {s}"),
            Self::ShaderEmit(s) => write!(f, "SPIR-V emission failed: {s}"),
            Self::Vulkan(r) => write!(f, "Vulkan call failed: {r:?}"),
            Self::NoDeviceLocalMemory => write!(f, "no DEVICE_LOCAL memory type available"),
            Self::NoHostVisibleMemory => {
                write!(f, "no HOST_VISIBLE | HOST_COHERENT memory type available")
            }
        }
    }
}

impl std::error::Error for LaunchError {}

pub fn measure_launch_overhead(ctx: &VulkanContext) -> Result<LaunchOverheadResult, LaunchError> {
    let null_launch_ns = measure_null_launch(ctx)?;
    let memcpy_1mib_ns = measure_memcpy_1mib(ctx)?;
    Ok(LaunchOverheadResult {
        null_launch_ns,
        memcpy_1mib_ns,
    })
}

fn measure_null_launch(ctx: &VulkanContext) -> Result<f64, LaunchError> {
    let spirv = compile_null_shader()?;
    let device = &ctx.device;

    let pl_info = vk::PipelineLayoutCreateInfo::default();
    let pipeline_layout =
        unsafe { device.create_pipeline_layout(&pl_info, None) }.map_err(LaunchError::Vulkan)?;

    let shader_info = vk::ShaderModuleCreateInfo::default().code(&spirv);
    let shader =
        unsafe { device.create_shader_module(&shader_info, None) }.map_err(LaunchError::Vulkan)?;

    let stage = vk::PipelineShaderStageCreateInfo::default()
        .stage(vk::ShaderStageFlags::COMPUTE)
        .module(shader)
        .name(c"main");
    let info = [vk::ComputePipelineCreateInfo::default()
        .stage(stage)
        .layout(pipeline_layout)];
    let pipelines =
        unsafe { device.create_compute_pipelines(vk::PipelineCache::null(), &info, None) }
            .map_err(|(_, r)| LaunchError::Vulkan(r))?;
    let pipeline = pipelines[0];

    let cmd_info = vk::CommandBufferAllocateInfo::default()
        .command_pool(ctx.command_pool)
        .level(vk::CommandBufferLevel::PRIMARY)
        .command_buffer_count(1);
    let cmd =
        unsafe { device.allocate_command_buffers(&cmd_info) }.map_err(LaunchError::Vulkan)?[0];

    let record = || -> Result<(), LaunchError> {
        let begin = vk::CommandBufferBeginInfo::default();
        unsafe {
            device
                .begin_command_buffer(cmd, &begin)
                .map_err(LaunchError::Vulkan)?;
            device.cmd_bind_pipeline(cmd, vk::PipelineBindPoint::COMPUTE, pipeline);
            device.cmd_dispatch(cmd, 1, 1, 1);
            device
                .end_command_buffer(cmd)
                .map_err(LaunchError::Vulkan)?;
        }
        Ok(())
    };

    let ns = sample_command(ctx, cmd, record)?;

    unsafe {
        device.free_command_buffers(ctx.command_pool, &[cmd]);
        device.destroy_pipeline(pipeline, None);
        device.destroy_shader_module(shader, None);
        device.destroy_pipeline_layout(pipeline_layout, None);
    }

    Ok(ns)
}

fn measure_memcpy_1mib(ctx: &VulkanContext) -> Result<f64, LaunchError> {
    let device = &ctx.device;

    let device_mem_type = ctx
        .find_memory_type(u32::MAX, vk::MemoryPropertyFlags::DEVICE_LOCAL)
        .ok_or(LaunchError::NoDeviceLocalMemory)?;
    let host_mem_type = ctx
        .find_memory_type(
            u32::MAX,
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
        )
        .ok_or(LaunchError::NoHostVisibleMemory)?;

    let src = create_buffer(
        device,
        MEMCPY_BYTES,
        vk::BufferUsageFlags::TRANSFER_SRC,
        host_mem_type,
    )?;
    let dst = match create_buffer(
        device,
        MEMCPY_BYTES,
        vk::BufferUsageFlags::TRANSFER_DST,
        device_mem_type,
    ) {
        Ok(d) => d,
        Err(e) => {
            unsafe { destroy_buffer(device, src) };
            return Err(e);
        }
    };

    let cmd_info = vk::CommandBufferAllocateInfo::default()
        .command_pool(ctx.command_pool)
        .level(vk::CommandBufferLevel::PRIMARY)
        .command_buffer_count(1);
    let cmd_result = unsafe { device.allocate_command_buffers(&cmd_info) };
    let cmd = match cmd_result {
        Ok(v) => v[0],
        Err(e) => {
            unsafe {
                destroy_buffer(device, dst);
                destroy_buffer(device, src);
            }
            return Err(LaunchError::Vulkan(e));
        }
    };

    let region = [vk::BufferCopy::default()
        .src_offset(0)
        .dst_offset(0)
        .size(MEMCPY_BYTES)];

    let record = || -> Result<(), LaunchError> {
        let begin = vk::CommandBufferBeginInfo::default();
        unsafe {
            device
                .begin_command_buffer(cmd, &begin)
                .map_err(LaunchError::Vulkan)?;
            device.cmd_copy_buffer(cmd, src.buffer, dst.buffer, &region);
            device
                .end_command_buffer(cmd)
                .map_err(LaunchError::Vulkan)?;
        }
        Ok(())
    };

    let ns = sample_command(ctx, cmd, record);

    unsafe {
        device.free_command_buffers(ctx.command_pool, &[cmd]);
        destroy_buffer(device, dst);
        destroy_buffer(device, src);
    }

    ns
}

struct AllocatedBuffer {
    buffer: vk::Buffer,
    memory: vk::DeviceMemory,
}

fn create_buffer(
    device: &ash::Device,
    size: vk::DeviceSize,
    usage: vk::BufferUsageFlags,
    memory_type: u32,
) -> Result<AllocatedBuffer, LaunchError> {
    let info = vk::BufferCreateInfo::default()
        .size(size)
        .usage(usage)
        .sharing_mode(vk::SharingMode::EXCLUSIVE);
    let buffer = unsafe { device.create_buffer(&info, None) }.map_err(LaunchError::Vulkan)?;
    let reqs = unsafe { device.get_buffer_memory_requirements(buffer) };
    let alloc = vk::MemoryAllocateInfo::default()
        .allocation_size(reqs.size)
        .memory_type_index(memory_type);
    let memory = match unsafe { device.allocate_memory(&alloc, None) } {
        Ok(m) => m,
        Err(e) => {
            unsafe { device.destroy_buffer(buffer, None) };
            return Err(LaunchError::Vulkan(e));
        }
    };
    if let Err(e) = unsafe { device.bind_buffer_memory(buffer, memory, 0) } {
        unsafe {
            device.free_memory(memory, None);
            device.destroy_buffer(buffer, None);
        }
        return Err(LaunchError::Vulkan(e));
    }
    Ok(AllocatedBuffer { buffer, memory })
}

// Takes ownership intentionally: the buffer is no longer usable after this.
#[allow(clippy::needless_pass_by_value)]
unsafe fn destroy_buffer(device: &ash::Device, b: AllocatedBuffer) {
    unsafe {
        device.destroy_buffer(b.buffer, None);
        device.free_memory(b.memory, None);
    }
}

fn sample_command<R>(
    ctx: &VulkanContext,
    cmd: vk::CommandBuffer,
    record: R,
) -> Result<f64, LaunchError>
where
    R: Fn() -> Result<(), LaunchError>,
{
    let device = &ctx.device;
    record()?;

    for _ in 0..WARMUP_COUNT {
        submit_and_wait(ctx, cmd)?;
        unsafe {
            device
                .reset_command_buffer(cmd, vk::CommandBufferResetFlags::empty())
                .map_err(LaunchError::Vulkan)?;
        }
        record()?;
    }

    let mut samples = Vec::with_capacity(SAMPLE_COUNT);
    for _ in 0..SAMPLE_COUNT {
        let start = Instant::now();
        submit_and_wait(ctx, cmd)?;
        #[allow(clippy::cast_precision_loss)]
        let ns = start.elapsed().as_nanos() as f64;
        samples.push(ns);
        unsafe {
            device
                .reset_command_buffer(cmd, vk::CommandBufferResetFlags::empty())
                .map_err(LaunchError::Vulkan)?;
        }
        record()?;
    }

    Ok(median(&mut samples))
}

fn submit_and_wait(ctx: &VulkanContext, cmd: vk::CommandBuffer) -> Result<(), LaunchError> {
    let info = vk::SubmitInfo::default().command_buffers(std::slice::from_ref(&cmd));
    unsafe {
        ctx.device
            .queue_submit(ctx.queue, &[info], vk::Fence::null())
            .map_err(LaunchError::Vulkan)?;
        ctx.device
            .queue_wait_idle(ctx.queue)
            .map_err(LaunchError::Vulkan)?;
    }
    Ok(())
}

fn median(samples: &mut [f64]) -> f64 {
    samples.sort_by(f64::total_cmp);
    let mid = samples.len() / 2;
    if samples.len() % 2 == 0 {
        f64::midpoint(samples[mid - 1], samples[mid])
    } else {
        samples[mid]
    }
}

fn compile_null_shader() -> Result<Vec<u32>, LaunchError> {
    let module = wgsl::parse_str(NULL_LAUNCH_WGSL)
        .map_err(|e| LaunchError::ShaderParse(format!("{e:?}")))?;
    let info = valid::Validator::new(valid::ValidationFlags::all(), valid::Capabilities::all())
        .validate(&module)
        .map_err(|e| LaunchError::ShaderValidate(format!("{e:?}")))?;
    let options = spv::Options::default();
    spv::write_vec(&module, &info, &options, None)
        .map_err(|e| LaunchError::ShaderEmit(format!("{e:?}")))
}