use rand::Rng;
use raw_window_handle::{
    RawDisplayHandle, RawWindowHandle, WaylandDisplayHandle, WaylandWindowHandle,
};
use smithay_client_toolkit::{
    compositor::CompositorState,
    output::OutputState,
    registry::RegistryState,
    seat::SeatState,
    shell::{
        WaylandSurface,
        wlr_layer::{Anchor, LayerShell, LayerSurface},
    },
};
use std::{borrow::Cow, env::args, ptr::NonNull, str::FromStr, time::Instant};
use wayland_client::{Connection, Proxy, QueueHandle, globals::registry_queue_init};
use wgpu::{BindGroup, Buffer, util::DeviceExt};

use crate::color_palette::ColorPalette;

mod color_palette;
mod implementations;

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct Vertex {
    position: [f32; 2],
}
impl Vertex {
    #[allow(dead_code)]
    fn triangle(x: f32, y: f32, w: f32, h: f32) -> [Self; 3] {
        let x0 = x;
        let x1 = x + w;
        let y0 = y;
        let y1 = y + h;

        [
            Vertex { position: [x0, y0] },
            Vertex { position: [x1, y0] },
            Vertex { position: [x1, y1] },
        ]
    }
    #[allow(dead_code)]
    fn rectangle(x: f32, y: f32, w: f32, h: f32) -> [Self; 6] {
        let x0 = x;
        let x1 = x + w;
        let y0 = y;
        let y1 = y + h;

        [
            Vertex { position: [x0, y0] }, // Triangle 1
            Vertex { position: [x1, y0] },
            Vertex { position: [x1, y1] },
            Vertex { position: [x0, y0] }, // Triangle 2
            Vertex { position: [x1, y1] },
            Vertex { position: [x0, y1] },
        ]
    }
}

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct InstanceData {
    direction: [f32; 2],
    color: [f32; 3],
}

#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct Uniforms {
    time: f32,
}

impl Uniforms {
    fn new() -> Self {
        Self { time: 0.0 }
    }
}

fn main() {
    env_logger::init();

    let conn = Connection::connect_to_env().unwrap();
    let (globals, mut event_queue) = registry_queue_init(&conn).unwrap();
    let qh = event_queue.handle();

    // Initialize xdg_shell handlers so we can select the correct adapter
    let compositor_state =
        CompositorState::bind(&globals, &qh).expect("wl_compositor not available");
    let layer_state = LayerShell::bind(&globals, &qh).expect("layer_shell not available");

    let surface = compositor_state.create_surface(&qh);
    // Create the window for adapter selection
    let layer = layer_state.create_layer_surface(
        &qh,
        surface,
        smithay_client_toolkit::shell::wlr_layer::Layer::Top,
        Some(""),
        None,
    );
    layer.set_anchor(Anchor::TOP | Anchor::BOTTOM | Anchor::LEFT | Anchor::RIGHT);
    let (width, height) = (400, 400);
    layer.set_size(0, 0); // 0 width = stretch to full width
    layer.set_opaque_region(None);
    layer.commit();

    // Initialize wgpu
    let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
        backends: wgpu::Backends::all(),
        ..Default::default()
    });

    // Create the raw window handle for the surface.
    let raw_display_handle = RawDisplayHandle::Wayland(WaylandDisplayHandle::new(
        NonNull::new(conn.backend().display_ptr() as *mut _).unwrap(),
    ));
    let raw_window_handle = RawWindowHandle::Wayland(WaylandWindowHandle::new(
        NonNull::new(layer.wl_surface().id().as_ptr() as *mut _).unwrap(),
    ));

    let surface = unsafe {
        instance
            .create_surface_unsafe(wgpu::SurfaceTargetUnsafe::RawHandle {
                raw_display_handle,
                raw_window_handle,
            })
            .unwrap()
    };

    // Pick a supported adapter
    let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
        compatible_surface: Some(&surface),
        ..Default::default()
    }))
    .expect("Failed to find suitable adapter");

    let (device, queue) = pollster::block_on(adapter.request_device(&Default::default()))
        .expect("Failed to request device");

    let mut surface_config = surface.get_default_config(&adapter, 200, 200).unwrap();
    surface_config.alpha_mode = wgpu::CompositeAlphaMode::PreMultiplied;
    surface_config.format = wgpu::TextureFormat::Bgra8Unorm;

    let (layout, group, uniform_buffer, uniforms) = create_uniforms(&device);
    let (vertex_buffer, instance_buffer, vertex_count, instance_count) =
        create_vertex_buffer(&device, width as f32, height as f32);

    let render_pipeline = create_pipeline(&device, surface_config.format, &layout);

    let mut wgpu = Wgpu {
        registry_state: RegistryState::new(&globals),
        seat_state: SeatState::new(&globals, &qh),
        output_state: OutputState::new(&globals, &qh),

        start_time: Instant::now(),
        first_configure: true,
        exit: false,
        width: 256,
        height: 256,
        window: layer,
        device,
        surface,
        queue,
        render_pipeline,
        group,

        uniforms,
        uniform_buffer,
        vertex_buffer,
        instance_buffer,
        vertex_count,
        instance_count,
    };

    // We don't draw immediately, the configure will notify us when to first draw.
    loop {
        event_queue.blocking_dispatch(&mut wgpu).unwrap();

        if wgpu.exit {
            break;
        }
    }

    // On exit we must destroy the surface before the window is destroyed.
    drop(wgpu.surface);
    drop(wgpu.window);
}
struct Wgpu {
    registry_state: RegistryState,
    seat_state: SeatState,
    output_state: OutputState,

    start_time: Instant,
    exit: bool,
    first_configure: bool,
    width: u32,
    height: u32,
    window: LayerSurface,

    device: wgpu::Device,
    queue: wgpu::Queue,
    surface: wgpu::Surface<'static>,
    render_pipeline: wgpu::RenderPipeline,
    group: wgpu::BindGroup,

    uniforms: Uniforms,
    uniform_buffer: Buffer,
    vertex_buffer: Buffer,
    instance_buffer: Buffer,
    vertex_count: u32,
    instance_count: u32,
}

impl Wgpu {
    fn draw(&mut self, _qh: &QueueHandle<Self>) {
        let elapsed = self.start_time.elapsed().as_secs_f32();
        self.update_time(elapsed);
        if elapsed > 3.0 {
            self.exit = true
        }
        let surface_texture = self
            .surface
            .get_current_texture()
            .expect("Failed to acquire next swap chain texture");
        let texture_view = surface_texture
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default());
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: None,
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &texture_view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
                        store: wgpu::StoreOp::Store,
                    },
                    depth_slice: None,
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            });
            rpass.set_bind_group(0, &self.group, &[]);
            rpass.set_pipeline(&self.render_pipeline);
            rpass.set_vertex_buffer(0, self.vertex_buffer.slice(..));
            rpass.set_vertex_buffer(1, self.instance_buffer.slice(..));
            rpass.draw(0..self.vertex_count, 0..self.instance_count);
        }
        self.queue.submit(Some(encoder.finish()));
        surface_texture.present();
    }
    pub fn update_time(&mut self, time: f32) {
        self.uniforms.time = time;
        self.queue
            .write_buffer(&self.uniform_buffer, 0, bytemuck::bytes_of(&self.uniforms));
    }
}

fn create_pipeline(
    device: &wgpu::Device,
    swap_chain_format: wgpu::TextureFormat,
    bind_group_layout: &wgpu::BindGroupLayout,
) -> wgpu::RenderPipeline {
    // Load theushaders from disk
    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: None,
        source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("shader.wgsl"))),
    });
    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: Some("Render Pipeline Layout"),
        bind_group_layouts: &[bind_group_layout],
        push_constant_ranges: &[],
    });
    let vertex_buffer_layout = wgpu::VertexBufferLayout {
        array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
        step_mode: wgpu::VertexStepMode::Vertex,
        attributes: &[wgpu::VertexAttribute {
            format: wgpu::VertexFormat::Float32x2,
            offset: 0,
            shader_location: 0,
        }],
    };
    let instance_buffer_layout = wgpu::VertexBufferLayout {
        array_stride: std::mem::size_of::<InstanceData>() as wgpu::BufferAddress,
        step_mode: wgpu::VertexStepMode::Instance,
        attributes: &[
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Float32x2, // v_start
                offset: 0,
                shader_location: 1,
            },
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Float32x3, // color
                offset: std::mem::size_of::<[f32; 2]>() as wgpu::BufferAddress,
                shader_location: 2,
            },
        ],
    };
    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: None,
        layout: Some(&pipeline_layout),
        vertex: wgpu::VertexState {
            module: &shader,
            entry_point: Some("vs_main"),
            buffers: &[vertex_buffer_layout, instance_buffer_layout],
            compilation_options: Default::default(),
        },
        fragment: Some(wgpu::FragmentState {
            module: &shader,
            entry_point: Some("fs_main"),
            compilation_options: Default::default(),
            targets: &[Some(wgpu::ColorTargetState {
                format: swap_chain_format,
                blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                write_mask: wgpu::ColorWrites::ALL,
            })],
        }),
        primitive: wgpu::PrimitiveState {
            topology: wgpu::PrimitiveTopology::TriangleList,
            // strip_index_format: None,
            // front_face: wgpu::FrontFace::Ccw,
            // cull_mode: Some(wgpu::Face::Back),
            // // Setting this to anything other than Fill requires Features::POLYGON_MODE_LINE
            // // or Features::POLYGON_MODE_POINT
            // polygon_mode: wgpu::PolygonMode::Fill,
            // // Requires Features::DEPTH_CLIP_CONTROL
            // unclipped_depth: false,
            // // Requires Features::CONSERVATIVE_RASTERIZATION
            // conservative: false,
            ..Default::default()
        },
        depth_stencil: None,
        multisample: wgpu::MultisampleState::default(),
        multiview: None,
        cache: None,
    })
}

fn create_uniforms(device: &wgpu::Device) -> (wgpu::BindGroupLayout, BindGroup, Buffer, Uniforms) {
    let uniforms = Uniforms::new();
    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("Uniform Buffer"),
        contents: bytemuck::bytes_of(&uniforms),
        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    });

    let uniform_bind_group_layout =
        device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
            label: Some("uniform_bind_group_layout"),
        });

    let uniform_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        layout: &uniform_bind_group_layout,
        entries: &[wgpu::BindGroupEntry {
            binding: 0,
            resource: uniform_buffer.as_entire_binding(),
        }],
        label: Some("uniform_bind_group"),
    });

    (
        uniform_bind_group_layout,
        uniform_bind_group,
        uniform_buffer,
        uniforms,
    )
}

fn create_vertex_buffer(
    device: &wgpu::Device,
    width: f32,
    height: f32,
) -> (Buffer, Buffer, u32, u32) {
    // Only 1 rectangle vertices here, since instances define position:
    let rectangle = Vertex::rectangle(0.0, 0.0, 0.00002 * height, 0.00002 * width);
    let args = args().collect::<Vec<String>>();
    let pallette = extreact_flag_value::<ColorPalette>(&args, "--pallette").unwrap_or_default();
    let colors = pallette.get_colors();
    let color_count = colors.len();

    let mut rng = rand::rng();
    let instances = (0..200)
        .map(|_| {
            let x = rng.random_range(-1.0..1.0) as f32;
            let y_max = (1.0 - x * x).sqrt() * 2.5;
            let y = rng.random_range(-0.5..y_max);
            InstanceData {
                direction: [x * 1.2, y],
                color: colors
                    .get(rng.random_range(0..color_count))
                    .unwrap()
                    .clone(),
            }
        })
        .collect::<Vec<InstanceData>>();

    let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("Rectangle Vertex Buffer"),
        contents: bytemuck::cast_slice(&rectangle),
        usage: wgpu::BufferUsages::VERTEX,
    });

    let instance_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("Instance Buffer"),
        contents: bytemuck::cast_slice(&instances),
        usage: wgpu::BufferUsages::VERTEX,
    });

    (
        vertex_buffer,
        instance_buffer,
        rectangle.len() as u32,
        instances.len() as u32,
    )
}

fn extreact_flag_value<T: FromStr>(args: &Vec<String>, name: &str) -> Option<T> {
    let pos = args.iter().position(|arg| arg == name)?;
    let val = args.get(pos + 1)?;
    if val.starts_with("-") {
        return None;
    }
    val.parse::<T>().ok()
}