use encase::ShaderType;
use hashbrown::HashSet;
use roast2d::prelude::*;
use roast2d::render::BackendState;
use roast2d::renderer::shader2d::default::DefaultShader;
use roast2d::renderer::traits::{Shader, SpritePipeline};
use roast2d::utils::bytes::to_wgsl_bytes;
use std::borrow::Cow;
use std::rc::Rc;
use wgpu::util::DeviceExt;

use roast2d::engine::Engine;
use roast2d::prelude::renderer_types::{CameraUniform, TextureResource, Vertex};
use roast2d::prelude::wgpu::{self, Device, RenderPass};
use wgpu::{BindGroup, BindGroupLayout, BufferUsages, RenderPipeline, ShaderModule, TextureFormat};

const SIZE: Vec2 = Vec2::splat(20.0);
const GAP: f32 = 10.0;
const COUNT: usize = 400;
const COLS: usize = 20;

pub struct Demo {
    blink_renderer: Rc<BlinkShader>,
    timer: f32,
    blink_timer: Vec<f32>,
    visited: HashSet<(usize, usize)>,
    queue: Vec<(usize, usize)>,
}

impl Scene for Demo {
    fn update(&mut self, g: &mut Engine) {
        self.timer += g.tick;

        // trigger blink
        let speed = 0.1;
        while self.timer > speed {
            self.timer -= speed;
            let queue = std::mem::take(&mut self.queue);
            for (x, y) in queue {
                let index = x + y * COLS;
                self.blink_timer[index] = 1.0;
                self.visited.insert((x, y));
                for (dx, dy) in [(-1, 0), (1, 0), (0, -1), (0, 1)] {
                    let nx = (x as i32 + dx) as usize % COLS;
                    let ny = (y as i32 + dy) as usize % COLS;
                    if !self.visited.contains(&(nx, ny)) {
                        self.queue.push((nx, ny));
                    }
                }
            }
        }

        if self.queue.is_empty() {
            self.queue.push((0, 0));
            self.visited.clear();
        }

        // update blink timer
        for timer in self.blink_timer.iter_mut() {
            *timer -= g.tick;
        }
    }

    fn draw(&mut self, g: &mut Engine) {
        let length = (SIZE + GAP) * COLS as f32 - 2.0 * GAP;
        let offset = (g.view_size() - length) * 0.5;
        for i in 0..COUNT {
            let x = (i % COLS) as f32 * (SIZE.x + GAP);
            let y = (i / COLS) as f32 * (SIZE.y + GAP);
            let pos = Vec2::new(x, y);
            let transform = Transform::new((pos + offset).extend(0.0));
            let blink_timer = self.blink_timer[i];
            g.with_renderer(self.blink_renderer.clone(), |g| {
                g.send_renderer_data(to_wgsl_bytes(&BlinkUniform::new(blink_timer)));
                g.draw_rect(BLUE, SIZE, None, transform);
            });
        }
    }
}

fn setup(g: &mut Engine) {
    // set resize and scale
    g.set_scale_mode(ScaleMode::Fixed {
        width: 800,
        height: 600,
    });
    let blink_renderer = Rc::new(BlinkShader::new(g.backend_state()));
    g.set_scene(Demo {
        blink_renderer,
        timer: 0.0,
        blink_timer: vec![0.0; COUNT],
        visited: HashSet::new(),
        queue: Vec::new(),
    });
}

fn run() {
    let app = App::default()
        .title("TextBox".to_string())
        .window(UVec2::new(800, 600))
        .vsync(true)
        .resizable(true);
    app.run(setup).expect("Start game");
}

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

#[derive(Debug, Copy, Clone, ShaderType)]
pub struct BlinkUniform {
    #[align(256)]
    pub blink: f32,
}

impl BlinkUniform {
    pub fn new(blink: f32) -> Self {
        Self { blink }
    }
}

pub(crate) struct BlinkPipeline {
    pipeline: RenderPipeline,
    camera_bind_group: BindGroup,
    texture_bind_group_layout: BindGroupLayout,
    device: Device,
}

impl SpritePipeline for BlinkPipeline {
    fn render(
        &self,
        pass: &mut RenderPass,
        texture: Option<&TextureResource>,
        shader_data: Option<&[u8]>,
    ) {
        let texture = texture.expect("BlinkPipeline requires a texture");
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, &self.camera_bind_group, &[]);

        // Create per-draw buffer for blink value (single f32 with alignment from encase)
        let buffer = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("blink_buffer"),
                usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
                contents: shader_data.expect("blink shader data required"),
            });
        let texture_bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            layout: &self.texture_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&texture.view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&texture.sampler),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: buffer.as_entire_binding(),
                },
            ],
            label: Some("Texture Bind Group"),
        });
        pass.set_bind_group(1, &texture_bind_group, &[]);
    }
}

pub(crate) struct BlinkShader {
    device: Device,
    shader: ShaderModule,
    vertex_shader: ShaderModule,
    camera_bind_group_layout: BindGroupLayout,
    texture_bind_group_layout: BindGroupLayout,
}

impl BlinkShader {
    const KEY: &'static str = "blink";
    const SHADER: &'static str = include_str!("../assets/shaders/blink.wgsl");

    pub fn new(state: &BackendState) -> Self {
        let device = state.device.clone();
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some(Self::KEY),
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(Self::SHADER)),
        });
        let vertex_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("vertex_shader"),
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(DefaultShader::SHADER)),
        });
        let camera_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: wgpu::BufferSize::new(CameraUniform::min_size().into()),
                    },
                    count: None,
                }],
                label: Some("camera_bind_group_layout"),
            });

        let texture_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Texture Bind Group Layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 2,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Uniform,
                            has_dynamic_offset: false,
                            min_binding_size: wgpu::BufferSize::new(
                                BlinkUniform::min_size().into(),
                            ),
                        },
                        count: None,
                    },
                ],
            });

        Self {
            shader,
            vertex_shader,
            camera_bind_group_layout,
            texture_bind_group_layout,
            device,
        }
    }
}

impl Shader for BlinkShader {
    fn key(&self) -> &'static str {
        Self::KEY
    }

    fn new_pipeline(
        &self,
        format: TextureFormat,
        _depth_format: Option<TextureFormat>,
        camera_buffer: &wgpu::Buffer,
        _transparent: bool,
    ) -> Box<dyn SpritePipeline> {
        let camera_bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            layout: &self.camera_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: camera_buffer.as_entire_binding(),
            }],
            label: Some("camera_bind_group"),
        });
        let pipeline_layout = self
            .device
            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Render Pipeline Layout"),
                bind_group_layouts: &[
                    &self.camera_bind_group_layout,
                    &self.texture_bind_group_layout,
                ],
                push_constant_ranges: &[],
            });
        let pipeline = self
            .device
            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
                cache: None,
                label: Some("Render Pipeline"),
                layout: Some(&pipeline_layout),
                vertex: wgpu::VertexState {
                    module: &self.vertex_shader,
                    entry_point: Some("vs_main"),
                    buffers: &[Vertex::desc()],
                    compilation_options: Default::default(),
                },
                fragment: Some(wgpu::FragmentState {
                    module: &self.shader,
                    entry_point: Some("fs_main"),
                    targets: &[Some(wgpu::ColorTargetState {
                        format,
                        blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                        write_mask: wgpu::ColorWrites::ALL,
                    })],
                    compilation_options: Default::default(),
                }),
                primitive: wgpu::PrimitiveState {
                    topology: wgpu::PrimitiveTopology::TriangleList,
                    strip_index_format: None,
                    front_face: wgpu::FrontFace::Ccw,
                    cull_mode: Some(wgpu::Face::Back),
                    polygon_mode: wgpu::PolygonMode::Fill,
                    unclipped_depth: false,
                    conservative: false,
                },
                depth_stencil: Some(wgpu::DepthStencilState {
                    format: _depth_format.unwrap(),
                    depth_write_enabled: true,
                    depth_compare: wgpu::CompareFunction::Less,
                    stencil: wgpu::StencilState::default(),
                    bias: wgpu::DepthBiasState::default(),
                }),
                multisample: wgpu::MultisampleState::default(),
                multiview: None,
            });
        let pipeline = BlinkPipeline {
            pipeline,
            camera_bind_group,
            texture_bind_group_layout: self.texture_bind_group_layout.clone(),
            device: self.device.clone(),
        };
        Box::new(pipeline)
    }
}
