use glam::{Mat2, Vec2};
use std::{borrow::Cow, mem, num::NonZeroU32, slice, sync::Arc};
use wgpu::*;
use winit::{dpi::PhysicalSize, event_loop::OwnedDisplayHandle, window::Window};

pub mod command;
mod font;
pub mod image;
mod texture;

use command::RenderCommand;
use font::CharAtlas;
use image::ImageId;
use texture::TextureWrapper;

pub const INSTANCE_MAX_COUNT: u64 = 128;
const IMAGE_SLOT_COUNT: u32 = 5;

#[repr(C)]
struct InstanceData {
    // 拡縮・回転行列 NDCを考慮していること
    transform: Mat2,
    // 中央原点 左→右[-1, 1] 下→上[-1, 1]
    translation: Vec2,
    uv_offset: Vec2,
    uv_size: Vec2,
    tex_id: u32,
    _pad: u32,
}

pub enum AcquireSurfaceTextureResult {
    Ok(SurfaceTexture),
    ShouldSkip,
    Lost,
}

pub struct Renderer {
    device: Device,
    queue: Queue,
    surface: Surface<'static>,
    surface_format: TextureFormat,
    view_format: TextureFormat,
    present_mode: PresentMode,
    pipeline: RenderPipeline,

    instances: Buffer,
    trans_bind_group: BindGroup,
    char_atlas: CharAtlas,
    images: Vec<TextureWrapper>,
    fixed_bind_group: BindGroup,
}

impl Renderer {
    pub async fn new(
        display: Box<OwnedDisplayHandle>,
        window: Arc<Window>,
    ) -> Result<Self, String> {
        let size = window.inner_size();
        let scale_factor = window.scale_factor() as f32;

        let instance = Instance::new(InstanceDescriptor::new_with_display_handle(display));
        let surface = instance.create_surface(window).map_err(|e| e.to_string())?;
        let adapter = instance
            .request_adapter(&RequestAdapterOptions {
                compatible_surface: Some(&surface),
                ..Default::default()
            })
            .await
            .map_err(|e| e.to_string())?;
        let (device, queue) = adapter
            .request_device(&DeviceDescriptor {
                required_features: Features::TEXTURE_BINDING_ARRAY
                    | Features::SAMPLED_TEXTURE_AND_STORAGE_BUFFER_ARRAY_NON_UNIFORM_INDEXING,
                required_limits: Limits {
                    max_binding_array_elements_per_shader_stage: IMAGE_SLOT_COUNT,
                    ..Default::default()
                },
                ..Default::default()
            })
            .await
            .map_err(|e| e.to_string())?;

        let surface_cap = surface.get_capabilities(&adapter);
        if surface_cap.formats.is_empty() {
            return Err("no surface format found".to_string());
        }

        let (surface_format, view_format) = decide_surface_and_view_format(&adapter, &surface_cap);
        let present_mode = decide_presentation_mode(&surface_cap);

        let shader = device.create_shader_module(ShaderModuleDescriptor {
            label: None,
            source: ShaderSource::Wgsl(Cow::Borrowed(include_str!("shader.wgsl"))),
        });

        let trans_bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            label: None,
            entries: &[BindGroupLayoutEntry {
                binding: 0,
                visibility: ShaderStages::VERTEX,
                ty: BindingType::Buffer {
                    ty: BufferBindingType::Storage { read_only: true },
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        });
        let fixed_bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            label: None,
            entries: &[
                BindGroupLayoutEntry {
                    binding: 0,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Sampler(SamplerBindingType::Filtering),
                    count: None,
                },
                BindGroupLayoutEntry {
                    binding: 1,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Texture {
                        sample_type: TextureSampleType::Float { filterable: true },
                        view_dimension: TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                BindGroupLayoutEntry {
                    binding: 2,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Texture {
                        sample_type: TextureSampleType::Float { filterable: true },
                        view_dimension: TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: NonZeroU32::new(IMAGE_SLOT_COUNT),
                },
            ],
        });
        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
            label: None,
            bind_group_layouts: &[
                Some(&trans_bind_group_layout),
                Some(&fixed_bind_group_layout),
            ],
            immediate_size: 0,
        });
        let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
            label: None,
            layout: Some(&pipeline_layout),
            vertex: VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[],
                compilation_options: Default::default(),
            },
            fragment: Some(FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                compilation_options: Default::default(),
                targets: &[Some(ColorTargetState {
                    format: view_format,
                    blend: Some(BlendState::ALPHA_BLENDING),
                    write_mask: ColorWrites::ALL,
                })],
            }),
            primitive: PrimitiveState {
                topology: PrimitiveTopology::TriangleStrip,
                cull_mode: None,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: MultisampleState::default(),
            multiview_mask: None,
            cache: None,
        });

        // NOTE: インスタンスバッファは固定長で使いまわすので、ここで1個だけ作成すれば良い。
        let instances = device.create_buffer(&BufferDescriptor {
            label: None,
            size: mem::size_of::<InstanceData>() as u64 * INSTANCE_MAX_COUNT,
            usage: BufferUsages::STORAGE | BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let trans_bind_group = device.create_bind_group(&BindGroupDescriptor {
            label: None,
            layout: &trans_bind_group_layout,
            entries: &[BindGroupEntry {
                binding: 0,
                resource: BindingResource::Buffer(BufferBinding {
                    buffer: &instances,
                    offset: 0,
                    size: None,
                }),
            }],
        });

        // NOTE: 本ゲームではlinearなサンプラしか用意しないので、ここで1個だけ作成すれば良い。
        let sampler = device.create_sampler(&SamplerDescriptor {
            label: None,
            address_mode_u: AddressMode::ClampToEdge,
            address_mode_v: AddressMode::ClampToEdge,
            address_mode_w: AddressMode::ClampToEdge,
            mag_filter: FilterMode::Linear,
            min_filter: FilterMode::Linear,
            mipmap_filter: MipmapFilterMode::Linear,
            ..Default::default()
        });
        // NOTE: 文字テクスチャアトラスは使い回すので、ここで1個だけ作成すれば良い。
        let char_atlas = CharAtlas::new(scale_factor, &device);
        // NOTE: 画像テクスチャ配列は使いまわすので、ここで1個だけ作成すればいい。
        let images = image::create_textures(&device, &queue);
        let fixed_bind_group = device.create_bind_group(&BindGroupDescriptor {
            label: None,
            layout: &fixed_bind_group_layout,
            entries: &[
                BindGroupEntry {
                    binding: 0,
                    resource: BindingResource::Sampler(&sampler),
                },
                BindGroupEntry {
                    binding: 1,
                    resource: BindingResource::TextureView(&char_atlas.texture.view),
                },
                BindGroupEntry {
                    binding: 2,
                    resource: BindingResource::TextureViewArray(
                        &[&images[ImageId::Image as usize].view; IMAGE_SLOT_COUNT as usize],
                    ),
                },
            ],
        });

        let v = Self {
            device,
            queue,
            surface,
            surface_format,
            view_format,
            present_mode,
            pipeline,
            instances,
            trans_bind_group,
            char_atlas,
            images,
            fixed_bind_group,
        };

        v.configure_surface(size);

        Ok(v)
    }

    pub fn get_current_texture(&self, window: &Window) -> AcquireSurfaceTextureResult {
        match self.surface.get_current_texture() {
            CurrentSurfaceTexture::Success(v) => AcquireSurfaceTextureResult::Ok(v),
            CurrentSurfaceTexture::Suboptimal(v) => {
                // NOTE: 取得したテクスチャが生存した状態でサーフェスを再設定すると
                //       `PreviousOutputExists`エラーになるので、
                //       破棄した上で再設定し、描画をスキップする。
                drop(v);
                self.configure_surface(window.inner_size());
                AcquireSurfaceTextureResult::ShouldSkip
            }
            CurrentSurfaceTexture::Occluded => AcquireSurfaceTextureResult::ShouldSkip,
            CurrentSurfaceTexture::Timeout => AcquireSurfaceTextureResult::ShouldSkip,
            CurrentSurfaceTexture::Outdated => {
                self.configure_surface(window.inner_size());
                AcquireSurfaceTextureResult::ShouldSkip
            }
            CurrentSurfaceTexture::Lost => AcquireSurfaceTextureResult::Lost,
            // TODO: 一応対応する。
            CurrentSurfaceTexture::Validation => unreachable!(),
        }
    }

    pub fn render(
        &mut self,
        window: &Window,
        current_texture: SurfaceTexture,
        commands: &[RenderCommand],
    ) {
        let surface_texture_view = current_texture.texture.create_view(&TextureViewDescriptor {
            format: Some(self.view_format),
            ..Default::default()
        });

        let mut encoder = self.device.create_command_encoder(&Default::default());
        let mut render_pass = encoder.begin_render_pass(&RenderPassDescriptor {
            label: None,
            color_attachments: &[Some(RenderPassColorAttachment {
                view: &surface_texture_view,
                depth_slice: None,
                resolve_target: None,
                ops: Operations {
                    load: LoadOp::Clear(Color::BLACK),
                    store: StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });

        render_pass.set_pipeline(&self.pipeline);
        render_pass.set_bind_group(0, &self.trans_bind_group, &[]);
        render_pass.set_bind_group(1, &self.fixed_bind_group, &[]);

        let size = window.inner_size().to_logical::<f32>(window.scale_factor());
        let mut instances = Vec::new();
        for command in commands {
            command.build(
                size.width,
                size.height,
                window.scale_factor() as f32,
                &self.queue,
                &mut self.char_atlas,
                &self.images,
                &mut instances,
            );
        }
        instances.truncate(INSTANCE_MAX_COUNT as usize);
        let bytes = unsafe {
            slice::from_raw_parts(
                instances.as_ptr() as *const u8,
                mem::size_of_val(instances.as_slice()),
            )
        };
        self.queue.write_buffer(&self.instances, 0, bytes);
        render_pass.draw(0..4, 0..instances.len() as u32);

        drop(render_pass);

        self.queue.submit([encoder.finish()]);
        window.pre_present_notify();
        self.queue.present(current_texture);
    }

    /// サーフェスを設定するメソッド
    ///
    /// * size - 現在のクライアントサイズ
    ///
    /// クライアントサイズのリサイズがあったらこのメソッドを呼び出す。
    /// 幅・高さいずれかが`0`である場合は何もしない。
    pub fn configure_surface(&self, size: PhysicalSize<u32>) {
        if size.width == 0 || size.height == 0 {
            return;
        }

        let surface_config = SurfaceConfiguration {
            usage: TextureUsages::RENDER_ATTACHMENT,
            format: self.surface_format,
            color_space: SurfaceColorSpace::Auto,
            view_formats: if self.surface_format != self.view_format {
                vec![self.view_format]
            } else {
                vec![]
            },
            alpha_mode: CompositeAlphaMode::Auto,
            width: size.width,
            height: size.height,
            desired_maximum_frame_latency: 2,
            present_mode: self.present_mode,
        };
        self.surface.configure(&self.device, &surface_config);
    }

    /// Presentation modeがFIFOか確認するメソッド
    ///
    /// FIFOである場合はsleep+spinではなく垂直同期でスレッドを待機した方が良い。
    pub fn is_presentation_mode_fifo(&self) -> bool {
        self.present_mode == PresentMode::Fifo
    }
}

fn decide_surface_and_view_format(
    adapter: &Adapter,
    cap: &SurfaceCapabilities,
) -> (TextureFormat, TextureFormat) {
    // sRGBがサポートされているならsRGB
    if let Some(&v) = cap.formats.iter().find(|v| v.is_srgb()) {
        return (v, v);
    }

    let can_srgb_view = adapter
        .get_downlevel_capabilities()
        .flags
        .contains(DownlevelFlags::SURFACE_VIEW_FORMATS);
    let format = cap.formats[0];

    // sRGBがサポートされていない・かつテクスチャとビューのフォーマットが一致しなくて良いなら
    // テクスチャを`formats[0]`・ビューをsRGB
    //
    // NOTE: ビューのフォーマットに従って色変換が行われるので、
    //       テクスチャの実際のフォーマットが違うだけでsRGBと同様の結果になる。
    if can_srgb_view {
        (format, format.add_srgb_suffix())
    }
    // sRGBがサポートされていない・かつテクスチャとビューのフォーマットが一致する必要があるなら
    // テクスチャもビューも`formats[0]`
    //
    // NOTE: 色変換が行われないので中間色が暗くなる。
    //       シェーダで自力補正してやれば解決するが、まあ余力があったらね。
    else {
        (format, format)
    }
}

fn decide_presentation_mode(cap: &SurfaceCapabilities) -> PresentMode {
    // REVIEW: 高負荷なので素直にFIFOにして垂直同期で回す方が良いかもしれない。
    if cap.present_modes.contains(&PresentMode::Mailbox) {
        // 垂直同期を待つフレーム上書きモード
        PresentMode::Mailbox
    } else if cap.present_modes.contains(&PresentMode::Immediate) {
        // 垂直同期を待たないフレーム上書きモード
        PresentMode::Immediate
    } else {
        // 垂直同期を待つフレーム保守モード
        //
        // NOTE: 殆どすべての環境でサポートされていることを期待して良い。
        //       https://wgpu.rs/doc/wgpu/enum.PresentMode.html
        PresentMode::Fifo
    }
}