use std::collections::HashSet;
use std::f32::consts::PI;
use std::sync::Arc;
use std::time::Instant;

use bytemuck::bytes_of;
use egui_wgpu::ScreenDescriptor;
use gpu_shared::RayMarchingPushConstants;
use wgpu::Backends;
use wgpu::Device;
use wgpu::ExperimentalFeatures;
use wgpu::Extent3d;
use wgpu::FeaturesWGPU;
use wgpu::FeaturesWebGPU;
use wgpu::PushConstantRange;
use wgpu::RenderPipeline;
use wgpu::ShaderStages;
use wgpu::TextureDescriptor;
use wgpu::TextureFormat;
use wgpu::TextureUsages;
use wgpu::TextureView;
use wgpu::TextureViewDescriptor;
use wgpu::include_spirv;
use winit::event::MouseScrollDelta;
use winit::event::WindowEvent;
use winit::keyboard::KeyCode;
use winit::window::Window;

use crate::egui_renderer::EguiState;

pub struct State
{
    window: Arc<Window>,
    device: wgpu::Device,
    queue: wgpu::Queue,
    size: winit::dpi::PhysicalSize<u32>,
    surface: wgpu::Surface<'static>,
    surface_format: wgpu::TextureFormat,
    egui_state: EguiState,
    camera: Camera,

    // Input
    pressed_set: HashSet<KeyCode>,

    // Frame time
    last_frame: Instant,

    depth_buffer: TextureView,
    test_pipeline: RenderPipeline,
}

#[derive(Clone, Copy)]
pub struct Camera
{
    pub eye: glam::Vec3,
    pub up: glam::Vec3,
    pub aspect: f32,
    pub fovy: f32,
    pub znear: f32,
    pub zfar: f32,

    pub radius: f32,

    pub yaw: f32,
    pub pitch: f32,
    pub rotation_speed: f32,
    pub speed: f32,
}

impl State
{
    pub async fn new(window: Arc<Window>) -> State
    {
        // let shaders = ShaderModuleDescriptorPassthrough {
        //     label: Some("shaders_modules"),
        //     spirv: Some(make_spirv_raw(include_bytes!(env!("shaders.spv")))),
        //
        //     entry_point: "".to_owned(),
        //     // This is unused for SPIR-V
        //     num_workgroups: (0, 0, 0),
        //     runtime_checks: ShaderRuntimeChecks::unchecked(),
        //     dxil: None,
        //     msl: None,
        //     hlsl: None,
        //     glsl: None,
        //     wgsl: None,
        // };
        let shaders = include_spirv!(env!("shaders.spv"));
        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
            backends: Backends::VULKAN,
            ..wgpu::InstanceDescriptor::default()
        });
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions::default())
            .await
            .unwrap();
        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                required_features: wgpu::Features {
                    features_wgpu: FeaturesWGPU::PUSH_CONSTANTS
                        | FeaturesWGPU::BUFFER_BINDING_ARRAY
                        | FeaturesWGPU::TEXTURE_BINDING_ARRAY
                        | FeaturesWGPU::STORAGE_RESOURCE_BINDING_ARRAY
                        | FeaturesWGPU::MULTIVIEW,
                    //| FeaturesWGPU::EXPERIMENTAL_PASSTHROUGH_SHADERS,
                    features_webgpu: FeaturesWebGPU::empty(),
                },
                experimental_features: unsafe { ExperimentalFeatures::enabled() },
                required_limits: wgpu::Limits {
                    max_push_constant_size: size_of::<RayMarchingPushConstants>() as u32,
                    max_binding_array_elements_per_shader_stage: 4,
                    max_color_attachment_bytes_per_sample: 48,
                    ..wgpu::Limits::downlevel_defaults()
                },
                ..Default::default()
            })
            .await
            .unwrap();

        let size = window.inner_size();

        let surface = instance.create_surface(window.clone()).unwrap();
        let cap = surface.get_capabilities(&adapter);
        let surface_format = cap.formats[0];

        let shader_module = device.create_shader_module(shaders);
        let test_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("chunk_pipeline"),
            layout: Some(
                &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: None,
                    bind_group_layouts: &[],
                    push_constant_ranges: &[PushConstantRange {
                        stages: ShaderStages::VERTEX,
                        range: 0..(size_of::<RayMarchingPushConstants>() as u32),
                    }],
                }),
            ),
            vertex: wgpu::VertexState {
                module: &shader_module,
                entry_point: Some("main_vs"),
                buffers: &[],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader_module,
                entry_point: Some("main_fs"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_format,
                    blend: None,
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                ..Default::default()
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: TextureFormat::Depth24PlusStencil8,
                depth_write_enabled: true,
                depth_compare: wgpu::CompareFunction::Less,
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            }),
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,

            cache: None,
        });

        let state = State {
            egui_state: EguiState::new(&device, surface_format, &window),

            test_pipeline,

            window,
            queue,
            size,
            surface,
            surface_format,
            camera: Camera {
                eye: glam::Vec3::new(0., 0., 0.),
                up: glam::Vec3::ZERO.with_y(1.),
                aspect: size.width as f32 / size.height as f32,
                fovy: 90.,
                znear: 0.001,
                zfar: 1000.,
                radius: 2.,
                yaw: 1.,
                pitch: 0.,
                rotation_speed: 0.005,
                speed: 0.005,
            },

            pressed_set: HashSet::new(),
            last_frame: Instant::now(),

            depth_buffer: create_gbuffer(
                &device,
                size.width,
                size.height,
                TextureFormat::Depth24PlusStencil8,
                TextureUsages::RENDER_ATTACHMENT,
            ),
            device,
        };

        // Configure surface for the first time
        state.configure_surface();

        state
    }

    pub fn get_window(&self) -> &Window
    {
        &self.window
    }

    pub fn configure_surface(&self)
    {
        let surface_config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: self.surface_format,
            // Request compatibility with the sRGB-format texture view we‘re going to create later.
            view_formats: vec![self.surface_format.add_srgb_suffix()],
            alpha_mode: wgpu::CompositeAlphaMode::Auto,
            width: self.size.width,
            height: self.size.height,
            desired_maximum_frame_latency: 2,
            present_mode: wgpu::PresentMode::AutoVsync,
        };
        self.surface.configure(&self.device, &surface_config);
    }

    pub fn resize(&mut self, new_size: winit::dpi::PhysicalSize<u32>)
    {
        self.size = new_size;

        self.camera.aspect = new_size.width as f32 / new_size.height as f32;
        self.depth_buffer = create_gbuffer(
            &self.device,
            new_size.width,
            new_size.height,
            TextureFormat::Depth24PlusStencil8,
            TextureUsages::RENDER_ATTACHMENT,
        );
        self.configure_surface();
    }

    pub fn handle_event(&mut self, event: &WindowEvent)
    {
        self.egui_state.handle_event(&self.window, event);

        if let WindowEvent::KeyboardInput { event, .. } = event
        {
            match (event.state, event.physical_key)
            {
                (winit::event::ElementState::Pressed, winit::keyboard::PhysicalKey::Code(c)) =>
                {
                    self.pressed_set.insert(c);
                }
                (winit::event::ElementState::Released, winit::keyboard::PhysicalKey::Code(c)) =>
                {
                    self.pressed_set.remove(&c);
                }
                _ =>
                {}
            }
        }
    }

    pub fn render(&mut self)
    {
        // Update camera
        self.update_camera();
        // Create texture view
        let surface_texture = self
            .surface
            .get_current_texture()
            .expect("failed to acquire next swapchain texture");
        let surface_view = surface_texture
            .texture
            .create_view(&wgpu::TextureViewDescriptor {
                // Without add_srgb_suffix() the image we will be working with
                // might not be "gamma correct".
                format: Some(self.surface_format.add_srgb_suffix()),
                ..Default::default()
            });
        let mut encoder = self.device.create_command_encoder(&Default::default());

        {
            let mut hm_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: None,
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &surface_view,
                    depth_slice: None,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                    view: &self.depth_buffer,
                    depth_ops: Some(wgpu::Operations {
                        load: wgpu::LoadOp::Clear(1.),
                        store: wgpu::StoreOp::Discard,
                    }),
                    stencil_ops: None,
                }),
                timestamp_writes: None,
                occlusion_query_set: None,
            });

            hm_pass.set_pipeline(&self.test_pipeline);
            hm_pass.set_push_constants(
                ShaderStages::VERTEX,
                0,
                bytes_of(&gpu_shared::RayMarchingPushConstants {
                    view_projection: self.camera.view_proj(),
                    transform: glam::Mat4::IDENTITY,
                    eye_position: self.camera.eye,
                    root_color: glam::Vec4::ZERO,
                    root_subdivided: 0,
                }),
            );
            hm_pass.draw(0..36, 0..1);
        }

        // Egui Pass
        {
            let screen_descriptor = ScreenDescriptor {
                size_in_pixels: [self.size.width, self.size.height],
                pixels_per_point: 1.,
            };
            self.egui_state.begin_frame(&self.window);

            egui::Window::new("Hello Window").resizable(true).show(
                self.egui_state.context(),
                |ui| {
                    ui.label(format!(
                        "Frame time: {:.2} ms",
                        (Instant::now() - self.last_frame).as_secs_f32() * 1000.
                    ));
                    ui.label(format!(
                        "Frame rate: {:.2} fps",
                        1. / (Instant::now() - self.last_frame).as_secs_f32()
                    ));
                    ui.label(format!(
                        "Camera pos {:.1},{:.1},{:.1}",
                        self.camera.eye.x, self.camera.eye.y, self.camera.eye.z,
                    ))
                },
            );

            self.egui_state.end_frame_and_draw(
                &self.device,
                &self.queue,
                &mut encoder,
                &self.window,
                &surface_view,
                screen_descriptor,
            );
        }

        // Submit the command in the queue to execute
        self.queue.submit([encoder.finish()]);
        self.window.pre_present_notify();
        surface_texture.present();

        self.last_frame = Instant::now();
    }

    fn update_camera(&mut self)
    {
        let mut movement = glam::Vec3::new(0., 0., 0.);
        if self.pressed_set.contains(&KeyCode::KeyW)
        {
            movement.z += self.camera.speed;
        }
        if self.pressed_set.contains(&KeyCode::KeyS)
        {
            movement.z -= self.camera.speed;
        }

        // Left rigth
        if self.pressed_set.contains(&KeyCode::KeyA)
        {
            movement.x += self.camera.speed;
        }
        if self.pressed_set.contains(&KeyCode::KeyD)
        {
            movement.x -= self.camera.speed;
        }

        let rot_movement = glam::Mat3::from_rotation_y(-self.camera.yaw)
            * glam::Mat3::from_rotation_x(-self.camera.pitch)
            * movement;
        self.camera.eye -= rot_movement;
    }

    pub fn cursor_moved(&mut self, x: f32, y: f32)
    {
        const SENSIBILITY: f32 = 0.0004;
        let position = glam::Vec2::new(x, y);
        let offset = position * SENSIBILITY;

        self.camera.yaw += offset.x;
        self.camera.pitch += offset.y;
    }

    pub fn mouse_wheel(&mut self, delta: MouseScrollDelta)
    {
        if let MouseScrollDelta::LineDelta(_, y) = delta
        {
            self.camera.speed += y * (self.camera.speed * 0.05);
        }

        //self.speed += delta.
    }
}

pub fn create_gbuffer(
    device: &Device,
    width: u32,
    height: u32,
    format: TextureFormat,
    usage: TextureUsages,
) -> TextureView
{
    device
        .create_texture(&TextureDescriptor {
            label: Some("gbuffer"),
            size: Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format,
            usage,
            view_formats: &[format],
        })
        .create_view(&TextureViewDescriptor::default())
}

// #[derive(AsStd430)]
// pub struct RayMarchingPushConstants
// {
//     view_projection: Matrix4<f32>,
//     transform: Matrix4<f32>,
//     eye_position: Vector3<f32>,
//     root_color: Vector4<f32>,
//     root_subdivided: u32,
// }

#[rustfmt::skip]
pub const OPENGL_TO_WGPU_MATRIX: glam::Mat4 = glam::Mat4::from_cols(
    glam::Vec4::new(1.0, 0.0, 0.0, 0.0),
    glam::Vec4::new(0.0, 1.0, 0.0, 0.0),
    glam::Vec4::new(0.0, 0.0, 0.5, 0.0),
    glam::Vec4::new(0.0, 0.0, 0.5, 1.0),
);

impl Camera
{
    pub fn view_proj(&self) -> glam::Mat4
    {
        let view = glam::Mat4::from_translation(self.eye)
            * glam::Mat4::from_rotation_y(-self.yaw)
            * glam::Mat4::from_rotation_x(-self.pitch);
        let proj =
            glam::Mat4::perspective_rh(PI * self.fovy / 180., self.aspect, self.znear, self.zfar);
        OPENGL_TO_WGPU_MATRIX * proj * view.inverse()
    }
}
