use glam::{EulerRot, Mat3, Mat4, Vec3, vec3};
use std::sync::Arc;
use weareshared::{packets::Resource, tree::SceneTree};
use wgpu::{
    BindGroupLayout, BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingType, BlendState,
    Color, ColorTargetState, ColorWrites, CommandEncoder, CompareFunction, DepthBiasState,
    DepthStencilState, Device, Extent3d, FragmentState, FrontFace, IndexFormat, LoadOp,
    MultisampleState, Operations, PipelineCompilationOptions, PipelineLayoutDescriptor,
    PolygonMode, PrimitiveState, PrimitiveTopology, PushConstantRange, RenderPassColorAttachment,
    RenderPassDepthStencilAttachment, RenderPassDescriptor, RenderPipeline,
    RenderPipelineDescriptor, SamplerBindingType, ShaderStages, StencilState, StoreOp,
    TextureDescriptor, TextureDimension, TextureFormat, TextureSampleType, TextureUsages,
    TextureView, TextureViewDescriptor, TextureViewDimension, VertexAttribute, VertexBufferLayout,
    VertexFormat, VertexState, VertexStepMode, include_wgsl,
};

use crate::scene_prepare::{DemandMap, RPrefab};

pub struct ScenePipeline {
    pipeline: RenderPipeline,
    depth: TextureView,
}

macro_rules! v_attr {
    ($($n:literal),*) => {
        [$(VertexBufferLayout {
            step_mode: VertexStepMode::Vertex,
            array_stride: 4,
            attributes: &[VertexAttribute {
                format: VertexFormat::Float32,
                offset: 0,
                shader_location: $n,
            }],
        }),*]
    };
}

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

        let depth = device.create_texture(&TextureDescriptor {
            label: None,
            size: Extent3d {
                height: 256,
                width: 256,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: TextureDimension::D2,
            format: TextureFormat::Depth32Float,
            usage: TextureUsages::RENDER_ATTACHMENT,
            view_formats: &[],
        });
        let depth = depth.create_view(&TextureViewDescriptor::default());

        let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            entries: &[
                BindGroupLayoutEntry {
                    binding: 0,
                    count: None,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Texture {
                        sample_type: TextureSampleType::Float { filterable: true },
                        view_dimension: TextureViewDimension::D2,
                        multisampled: false,
                    },
                },
                BindGroupLayoutEntry {
                    binding: 1,
                    count: None,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Sampler(SamplerBindingType::Filtering),
                },
            ],
            label: None,
        });
        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
            label: None,
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[PushConstantRange {
                range: 0..(4 * 4 * size_of::<f32>() as u32),
                stages: ShaderStages::VERTEX,
            }],
        });
        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: &v_attr!(0, 1, 2, 3, 4, 5, 6, 7),
                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: Some(DepthStencilState {
                depth_compare: CompareFunction::Greater,
                depth_write_enabled: true,
                format: TextureFormat::Depth32Float,
                bias: DepthBiasState::default(),
                stencil: StencilState::default(),
            }),
            multisample: MultisampleState::default(),
            multiview: None,
            cache: None,
        });
        (Self { pipeline, depth }, bind_group_layout)
    }

    pub fn resize(&mut self, device: &Device, width: u32, height: u32) {
        self.depth = device
            .create_texture(&TextureDescriptor {
                label: None,
                size: Extent3d {
                    height,
                    width,
                    depth_or_array_layers: 1,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: TextureDimension::D2,
                format: TextureFormat::Depth32Float,
                usage: TextureUsages::RENDER_ATTACHMENT,
                view_formats: &[],
            })
            .create_view(&TextureViewDescriptor::default());
    }

    pub fn draw(
        &mut self,
        commands: &mut CommandEncoder,
        target: &TextureView,
        scene: &SceneTree,
        prefabs: &mut DemandMap<Resource, Arc<RPrefab>>,
    ) {
        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.01,
                        g: 0.01,
                        b: 0.01,
                        a: 1.,
                    }),
                },
            })],
            depth_stencil_attachment: Some(RenderPassDepthStencilAttachment {
                view: &self.depth,
                depth_ops: Some(Operations {
                    load: LoadOp::Clear(0.),
                    store: StoreOp::Store,
                }),
                stencil_ops: None,
            }),
            ..Default::default()
        });

        let camera = Mat4::perspective_infinite_reverse_rh(1., 1., 0.1)
            * Mat4::look_at_rh(vec3(-5., 0., 0.), vec3(0., 0., 0.), Vec3::Y);

        for ob in scene.objects.values() {
            let prefab_projection = camera
                * Mat4::from_mat3(Mat3::from_euler(
                    EulerRot::YXZ,
                    ob.rot.x,
                    ob.rot.y,
                    ob.rot.z,
                ))
                * Mat4::from_translation(ob.pos.into());
            if let Some(prefab) = prefabs.try_get(ob.res) {
                for (affine, part) in &prefab.0 {
                    let part_projection = prefab_projection
                        * Mat4::from_mat3a(affine.matrix3)
                        * Mat4::from_translation(affine.translation.into());
                    let projection = part_projection.to_cols_array().map(|v| v.to_le_bytes());

                    rpass.set_bind_group(0, &*part.texture, &[]);
                    rpass.set_pipeline(&self.pipeline);
                    rpass.set_push_constants(ShaderStages::VERTEX, 0, projection.as_flattened());
                    rpass.set_index_buffer(part.index.slice(..), IndexFormat::Uint16);
                    rpass.set_vertex_buffer(0, part.position[0].slice(..));
                    rpass.set_vertex_buffer(1, part.position[1].slice(..));
                    rpass.set_vertex_buffer(2, part.position[2].slice(..));
                    rpass.set_vertex_buffer(3, part.normal[0].slice(..));
                    rpass.set_vertex_buffer(4, part.normal[1].slice(..));
                    rpass.set_vertex_buffer(5, part.normal[2].slice(..));
                    rpass.set_vertex_buffer(6, part.texcoord[0].slice(..));
                    rpass.set_vertex_buffer(7, part.texcoord[1].slice(..));
                    rpass.draw_indexed(0..part.index_count, 0, 0..1);
                }
            }
        }
    }
}