/*
    wearechat - generic multiplayer game with voip
    Copyright (C) 2025 metamuffin

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, version 3 of the License only.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
use super::{GraphicsConfig, ScenePreparer, TextureIdentityKind};
use anyhow::Result;
use image::ImageReader;
use log::debug;
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use std::{
    io::Cursor,
    sync::{Arc, atomic::AtomicUsize},
    time::Instant,
};
use wgpu::{
    AddressMode, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout, BindingResource,
    Color, ColorTargetState, ColorWrites, CommandEncoderDescriptor, Device, Extent3d, FilterMode,
    ImageDataLayout, LoadOp, Operations, Queue, RenderPassColorAttachment, RenderPassDescriptor,
    RenderPipeline, SamplerDescriptor, StoreOp, Texture, TextureAspect, TextureDescriptor,
    TextureDimension, TextureFormat, TextureUsages, TextureViewDescriptor, include_wgsl,
};

pub struct MipGenerationPipeline {
    pipeline: RenderPipeline,
}
impl MipGenerationPipeline {
    pub fn load(device: &Device, format: TextureFormat) -> Self {
        let shader = device.create_shader_module(include_wgsl!("../shaders/texture_copy.wgsl"));
        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("mip generator"),
            layout: None,
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                compilation_options: Default::default(),
                buffers: &[],
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                compilation_options: Default::default(),
                targets: &[Some(ColorTargetState {
                    format,
                    blend: None,
                    write_mask: ColorWrites::ALL,
                })],
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });
        Self { pipeline }
    }
}

impl ScenePreparer {
    pub fn update_textures(&self, num_done: &mut usize) -> Result<()> {
        for format in self.mip_generation_pipelines.needed() {
            self.mip_generation_pipelines.insert(
                format,
                Arc::new(MipGenerationPipeline::load(&self.device, format)),
                0,
            );
            *num_done += 1
        }

        for kind in self.placeholder_textures.needed() {
            let (linear, color) = match kind {
                TextureIdentityKind::Normal => (true, [128, 128, 255, 255]),
                TextureIdentityKind::Multiply => (false, [255, 255, 255, 255]),
            };
            let tex_bg = create_texture(
                &self.device,
                &self.queue,
                &self.layouts.texture,
                &color,
                1,
                1,
                if linear {
                    TextureFormat::Rgba8Unorm
                } else {
                    TextureFormat::Rgba8UnormSrgb
                },
                None,
                &self.config.read().unwrap().clone(),
            );
            self.placeholder_textures.insert(kind, tex_bg, 4);
            *num_done += 1;
        }

        let numdone = AtomicUsize::new(0);
        self.textures
            .needed()
            .into_par_iter()
            .try_for_each(|spec| {
                let start = Instant::now();
                let format = if spec.linear {
                    TextureFormat::Rgba8Unorm
                } else {
                    TextureFormat::Rgba8UnormSrgb
                };
                if let Some(mipgen) = self.mip_generation_pipelines.try_get(format) {
                    if let Some(buf) = self.downloader.try_get(spec.data.clone())? {
                        let image = ImageReader::new(Cursor::new(buf.0)).with_guessed_format()?;
                        let image = image.decode()?;
                        let dims = (image.width(), image.height());
                        let image = image.into_rgba8();
                        let image = image.into_vec();
                        let tex_bg = create_texture(
                            &self.device,
                            &self.queue,
                            &self.layouts.texture,
                            &image,
                            dims.0,
                            dims.1,
                            format,
                            Some(&mipgen),
                            &self.config.read().unwrap().clone(),
                        );
                        self.textures.insert(spec, tex_bg, image.len());
                        debug!(
                            "texture created (res={}x{}, took {:?})",
                            dims.0,
                            dims.1,
                            start.elapsed()
                        );
                        numdone.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    }
                }
                Ok::<(), anyhow::Error>(())
            })?;
        *num_done += numdone.load(std::sync::atomic::Ordering::Relaxed);
        Ok(())
    }
}

fn create_texture(
    device: &Device,
    queue: &Queue,
    bgl: &BindGroupLayout,
    data: &[u8],
    width: u32,
    height: u32,
    format: TextureFormat,
    mipgen: Option<&MipGenerationPipeline>,
    config: &GraphicsConfig,
) -> (Arc<Texture>, Arc<BindGroup>) {
    let mip_level_count = (width.ilog2().max(4) - 3).min(config.max_mip_count);

    let extent = Extent3d {
        depth_or_array_layers: 1,
        width,
        height,
    };
    let texture = device.create_texture(&TextureDescriptor {
        label: None,
        size: extent,
        mip_level_count,
        sample_count: 1,
        dimension: TextureDimension::D2,
        format,
        usage: TextureUsages::TEXTURE_BINDING
            | TextureUsages::COPY_DST
            | TextureUsages::RENDER_ATTACHMENT,
        view_formats: &[],
    });
    let textureview = texture.create_view(&TextureViewDescriptor::default());
    let sampler = device.create_sampler(&SamplerDescriptor {
        address_mode_u: AddressMode::Repeat,
        address_mode_v: AddressMode::Repeat,
        mag_filter: FilterMode::Linear,
        min_filter: FilterMode::Linear,
        mipmap_filter: FilterMode::Linear,
        anisotropy_clamp: config.max_anisotropy,
        ..Default::default()
    });
    let bind_group = device.create_bind_group(&BindGroupDescriptor {
        label: None,
        layout: &bgl,
        entries: &[
            BindGroupEntry {
                binding: 0,
                resource: BindingResource::TextureView(&textureview),
            },
            BindGroupEntry {
                binding: 1,
                resource: BindingResource::Sampler(&sampler),
            },
        ],
    });

    let level_views = (0..mip_level_count)
        .map(|mip| {
            texture.create_view(&TextureViewDescriptor {
                label: Some("mip generation level view"),
                format: None,
                dimension: None,
                aspect: TextureAspect::All,
                base_mip_level: mip,
                mip_level_count: Some(1),
                base_array_layer: 0,
                array_layer_count: None,
            })
        })
        .collect::<Vec<_>>();

    let mut encoder = device.create_command_encoder(&CommandEncoderDescriptor { label: None });

    // TODO why does copy_buffer_to_texture have more restrictive alignment requirements?!
    // let upload_buffer = device.create_buffer_init(&BufferInitDescriptor {
    //     label: Some("texture upload"),
    //     contents: data,
    //     usage: BufferUsages::COPY_DST | BufferUsages::COPY_SRC,
    // });
    // encoder.copy_buffer_to_texture(
    //     ImageCopyBuffer {
    //         buffer: &upload_buffer,
    //         layout: ImageDataLayout {
    //             offset: 0,
    //             bytes_per_row: Some(width * 4),
    //             rows_per_image: None,
    //         },
    //     },
    //     texture.as_image_copy(),
    //     extent,
    // );

    queue.write_texture(
        texture.as_image_copy(),
        data,
        ImageDataLayout {
            bytes_per_row: Some(width * 4),
            rows_per_image: None,
            offset: 0,
        },
        extent,
    );

    for level in 1..mip_level_count {
        let mip_pipeline = &mipgen.unwrap().pipeline;
        let source_view = &level_views[level as usize - 1];
        let target_view = &level_views[level as usize];
        let mip_bind_group = device.create_bind_group(&BindGroupDescriptor {
            layout: &mip_pipeline.get_bind_group_layout(0),
            entries: &[
                BindGroupEntry {
                    binding: 0,
                    resource: BindingResource::TextureView(source_view),
                },
                BindGroupEntry {
                    binding: 1,
                    resource: BindingResource::Sampler(&sampler),
                },
            ],
            label: None,
        });
        let mut rpass = encoder.begin_render_pass(&RenderPassDescriptor {
            label: None,
            color_attachments: &[Some(RenderPassColorAttachment {
                view: target_view,
                resolve_target: None,
                ops: Operations {
                    load: LoadOp::Clear(Color::WHITE),
                    store: StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
        });

        rpass.set_pipeline(&mip_pipeline);
        rpass.set_bind_group(0, &mip_bind_group, &[]);
        rpass.draw(0..3, 0..1);
    }

    queue.submit(Some(encoder.finish()));

    (Arc::new(texture), Arc::new(bind_group))
}