use std::cell::OnceCell;
use std::rc::Rc;

use glam::{Mat4, Vec3};
use nalgebra::DMatrix;
use rapier3d::parry::shape::HeightField;

use crate::{
    loaders::mesh::{InstanceRaw, Mesh, Vertex},
    physics::PhysicsManager,
    render::{self, DrawCall, Pipeline},
};
use bytemuck::cast_slice;
use wgpu::util::DeviceExt;

thread_local! {
    static WIREFRAME_BOX: OnceCell<Rc<Mesh>> = OnceCell::new();
    static DEBUG_HEIGHTFIELD: OnceCell<Option<Rc<Mesh>>> = OnceCell::new();
}

pub fn set_debug_heightfield(heightfield: &HeightField, scale: [f32; 3], offset: [f32; 3])
{
    DEBUG_HEIGHTFIELD.with(|cell| {
        cell.get_or_init(|| {
            render::with_device(|device| {
                Some(Rc::new(create_heightfield_wireframe(
                    device,
                    heightfield,
                    scale,
                    offset,
                )))
            })
        });
    });
}

fn create_heightfield_wireframe(
    device: &wgpu::Device,
    heightfield: &HeightField,
    scale: [f32; 3],
    offset: [f32; 3],
) -> Mesh
{
    let nrows = heightfield.nrows();
    let ncols = heightfield.ncols();

    let mut vertices = Vec::new();
    let mut indices = Vec::new();

    for row in 0..nrows
    {
        for col in 0..ncols
        {
            let x = col as f32 * scale[0];
            let y = heightfield.heights()[(row, col)] * scale[1];
            let z = row as f32 * scale[2];

            vertices.push(Vertex {
                position: [x + offset[0], y + offset[1], z + offset[2]],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            });
        }
    }

    for row in 0..nrows
    {
        for col in 0..ncols
        {
            let idx = (row * ncols + col) as u32;

            if col < ncols - 1
            {
                indices.push(idx);
                indices.push(idx + 1);
            }

            if row < nrows - 1
            {
                indices.push(idx);
                indices.push(idx + ncols as u32);
            }
        }
    }

    if !vertices.is_empty()
    {
        let first = &vertices[0].position;
        let last = &vertices[vertices.len() - 1].position;
        println!(
            "Heightfield bounds: ({}, {}, {}) to ({}, {}, {})",
            first[0], first[1], first[2], last[0], last[1], last[2]
        );
        println!(
            "Total vertices: {}, indices: {}",
            vertices.len(),
            indices.len()
        );
    }

    Mesh::new(device, &vertices, &indices)
}

pub fn render_collider_debug() -> Vec<DrawCall>
{
    let mut draw_calls = Vec::new();

    let aabbs = PhysicsManager::get_all_collider_aabbs();

    WIREFRAME_BOX.with(|cell| {
        let wireframe_box = cell.get_or_init(|| {
            render::with_device(|device| Rc::new(Mesh::create_wireframe_box(device)))
        });

        for (mins, maxs) in aabbs
        {
            let min = Vec3::from(mins);
            let max = Vec3::from(maxs);

            let center = (min + max) * 0.5;
            let size = max - min;

            let scale = Mat4::from_scale(size);
            let translation = Mat4::from_translation(center);
            let model = translation * scale;

            let instance_data = InstanceRaw {
                model: model.to_cols_array_2d(),
                dissolve_amount: 0.0,
                _padding: [0.0; 3],
            };

            let instance_buffer = render::with_device(|device| {
                device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("Debug Instance Buffer"),
                    contents: cast_slice(&[instance_data]),
                    usage: wgpu::BufferUsages::VERTEX,
                })
            });

            draw_calls.push(DrawCall {
                vertex_buffer: wireframe_box.vertex_buffer.clone(),
                index_buffer: wireframe_box.index_buffer.clone(),
                num_indices: wireframe_box.num_indices,
                model,
                pipeline: Pipeline::DebugLines,
                instance_buffer: Some(instance_buffer),
                num_instances: 1,
                tile_scale: 4.0,
                enable_dissolve: false,
                enable_snow_light: false,
                displacement_bind_group: None,
                entity: None,
            });
        }
    });

    DEBUG_HEIGHTFIELD.with(|cell| {
        if let Some(Some(heightfield_mesh)) = cell.get()
        {
            let instance_data = InstanceRaw {
                model: Mat4::IDENTITY.to_cols_array_2d(),
                dissolve_amount: 0.0,
                _padding: [0.0; 3],
            };

            let instance_buffer = render::with_device(|device| {
                device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("Heightfield Debug Instance Buffer"),
                    contents: cast_slice(&[instance_data]),
                    usage: wgpu::BufferUsages::VERTEX,
                })
            });

            draw_calls.push(DrawCall {
                vertex_buffer: heightfield_mesh.vertex_buffer.clone(),
                index_buffer: heightfield_mesh.index_buffer.clone(),
                num_indices: heightfield_mesh.num_indices,
                model: Mat4::IDENTITY,
                pipeline: Pipeline::DebugLines,
                instance_buffer: Some(instance_buffer),
                num_instances: 1,
                tile_scale: 4.0,
                enable_dissolve: false,
                enable_snow_light: false,
                displacement_bind_group: None,
                entity: None,
            });
        }
    });

    draw_calls
}
