use std::time::{Duration, Instant};
use wgpu::util::DeviceExt;

const SHADER: &str = r#"
@group(0) @binding(0) var<storage, read_write> buf: array<u32>;
@group(0) @binding(1) var<uniform> p: Params;
struct Params { iters: u32 };

@compute @workgroup_size(64)
fn cs_main(@builtin(global_invocation_id) gid: vec3<u32>) {
    var x: u32 = gid.x ^ 0xDEADBEEFu;
    for (var i: u32 = 0u; i < p.iters; i = i + 1u) {
        x = x ^ (x << 13u);
        x = x ^ (x >> 17u);
        x = x ^ (x << 5u);
    }
    buf[gid.x] = x;
}
"#;

fn main() { pollster::block_on(run()); }

async fn run() {
    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
        backends: wgpu::Backends::METAL,
        backend_options: Default::default(),
        flags: wgpu::InstanceFlags::default(),
        memory_budget_thresholds: Default::default(),
        display: None,
    });
    let adapter = instance.request_adapter(&wgpu::RequestAdapterOptions {
        power_preference: wgpu::PowerPreference::HighPerformance,
        force_fallback_adapter: false,
        compatible_surface: None,
    }).await.unwrap();

    let (device, queue) = adapter.request_device(&wgpu::DeviceDescriptor {
        label: None,
        required_features: wgpu::Features::default(),
        required_limits: wgpu::Limits::default(),
        memory_hints: Default::default(),
        trace: wgpu::Trace::Off,
        experimental_features: Default::default(),
    }).await.unwrap();

    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: None,
        source: wgpu::ShaderSource::Wgsl(SHADER.into()),
    });
    let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: None,
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0, visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Storage { read_only: false },
                    has_dynamic_offset: false, min_binding_size: None,
                },
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 1, visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false, min_binding_size: None,
                },
                count: None,
            },
        ],
    });
    let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: None,
        bind_group_layouts: &[Some(&bgl)],
        immediate_size: 0,
    });
    let pipe = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: None, layout: Some(&pl), module: &module,
        entry_point: Some("cs_main"),
        compilation_options: Default::default(), cache: None,
    });

    const N_THREADS: u32 = 1024 * 64;
    const ITERS: u32 = 5_000_000;
    let buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: None, size: (N_THREADS as u64) * 4,
        usage: wgpu::BufferUsages::STORAGE, mapped_at_creation: false,
    });
    let ubuf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None, contents: bytemuck::bytes_of(&ITERS),
        usage: wgpu::BufferUsages::UNIFORM,
    });
    let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None, layout: &bgl, entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: ubuf.as_entire_binding() },
        ],
    });
    let encode_and_submit = || {
        let mut enc = device.create_command_encoder(&Default::default());
        {
            let mut pass = enc.begin_compute_pass(&Default::default());
            pass.set_pipeline(&pipe);
            pass.set_bind_group(0, &bg, &[]);
            pass.dispatch_workgroups(N_THREADS / 64, 1, 1);
        }
        queue.submit(Some(enc.finish()));
    };

    // [1] Reference: bounded-timeout poll loop completes promptly.
    eprintln!("[1] bounded-timeout poll loop:");
    encode_and_submit();
    let t0 = Instant::now();
    loop {
        let r = device.poll(wgpu::PollType::Wait {
            submission_index: None, timeout: Some(Duration::from_millis(200)),
        });
        eprintln!("  poll @ {:>5.2}s -> {r:?}", t0.elapsed().as_secs_f32());
        match r {
            Ok(wgpu::PollStatus::QueueEmpty | wgpu::PollStatus::WaitSucceeded) => break,
            Err(wgpu::PollError::Timeout) => continue,
            other => panic!("unexpected: {other:?}"),
        }
    }
    eprintln!("[1] elapsed: {:.2}s", t0.elapsed().as_secs_f32());

    // [2] Bug: wait_indefinitely on the same workload deadlocks.
    eprintln!("\n[2] wait_indefinitely on the same workload:");
    encode_and_submit();
    let t0 = Instant::now();
    let _wd = std::thread::spawn(|| {
        std::thread::sleep(Duration::from_secs(60));
        eprintln!("*** watchdog: poll(wait_indefinitely) has not returned after 60s — DEADLOCK ***");
        std::process::exit(2);
    });
    let r = device.poll(wgpu::PollType::wait_indefinitely());
    eprintln!("[2] returned at {:.2}s -> {r:?}", t0.elapsed().as_secs_f32());
}

loop {
    if let MTLCommandBufferStatus::Completed = cmd_buf.status() {
        return Ok(true);
    }
    if let Some(timeout) = timeout {
        if start.elapsed() >= timeout { return Ok(false); }
    }
    thread::sleep(Duration::from_millis(1));
}