use std::marker::PhantomData;

use fxhash::FxHashMap;
use rand::{Rng, SeedableRng, rngs::SmallRng};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

use crate::{
    config::CHUNK_SIZE,
    sim::{
        cell::{cell::Cell, materials::MaterialDef},
        cell_sim::{chunk::Chunk, 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) }
    }
}

unsafe impl Sync for ChunkAccess<'_> {}

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

#[inline]
fn neighbourhood_index(x: i8, y: i8) -> usize {
    (x + 1 + (y + 1) * 3) as usize
}

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

        chunks[neighbourhood_index(dcx as i8, dcy as i8)]
            .as_ref()
            .map(|chunk| chunk.get_cell_at_local_position(nc_x, nc_y))
    } else {
        chunks[4]
            .as_ref()
            .map(|target| target.get_cell_at_local_position(x as u8, y as u8))
    }
}

fn adjacent_chunks(x: u8, y: u8) -> Vec<usize> {
    if x == 0 {
        if y == 0 {
            vec![
                // L
                neighbourhood_index(-1, 0),
                // U
                neighbourhood_index(0, -1),
                // LU
                neighbourhood_index(-1, -1),
            ]
        } else if y == (CHUNK_SIZE - 1) as u8 {
            vec![
                // L
                neighbourhood_index(-1, 0),
                // D
                neighbourhood_index(0, 1),
                // LD
                neighbourhood_index(-1, 1),
            ]
        } else {
            // L
            vec![neighbourhood_index(-1, 0)]
        }
    } else if x == (CHUNK_SIZE - 1) as u8 {
        if y == 0 {
            vec![
                // R
                neighbourhood_index(1, 0),
                // U
                neighbourhood_index(0, -1),
                // RU
                neighbourhood_index(1, -1),
            ]
        } else if y == (CHUNK_SIZE - 1) as u8 {
            vec![
                // R
                neighbourhood_index(1, 0),
                // D
                neighbourhood_index(0, 1),
                // RD
                neighbourhood_index(1, 1),
            ]
        } else {
            // R
            vec![neighbourhood_index(1, 0)]
        }
    } else if y == 0 {
        // U
        vec![neighbourhood_index(0, -1)]
    } else if y == (CHUNK_SIZE - 1) as u8 {
        // D
        vec![neighbourhood_index(0, 1)]
    } else {
        vec![]
    }
}

pub fn set_cell(chunks: &mut [Option<&mut Chunk>; 9], x: i32, y: i32, cell: Cell) {
    let cx = x.div_euclid(CHUNK_SIZE);
    let cy = y.div_euclid(CHUNK_SIZE);
    let lx = x.rem_euclid(CHUNK_SIZE) as u8;
    let ly = y.rem_euclid(CHUNK_SIZE) as u8;
    if cx != 0 || cy != 0 {
        // in a different chunk
        if let Some(chunk) = &mut chunks[neighbourhood_index(cx as i8, cy as i8)] {
            chunk.set_cell_at_local_position(lx, ly, cell);
            chunk.needs_texture_update = true;
            chunk.sleeping = false;
            // if we're at the boundaries of the chunk, wake the adjacent chunk(s)
            for idx in adjacent_chunks(lx, ly) {
                if let Some(chunk) = chunks[idx].as_mut() {
                    chunk.sleeping = false;
                }
            }
        }
    } else {
        if let Some(target) = &mut chunks[4] {
            target.set_cell_at_local_position(x as u8, y as u8, cell);
            target.needs_texture_update = true;
            target.sleeping = false;
            // if we're at the boundaries of the chunk, wake the adjacent chunk(s)
            for idx in adjacent_chunks(lx, ly) {
                if let Some(chunk) = chunks[idx].as_mut() {
                    chunk.sleeping = false;
                }
            }
        }
    }
}

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,

    pub rng: &'c mut dyn Rng,
    swapped: bool,
}

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)
    }

    pub fn set_cell(&mut self, dx: i32, dy: i32, cell: Cell) {
        // cannot move out of the neighbourhood, but also cannot move to the edge of the neighbourhood
        // as this would wake a chunk outside of the neighbourhood
        debug_assert!(dx > -CHUNK_SIZE + 1 && dx < CHUNK_SIZE - 1);
        debug_assert!(dy > -CHUNK_SIZE + 1 && dy < CHUNK_SIZE - 1);
        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 {
            if let Some(mut candidate_cell) = self.get_cell(dx, dy)
                && candidate_cell.material.def().density < self.material.density
            {
                candidate_cell.reset_settled();
                candidate_cell.match_parity(self.seqno + 1);
                self.cell.reset_settled();
                self.cell.match_parity(self.seqno + 1);
                self.set_cell(0, 0, candidate_cell);
                self.set_cell(dx, dy, *self.cell);
                self.swapped = true;
                return true;
            }
        }
        false
    }
}

pub fn sim_tick_chunk(chunks: &mut [Option<&mut Chunk>; 9], seqno: u64) {
    puffin::profile_function!();
    let seqno_parity = (seqno as u8) & 0b1;
    let mut rng = SmallRng::seed_from_u64(seqno);

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

                let mut cell = get_cell(chunks, x, y).unwrap();
                let material = cell.material.def();

                if let Some(update) = material.sim_update
                // NOTE we only update parity when cells are updated, which means that static cells
                // are only evaluated every other tick
                // it also means that the settled counter increases every other tick
                    && cell.parity() == seqno_parity
                {
                    let mut update_ctx = UpdateCtx {
                        chunks,
                        seqno,
                        seqno_parity,

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

                        // TODO this is platform-dependent, will break for multiplayer
                        rng: &mut rng,
                        swapped: false,
                    };

                    update(&mut update_ctx);

                    if !update_ctx.swapped {
                        // the cell didn't move, so it's more settled, and we also need to update its state for parity
                        cell.increment_settled();
                        set_cell(chunks, x, y, cell);
                    }
                }
            }
        }
    }
}

pub fn sim_tick(world: &mut World, seqno: u64, use_threading: bool) {
    puffin::profile_function!();

    let mut columns: FxHashMap<i32, Vec<i32>> = FxHashMap::default();
    for &(cx, cy) in world.chunk_position_to_chunk_idx.keys() {
        columns.entry(cx).or_default().push(cy);
    }

    // color columns s.t. columns of same color are separated by two columns
    // and sort the column bottom-to-top
    // --------------------
    // | 0, 1, 2, 0, 1, 2 |
    // | 0, 1, 2, 0, 1, 2 |
    // | 0, 1, 2, 0, 1, 2 |
    // --------------------
    let mut columns_by_color: [Vec<(i32, Vec<i32>)>; 3] = Default::default();
    for (cx, mut cys) in columns {
        cys.sort_unstable_by(|a, b| b.cmp(a));
        columns_by_color[cx.rem_euclid(3) as usize].push((cx, cys));
    }

    let access = ChunkAccess::new(&mut world.chunks);

    for color in &columns_by_color {
        puffin::profile_scope!("chunk_color");

        let chunk_closure = |(cx, cys): &(i32, Vec<i32>)| {
            let cx = *cx;
            for &cy in cys {
                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) })
                });

                if let Some(target) = &mut chunks[4] {
                    if target.sleeping {
                        continue;
                    }
                    target.sleeping = true;
                }

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

        if use_threading {
            // TODO use forte
            color.par_iter().for_each(chunk_closure);
        } else {
            color.iter().for_each(chunk_closure);
        };
    }
}