use std::marker::PhantomData;

use crate::{
    config::CHUNK_SIZE,
    sim::{cell::Cell, chunk::Chunk, materials::MaterialDef, world::World},
};

struct ChunkAccess<'a> {
    ptr: *mut Chunk,
    len: usize,
    _marker: PhantomData<&'a mut [Chunk]>,
}

impl<'a> ChunkAccess<'a> {
    pub fn new(chunks: &'a mut [Chunk]) -> Self {
        Self {
            ptr: chunks.as_mut_ptr(),
            len: chunks.len(),
            _marker: PhantomData,
        }
    }
    unsafe fn get(&self, i: usize) -> &'a mut Chunk {
        debug_assert!(i < self.len);
        unsafe { &mut *self.ptr.add(i) }
    }
}

pub struct UpdateCtx<'a, 'b, 'c> {
    pub chunks: &'a mut [Option<&'b mut Chunk>; 9],
    pub seqno: u64,
    pub seqno_parity: u8,

    pub x: i32,
    pub y: i32,
    pub cell: &'c mut Cell,
    pub material: &'c MaterialDef,
}

fn get_cell(chunks: &[Option<&mut Chunk>; 9], x: i32, y: i32) -> Option<Cell> {
    let dcx = x.div_euclid(CHUNK_SIZE as i32);
    let dcy = y.div_euclid(CHUNK_SIZE as i32);
    if dcx != 0 || dcy != 0 {
        // in a different chunk
        let nc_x = x.rem_euclid(CHUNK_SIZE as i32) as u8;
        let nc_y = y.rem_euclid(CHUNK_SIZE as i32) as u8;

        return if let Some(chunk) = &chunks[(dcx + 1 + (dcy + 1) * 3) as usize] {
            Some(chunk.get_cell_at_local_position(nc_x, nc_y))
        } else {
            None
        };
    } else {
        if let Some(target) = &chunks[4] {
            Some(target.get_cell_at_local_position(x as u8, y as u8))
        } else {
            None
        }
    }
}

pub fn set_cell(chunks: &mut [Option<&mut Chunk>; 9], x: i32, y: i32, cell: Cell) {
    let dcx = x.div_euclid(CHUNK_SIZE as i32);
    let dcy = y.div_euclid(CHUNK_SIZE as i32);
    if dcx != 0 || dcy != 0 {
        // in a different chunk
        let nc_x = x.rem_euclid(CHUNK_SIZE as i32) as u8;
        let nc_y = y.rem_euclid(CHUNK_SIZE as i32) as u8;

        if let Some(chunk) = &mut chunks[(dcx + 1 + (dcy + 1) * 3) as usize] {
            chunk.set_cell_at_local_position(nc_x, nc_y, cell);
        }
    } else {
        if let Some(target) = &mut chunks[4] {
            target.set_cell_at_local_position(x as u8, y as u8, cell);
        }
    }
}

impl UpdateCtx<'_, '_, '_> {
    pub fn get_cell(&self, dx: i32, dy: i32) -> Option<Cell> {
        let x = self.x + dx;
        let y = self.y + dy;
        get_cell(self.chunks, x, y)
    }

    fn set_cell(&mut self, dx: i32, dy: i32, cell: Cell) {
        let x = self.x + dx;
        let y = self.y + dy;
        set_cell(self.chunks, x, y, cell);
    }

    pub fn candidates_swap(&mut self, candidates: &[(i32, i32)]) -> bool {
        for &(dx, dy) in candidates {
            let candidate_cell = self.get_cell(dx, dy);
            if candidate_cell.is_some_and(|c| c.material.def().density < self.material.density) {
                self.set_cell(0, 0, candidate_cell.unwrap());
                self.set_cell(dx, dy, *self.cell);
                return true;
            }
        }
        false
    }
}

pub fn sim_tick_chunk(chunks: &mut [Option<&mut Chunk>; 9], seqno: u64) {
    puffin::profile_function!();
    // scan bottom to top to enable contiguous falling
    let seqno_parity = (seqno as u8) & 0b1;

    if chunks[4].is_some() {
        for y in (0..CHUNK_SIZE as i32).rev() {
            for i in 0..CHUNK_SIZE as i32 {
                let x = if seqno_parity == 0 {
                    i
                } else {
                    (CHUNK_SIZE as i32) - i - 1
                };

                let mut cell = get_cell(chunks, x, y).unwrap();
                if cell.flags & 0b1 == seqno_parity {
                    // flip the parity bit
                    // TODO if the cell doesn't move this doesn't stay
                    cell.flags = cell.flags ^ 0b1;
                    let material = cell.material.def();

                    let mut update_ctx = UpdateCtx {
                        chunks,
                        seqno,
                        seqno_parity,

                        x,
                        y,
                        cell: &mut cell,
                        material,
                    };

                    if let Some(update) = material.sim_update {
                        update(&mut update_ctx);
                    }
                }
            }
        }
    }
}

const NEIGHBORHOOD_OFFSETS: [(i32, i32); 9] = [
    (-1, -1),
    (0, -1),
    (1, -1),
    (-1, 0),
    (0, 0),
    (1, 0),
    (-1, 1),
    (0, 1),
    (1, 1),
];

pub fn sim_tick(world: &mut World, seqno: u64) {
    puffin::profile_function!();
    let mut update_groups: [Vec<(i32, i32)>; 9] = Default::default();

    // assign a color to each chunk s.t. every chunk is surrounded by <= 8 chunks of different colors
    // --------------------
    // | 0, 1, 2, 0, 1, 2 |
    // | 3, 4, 5, 3, 4, 5 |
    // | 6, 7, 8, 6, 7, 8 |
    // | 0, 1, 2, 0, 1, 2 |
    // | 3, 4, 5, 3, 4, 5 |
    // | 6, 7, 8, 6, 7, 8 |
    // --------------------

    for (&(cx, cy), _) in &world.chunk_position_to_chunk_idx {
        let color = (cx.rem_euclid(3) * 3 + cy.rem_euclid(3)) as usize;
        update_groups[color].push((cx, cy));
    }

    for group in &update_groups {
        let access = ChunkAccess::new(&mut world.chunks);
        // TODO this can be parallelized since they will never share neighbours
        for &(cx, cy) in group {
            let mut chunks: [Option<&mut Chunk>; 9] = NEIGHBORHOOD_OFFSETS.map(|(dx, dy)| {
                world
                    .chunk_position_to_chunk_idx
                    .get(&(cx + dx, cy + dy))
                    .map(|&idx| unsafe { access.get(idx) })
            });

            sim_tick_chunk(&mut chunks, seqno);
        }
    }
}