use std::rc::Rc;

use exr::prelude::{ReadChannels, ReadLayers};
use glam::Vec2;
use nalgebra::{vector, DMatrix};
use rapier3d::{
    math::Isometry,
    prelude::{ColliderBuilder, RigidBodyBuilder},
};
use wesl::Wesl;

use crate::{
    components::{MeshComponent, PhysicsComponent},
    entity::EntityHandle,
    mesh::{InstanceData, InstanceRaw, Mesh, Vertex},
    physics::PhysicsManager,
    render,
    world::{Transform, World},
};

pub struct TerrainConfig
{
    pub gltf_path: String,
    pub heightmap_path: String,
    pub size: Vec2,
}

impl TerrainConfig
{
    pub fn new(gltf_path: &str, heightmap_path: &str, size: Vec2) -> Self
    {
        Self {
            gltf_path: gltf_path.to_string(),
            heightmap_path: heightmap_path.to_string(),
            size,
        }
    }

    pub fn default() -> Self
    {
        Self {
            gltf_path: "meshes/terrain.gltf".to_string(),
            heightmap_path: "textures/terrain_heightmap.exr".to_string(),
            size: Vec2::new(1000.0, 1000.0),
        }
    }
}

pub struct Terrain;

impl Terrain
{
    pub fn spawn(
        world: &mut World,
        mesh_data: Vec<(Mesh, Vec<InstanceData>)>,
        config: &TerrainConfig,
    ) -> anyhow::Result<EntityHandle>
    {
        let transform = Transform::IDENTITY;
        let mut first_entity = None;
        let mut physics_added = false;

        for (mesh, instances) in mesh_data
        {
            let entity = world.spawn();

            if first_entity.is_none()
            {
                first_entity = Some(entity);
            }

            if instances.is_empty()
            {
                world.transforms.insert(entity, transform);
                world.meshes.insert(
                    entity,
                    MeshComponent {
                        mesh: Rc::new(mesh),
                        pipeline: render::Pipeline::Terrain,
                        instance_buffer: None,
                        num_instances: 1,
                    },
                );

                if !physics_added
                {
                    let heights = Self::load_heightfield_from_exr(&config.heightmap_path)?;

                    let height_scale = 1.0;
                    let scale = vector![config.size.x, height_scale, config.size.y];

                    let body = RigidBodyBuilder::fixed()
                        .translation(transform.get_position().into())
                        .build();

                    let rigidbody_handle = PhysicsManager::add_rigidbody(body);

                    let collider = ColliderBuilder::heightfield(heights.clone(), scale).build();

                    let collider_handle =
                        PhysicsManager::add_collider(collider, Some(rigidbody_handle));

                    PhysicsManager::set_heightfield_data(
                        heights,
                        scale,
                        transform.get_position().into(),
                    );

                    world.physics.insert(
                        entity,
                        PhysicsComponent {
                            rigidbody: rigidbody_handle,
                            collider: Some(collider_handle),
                        },
                    );

                    physics_added = true;
                }
            }
            else
            {
                let num_instances = instances.len();

                let instance_raw: Vec<InstanceRaw> = instances.iter().map(|i| i.to_raw()).collect();

                let instance_buffer = render::with_device(|device| {
                    use wgpu::util::DeviceExt;
                    device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: Some("Tree Instance Buffer"),
                        contents: bytemuck::cast_slice(&instance_raw),
                        usage: wgpu::BufferUsages::VERTEX,
                    })
                });

                world.transforms.insert(entity, Transform::IDENTITY);
                world.meshes.insert(
                    entity,
                    MeshComponent {
                        mesh: Rc::new(mesh),
                        pipeline: render::Pipeline::Environment,
                        instance_buffer: Some(instance_buffer),
                        num_instances: num_instances as u32,
                    },
                );
            }
        }

        first_entity.ok_or_else(|| anyhow::anyhow!("No meshes found in glTF file"))
    }

    fn load_heightfield_from_exr(path: &str) -> anyhow::Result<DMatrix<f32>>
    {
        let image = exr::prelude::read()
            .no_deep_data()
            .largest_resolution_level()
            .all_channels()
            .all_layers()
            .all_attributes()
            .from_file(path)?;

        let layer = &image.layer_data[0];
        let channel = &layer.channel_data.list[0];

        let width = layer.size.width();
        let height = layer.size.height();

        let heights: Vec<f32> = channel.sample_data.values_as_f32().collect();

        Ok(DMatrix::from_row_slice(height, width, &heights))
    }
}

pub fn create_terrain_render_pipeline(
    device: &wgpu::Device,
    config: &wgpu::SurfaceConfiguration,
    bind_group_layout: &wgpu::BindGroupLayout,
) -> wgpu::RenderPipeline
{
    let compiler = Wesl::new("src/shaders");
    let shader_source = compiler
        .compile(&"package::terrain".parse().unwrap())
        .inspect_err(|e| eprintln!("WESL error: {e}"))
        .unwrap()
        .to_string();

    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("Terrain Shader"),
        source: wgpu::ShaderSource::Wgsl(shader_source.into()),
    });

    let render_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: Some("Terrain Render Pipeline Layout"),
        bind_group_layouts: &[bind_group_layout],
        push_constant_ranges: &[],
    });

    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: Some("Terrain Render Pipeline"),
        layout: Some(&render_pipeline_layout),
        vertex: wgpu::VertexState {
            module: &shader,
            entry_point: Some("vs_main"),
            buffers: &[Vertex::desc(), InstanceRaw::desc()],
            compilation_options: Default::default(),
        },
        fragment: Some(wgpu::FragmentState {
            module: &shader,
            entry_point: Some("fs_main"),
            targets: &[Some(wgpu::ColorTargetState {
                format: config.format,
                blend: Some(wgpu::BlendState::REPLACE),
                write_mask: wgpu::ColorWrites::ALL,
            })],
            compilation_options: Default::default(),
        }),
        primitive: wgpu::PrimitiveState {
            topology: wgpu::PrimitiveTopology::TriangleList,
            strip_index_format: None,
            front_face: wgpu::FrontFace::Ccw,
            cull_mode: Some(wgpu::Face::Back),
            polygon_mode: wgpu::PolygonMode::Fill,
            unclipped_depth: false,
            conservative: false,
        },
        depth_stencil: Some(wgpu::DepthStencilState {
            format: wgpu::TextureFormat::Depth32Float,
            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,
    })
}
