pub struct Grid {
    pub materials: Vec<u8>,
    pub healths: Vec<u8>,
    pub temps: Vec<u8>,
    pub flags: Vec<u8>,
    pub width: u32,
    pub height: u32,
}

pub const FLAG_UPDATED: u8 = 0b0000_0001;
pub const FLAG_PLAYER: u8 = 0b0000_0010;
pub const FLAG_MOMENTUM: u8 = 0b0000_0100;

impl Grid {
    pub fn new(width: u32, height: u32) -> Self {
        let size = (width * height) as usize;
        Self {
            materials: vec![0; size],
            healths: vec![255; size],
            temps: vec![0; size],
            flags: vec![0; size],
            width,
            height,
        }
    }

    #[inline]
    pub fn index(&self, x: u32, y: u32) -> usize {
        (y * self.width + x) as usize
    }

    pub fn in_bounds(&self, x: i32, y: i32) -> bool {
        x >= 0 && x < self.width as i32 && y >= 0 && y < self.height as i32
    }

    #[inline]
    pub fn get(&self, x: u32, y: u32) -> u8 {
        self.materials[self.index(x, y)]
    }

    #[inline]
    pub fn get_signed(&self, x: i32, y: i32) -> u8 {
        if self.in_bounds(x, y) {
            self.materials[self.index(x as u32, y as u32)]
        } else {
            0
        }
    }

    #[inline]
    pub fn set(&mut self, x: u32, y: u32, material: u8) {
        let idx = self.index(x, y);
        self.materials[idx] = material;
        self.flags[idx] |= FLAG_UPDATED;
        if material == 6 {
            self.temps[idx] = 200;
            self.healths[idx] = 255;
        } else if material == 9 {
            self.temps[idx] = 200;
            self.healths[idx] = 200;
        } else if material == 11 {
            self.temps[idx] = 0;
        }
    }

    #[inline]
    pub fn swap(&mut self, x1: u32, y1: u32, x2: u32, y2: u32) {
        let i1 = self.index(x1, y1);
        let i2 = self.index(x2, y2);
        self.materials.swap(i1, i2);
        self.healths.swap(i1, i2);
        self.temps.swap(i1, i2);
        self.flags[i1] |= FLAG_UPDATED;
        self.flags[i2] |= FLAG_UPDATED;
    }

    pub fn is_empty(&self, x: u32, y: u32) -> bool {
        self.get(x, y) == 0
    }

    pub fn clear_flags(&mut self) {
        for f in &mut self.flags {
            *f &= !FLAG_UPDATED;
        }
    }

    pub fn fill_rect(&mut self, x: u32, y: u32, w: u32, h: u32, material: u8) {
        for dy in 0..h {
            for dx in 0..w {
                let px = x + dx;
                let py = y + dy;
                if px < self.width && py < self.height {
                    self.set(px, py, material);
                }
            }
        }
    }

    pub fn fill_circle(&mut self, cx: u32, cy: u32, radius: u32, material: u8) {
        let r = radius as i32;
        for dy in -r..=r {
            for dx in -r..=r {
                if dx * dx + dy * dy <= r * r {
                    let px = (cx as i32 + dx) as u32;
                    let py = (cy as i32 + dy) as u32;
                    if px < self.width && py < self.height {
                        self.set(px, py, material);
                    }
                }
            }
        }
    }
}