use std::collections::{HashMap, HashSet};
use weareshared::{packets::Resource, tree::SceneTree};
use wgpu::{
    BindGroup, BindGroupDescriptor, BindGroupLayoutDescriptor, BlendState, Color, ColorTargetState,
    ColorWrites, CommandEncoder, Device, FragmentState, FrontFace, IndexFormat, LoadOp,
    MultisampleState, Operations, PipelineCompilationOptions, PipelineLayoutDescriptor,
    PolygonMode, PrimitiveState, PrimitiveTopology, RenderPassColorAttachment,
    RenderPassDescriptor, RenderPipeline, RenderPipelineDescriptor, ShaderStages, StoreOp,
    TextureFormat, TextureView, VertexAttribute, VertexBufferLayout, VertexFormat, VertexState,
    VertexStepMode, include_wgsl,
};

use crate::scene_prepare::RPrefab;

pub struct ScenePipeline {
    pipeline: RenderPipeline,
    bind_group: BindGroup,
    prefabs: HashMap<Resource, RPrefab>,
    prefabs_needed: HashSet<Resource>,
}

impl ScenePipeline {
    pub fn new(device: &Device, format: TextureFormat) -> Self {
        let module = device.create_shader_module(include_wgsl!("shader.wgsl"));

        let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            entries: &[],
            label: None,
        });
        let bind_group = device.create_bind_group(&BindGroupDescriptor {
            label: None,
            layout: &bind_group_layout,
            entries: &[],
        });
        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
            label: None,
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[],
        });
        let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
            label: None,
            layout: Some(&pipeline_layout),
            fragment: Some(FragmentState {
                module: &module,
                entry_point: Some("fs_main"),
                targets: &[Some(ColorTargetState {
                    blend: Some(BlendState::PREMULTIPLIED_ALPHA_BLENDING),
                    format,
                    write_mask: ColorWrites::all(),
                })],
                compilation_options: PipelineCompilationOptions::default(),
            }),
            vertex: VertexState {
                module: &module,
                entry_point: Some("vs_main"),
                buffers: &[VertexBufferLayout {
                    step_mode: VertexStepMode::Vertex,
                    array_stride: 2 * 4 * 3,
                    attributes: &[
                        VertexAttribute {
                            format: VertexFormat::Float32x3,
                            offset: 0,
                            shader_location: 0,
                        },
                        VertexAttribute {
                            format: VertexFormat::Float32x3,
                            offset: 3 * 4,
                            shader_location: 1,
                        },
                    ],
                }],
                compilation_options: PipelineCompilationOptions::default(),
            },
            primitive: PrimitiveState {
                topology: PrimitiveTopology::TriangleList,
                front_face: FrontFace::Ccw,
                cull_mode: None, //Some(Face::Back),
                polygon_mode: PolygonMode::Fill,
                ..Default::default()
            },
            depth_stencil: Default::default(),
            multisample: MultisampleState::default(),
            multiview: None,
            cache: None,
        });
        Self {
            bind_group,
            pipeline,
            prefabs_needed: HashSet::new(),
            prefabs: HashMap::new(),
        }
    }
    pub fn draw(&mut self, commands: &mut CommandEncoder, target: &TextureView, scene: &SceneTree) {
        let mut rpass = commands.begin_render_pass(&RenderPassDescriptor {
            label: None,
            color_attachments: &[Some(RenderPassColorAttachment {
                view: target,
                resolve_target: None,
                ops: Operations {
                    store: StoreOp::Store,
                    load: LoadOp::Clear(Color {
                        r: 0.1,
                        g: 0.1,
                        b: 0.1,
                        a: 1.,
                    }),
                },
            })],
            ..Default::default()
        });
        rpass.set_bind_group(0, &self.bind_group, &[]);
        rpass.set_pipeline(&self.pipeline);

        for ob in scene.objects.values() {
            if let Some(prefab) = self.prefabs.get(&ob.res) {
                for (affine, part) in &prefab.0 {
                    let affine = affine.to_cols_array().map(|v| v.to_le_bytes());
                    rpass.set_push_constants(ShaderStages::VERTEX, 0, affine.as_flattened());
                    rpass.set_index_buffer(part.index.slice(..), IndexFormat::Uint16);
                    rpass.set_vertex_buffer(0, part.positions.slice(..));
                    rpass.set_vertex_buffer(1, part.normals.slice(..));
                    rpass.draw_indexed(0..part.index_count, 0, 0..1);
                }
            } else {
                self.prefabs_needed.insert(ob.res);
            }
        }
    }
}