use std::f32::consts::TAU;

use bevy::prelude::*;
use rand::Rng;

pub static MAX_MAP_WIDTH: usize = 1024;
use crate::physics::physics_body::PhysicsBody;

const CELL_SIZE: usize = 8;
const MAX_MAP_SIZE: usize = MAX_MAP_WIDTH * 2;
const ARRAY_WIDTH: usize = MAX_MAP_SIZE / CELL_SIZE;
const GRID_ARRAY_SIZE: usize = ARRAY_WIDTH * ARRAY_WIDTH;

const DENSITY_FORCE_SCALE: f32 = 6000.0;

#[derive(Resource)]
pub struct DensityGrid {
    grid: Box<[i32; GRID_ARRAY_SIZE]>,
}

impl DensityGrid {
    pub fn new() -> Self {
        DensityGrid { grid: Box::new([0; GRID_ARRAY_SIZE]) }
    }
}

pub fn pos_to_grid_cell(pos: Vec2) -> UVec2 {
    let ix = ((pos.x + MAX_MAP_WIDTH as f32) / CELL_SIZE as f32) as u32;
    let iy = ((pos.y + MAX_MAP_WIDTH as f32) / CELL_SIZE as f32) as u32;
    UVec2::new(ix, iy)
}

fn get_index_from_uvec2(grid_cell: UVec2) -> usize {
    if grid_cell.x >= ARRAY_WIDTH as u32 || grid_cell.y >= ARRAY_WIDTH as u32 {
        panic!("Index out of bounds: {:?}", grid_cell);
    }
    (grid_cell.x + grid_cell.y * ARRAY_WIDTH as u32) as usize
}

#[derive(Component, Default)]
pub struct DensityObject {
    random_offset: Vec2,
}

pub fn init_density_object(mut density_grid: ResMut<DensityGrid>, mut query: Query<(&Transform, &mut DensityObject)>) {
    let mut rng = rand::rng();
    for (transform, mut density_obj) in &mut query {
        let grid_cell = pos_to_grid_cell(transform.translation.xy());
        // println!("grid cell {:?}", grid_cell);
        density_obj.random_offset = Vec2::from_angle(rng.random_range(0.0..TAU)) * 0.2;
        let index = get_index_from_uvec2(grid_cell);
        density_grid.grid[index] += 1;
    }
}

pub fn density_grid_update(mut density_grid: ResMut<DensityGrid>, query: Query<&Transform, With<DensityObject>>) {
    density_grid.grid.fill(0);

    for transform in &query {
        let cell = pos_to_grid_cell(transform.translation.xy());
        let idx = get_index_from_uvec2(cell);
        density_grid.grid[idx] += 1;
    }
}

pub fn density_grid_force(
    density_grid: Res<DensityGrid>,
    mut query: Query<(&mut PhysicsBody, &Transform, &DensityObject)>,
) {
    for (mut physics_body, transform, density_obj) in &mut query {
        let grid_cell = pos_to_grid_cell(transform.translation.xy());
        let mut density_gradient = Vec2::ZERO;

        for dx in [-1, 0, 1] {
            for dy in [-1, 0, 1] {
                let offset = IVec2 { x: dx, y: dy };
                let neighbor_cell = (offset + grid_cell.as_ivec2()).as_uvec2();
                let index = get_index_from_uvec2(neighbor_cell);
                let density = density_grid.grid[index];

                // println!("offset : {:?} density: {:?}", offset.as_vec2(), density);
                density_gradient += offset.as_vec2() * density as f32;
            }
        }
        // println!("density : {:?}", density_gradient);
        density_gradient += density_obj.random_offset; // To not have a "grid" thingy
        physics_body.force -= density_gradient * DENSITY_FORCE_SCALE;
    }
}

pub struct DensityGridPlugin;

impl Plugin for DensityGridPlugin {
    fn build(&self, app: &mut App) {
        app.insert_resource(DensityGrid::new())
            .add_systems(PostStartup, init_density_object)
            .add_systems(FixedUpdate, (density_grid_update, density_grid_force).chain());
    }
}