#![no_std]
#![no_main]
#![allow(incomplete_features)]
#![feature(generic_const_exprs)]

use core::{
    mem::MaybeUninit,
    ops::{Index, IndexMut},
};

use fpga_lib::{
    panic_handler,
    vga::{Color, BLACK, FONT_HEIGHT, FONT_WIDTH, RED, WHITE},
};
use rand::{seq::IndexedRandom, Rng, SeedableRng};

const fn get_color(x: u64) -> Color {
    match x {
        2 => Color::from_rgb(238, 228, 222),
        4 => Color::from_rgb(237, 224, 200),
        8 => Color::from_rgb(242, 177, 120),
        16 => Color::from_rgb(245, 149, 99),
        32 => Color::from_rgb(246, 124, 95),
        64 => Color::from_rgb(246, 94, 59),
        128 => Color::from_rgb(237, 207, 114),
        256 => Color::from_rgb(237, 204, 97),
        512 => Color::from_rgb(237, 200, 80),
        1024 => Color::from_rgb(237, 197, 63),
        2048 => Color::from_rgb(237, 194, 46),
        4096 => Color::from_rgb(238, 102, 106),
        8192 => Color::from_rgb(238, 78, 90),
        16384 => Color::from_rgb(241, 65, 65),
        32768 => Color::from_rgb(114, 178, 212),
        65536 => Color::from_rgb(88, 159, 226),
        131072 => Color::from_rgb(23, 129, 204),
        _ => Color::from_rgb(255, 255, 255),
    }
}

pub struct NZip<T, const N: usize> {
    iters: [T; N],
}

impl<T: Iterator, const N: usize> Iterator for NZip<T, N> {
    type Item = [T::Item; N];

    fn next(&mut self) -> Option<Self::Item> {
        let mut res = MaybeUninit::<[T::Item; N]>::uninit();
        for (cpt, iter) in self.iters.iter_mut().enumerate() {
            unsafe { *(res.as_mut_ptr() as *mut T::Item).add(cpt) = iter.next()? };
        }
        Some(unsafe { res.assume_init() })
    }
}

impl<T: DoubleEndedIterator, const N: usize> DoubleEndedIterator for NZip<T, N> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let mut res = MaybeUninit::<[T::Item; N]>::uninit();
        for (cpt, iter) in self.iters.iter_mut().enumerate() {
            unsafe { *(res.as_mut_ptr() as *mut T::Item).add(cpt) = iter.next_back()? };
        }
        Some(unsafe { res.assume_init() })
    }
}

#[derive(Clone)]
pub struct Game<const SIZE: usize>
where
    [(); SIZE * SIZE]:,
{
    grid: [u64; SIZE * SIZE],
    rng: rand::rngs::SmallRng,
    score: u64,
}

impl<const SIZE: usize> Index<(usize, usize)> for Game<SIZE>
where
    [(); SIZE * SIZE]:,
{
    type Output = u64;

    fn index(&self, index: (usize, usize)) -> &Self::Output {
        &self.grid[index.1 * SIZE + index.0]
    }
}
impl<const SIZE: usize> IndexMut<(usize, usize)> for Game<SIZE>
where
    [(); SIZE * SIZE]:,
{
    fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
        &mut self.grid[index.1 * SIZE + index.0]
    }
}

impl<const SIZE: usize> Game<SIZE>
where
    [(); SIZE * SIZE]:,
{
    pub fn new() -> Self {
        Self {
            grid: [0; SIZE * SIZE],
            rng: rand::rngs::SmallRng::seed_from_u64(compile_rand::compile_rand!(u64)),
            score: 0,
        }
    }
    pub fn init() -> Self {
        let mut g = Self::new();
        g.add();
        g.add();
        g
    }
    pub fn reset(&mut self) {
        *self = Self::new()
    }
    pub fn add(&mut self) {
        let mut pos = [(0, 0); SIZE * SIZE];
        let mut cpt = 0;
        for y in 0..SIZE {
            for x in 0..SIZE {
                if self[(x, y)] == 0 {
                    pos[cpt] = (x, y);
                    cpt += 1;
                }
            }
        }
        let pos = *pos[0..cpt].choose(&mut self.rng).unwrap();
        self[pos] = (self.rng.random_range(0..=1) + 1) * 2;
    }
    pub fn lost(&self) -> bool {
        let mut copy = self.clone();
        !(copy.left() || copy.down() || copy.right() || copy.up())
    }
    fn rotate(&mut self) {
        let mut out = self.grid;

        for i in 0..SIZE {
            for j in 0..SIZE {
                out[j * SIZE + (SIZE - i - 1)] = self.grid[i * SIZE + j];
            }
        }

        self.grid = out;
    }

    fn rotate_back(&mut self) {
        let mut out = self.grid;

        for i in 0..SIZE {
            for j in 0..SIZE {
                out[(SIZE - j - 1) * SIZE + i] = self.grid[i * SIZE + j];
            }
        }

        self.grid = out;
    }
    pub fn left(&mut self) -> bool {
        let mut changed = false;
        for y in 0..SIZE {
            let mut index = 0;
            let mut fixed = false;
            let mut zero_found = false;
            for x in 0..SIZE {
                if self[(x, y)] == 0 {
                    zero_found = true;
                    continue;
                } else if zero_found {
                    changed = true;
                }
                if index != 0 && self[(index - 1, y)] == (self[(x, y)]) && !fixed {
                    self[(index - 1, y)] = self[(x, y)] * 2;
                    self[(x, y)] = 0;
                    self.score += self[(index - 1, y)];
                    fixed = true;
                    changed = true;
                } else {
                    self[(index, y)] = self[(x, y)];
                    if index != x {
                        self[(x, y)] = 0;
                    }
                    fixed = false;
                    index += 1;
                }
            }
        }
        changed
    }
    pub fn down(&mut self) -> bool {
        self.rotate();
        let res = self.left();
        self.rotate_back();
        res
    }
    pub fn right(&mut self) -> bool {
        self.rotate();
        self.rotate();
        let res = self.left();
        self.rotate_back();
        self.rotate_back();
        res
    }
    pub fn up(&mut self) -> bool {
        self.rotate_back();
        let res = self.left();
        self.rotate();
        res
    }
    pub fn run(&mut self) -> ! {
        let mut buttons = [false; 4];
        let funcs = [Self::left, Self::right, Self::up, Self::down];
        let mut lost = false;
        self.draw();
        loop {
            for direction in 0..4 {
                let button = unsafe { fpga_lib::interface::Buttons::get_button(direction as u8) };
                if button && !buttons[direction] {
                    buttons[direction] = true;

                    if funcs[direction](self) {
                        self.add();
                        self.draw();
                        if self.lost() {
                            lost = true;
                            break;
                        }
                    }
                } else if !button && buttons[direction] {
                    buttons[direction] = false
                }
            }
            if lost {
                unsafe {
                    fpga_lib::vga::VGA::draw_string(
                        OFFSET,
                        OFFSET + (SIZE as u16) * BLOCK_SIZE + FONT_HEIGHT as u16 + 10,
                        "Game Over !",
                        RED,
                    )
                };
                loop {}
            }
        }
    }
}

impl<const SIZE: usize> Default for Game<SIZE>
where
    [(); SIZE * SIZE]:,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<const SIZE: usize> Game<SIZE>
where
    [(); SIZE * SIZE]:,
{
    fn draw(&self) {
        fpga_lib::vga::VGA::clear_screen(BLACK);
        for i in 0..SIZE as u16 + 1 {
            for x in 0..BLOCK_SIZE * SIZE as u16 {
                unsafe {
                    fpga_lib::vga::VGA::write_pixel_unsafe(
                        OFFSET + x,
                        OFFSET + i * BLOCK_SIZE,
                        WHITE,
                    )
                };
                unsafe {
                    fpga_lib::vga::VGA::write_pixel_unsafe(
                        OFFSET + i * BLOCK_SIZE,
                        OFFSET + x,
                        WHITE,
                    )
                };
            }
        }
        for y in 0..SIZE {
            for x in 0..SIZE {
                if self[(x, y)] != 0 {
                    unsafe {
                        fpga_lib::vga::VGA::draw_u64(
                            OFFSET + x as u16 * BLOCK_SIZE + BLOCK_SIZE / 2
                                - (fpga_lib::vga::FONT_WIDTH as u16
                                    * (self[(x, y)].ilog10() as u16 + 1))
                                    / 2,
                            OFFSET + y as u16 * BLOCK_SIZE + BLOCK_SIZE / 2
                                - fpga_lib::vga::FONT_HEIGHT as u16 / 2,
                            self[(x, y)],
                            get_color(self[(x, y)]),
                            BLACK,
                        )
                    };
                }
            }
        }
        unsafe {
            fpga_lib::vga::VGA::draw_string(
                OFFSET,
                OFFSET + (SIZE as u16) * BLOCK_SIZE + 5,
                "Score :",
                WHITE,
            );
            fpga_lib::vga::VGA::draw_u64(
                OFFSET + FONT_WIDTH as u16 * 8,
                OFFSET + (SIZE as u16) * BLOCK_SIZE + 5,
                self.score,
                WHITE,
                BLACK,
            );
        }
    }
}

panic_handler! {}

#[unsafe(no_mangle)]
pub extern "C" fn main() -> ! {
    let mut g = Game::<4>::init();
    g.run()
}

const BLOCK_SIZE: u16 = 40;
const OFFSET: u16 = 20;
