mod maps;
mod rules;

use std::{
    collections::VecDeque,
    fmt::Display,
    num::NonZeroUsize,
    ops::{Deref, DerefMut},
    sync::Arc,
};

use az::SaturatingAs;
use bitvec::prelude::*;
use log::{error, warn};
use maps::{Map, Maps};
use rand::prelude::*;
use rules::{Ruleset, Rulesets};

use super::{Coord, Direction, wire::Request};

#[derive(Debug, PartialEq, Eq, Clone)]
enum SmallBitBox {
    Stack {
        storage: BitArr!(for Self::STACK_BITS),
        len: NonZeroUsize,
    },
    Heap(BitBox),
}

impl SmallBitBox {
    const STACK_BITS: usize = usize::BITS as usize * 3;

    fn new(initial: bool, len: usize) -> Self {
        if let len @ 1..Self::STACK_BITS = len {
            let Some(len) = NonZeroUsize::new(len) else {
                unreachable!()
            };
            let mut storage = BitArray::ZERO;
            storage.fill(initial);
            Self::Stack { storage, len }
        } else {
            Self::Heap(bitbox![u8::from(initial); len])
        }
    }
}

impl Deref for SmallBitBox {
    type Target = BitSlice;

    fn deref(&self) -> &Self::Target {
        match self {
            Self::Stack { storage, len } => &storage[..len.get()],
            Self::Heap(bit_box) => bit_box,
        }
    }
}

impl DerefMut for SmallBitBox {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            Self::Stack { storage, len } => &mut storage[..len.get()],
            Self::Heap(bit_box) => bit_box,
        }
    }
}

#[derive(Clone)]
pub struct Game {
    pub board: Board,
    map: Maps,
    ruleset: Rulesets,
}

impl From<&Request> for Game {
    fn from(value: &Request) -> Self {
        let board = value.into();
        let map = (&value.game).into();
        let ruleset = (&value.game.ruleset).into();
        Self {
            board,
            map,
            ruleset,
        }
    }
}

impl Game {
    pub fn simulate_random(&mut self, rng: &mut impl Rng) {
        loop {
            let random_actions: Vec<_> = (0..self.board.snakes.len())
                .filter_map(|i| {
                    self.board
                        .valid_actions_index(i)
                        .choose(rng)
                        .map(|direction| (i.saturating_as(), direction))
                })
                .collect();
            if self.next_turn(&random_actions, rng) {
                break;
            }
        }
    }

    pub fn next_turn_random(&mut self, actions: &[(u8, Direction)], rng: &mut impl Rng) -> bool {
        let random_actions: Vec<_> = (0..self.board.snakes.len())
            .filter_map(|i| {
                actions
                    .iter()
                    .find(|(j, _)| i == usize::from(*j))
                    .copied()
                    .or_else(|| {
                        self.board
                            .valid_actions_index(i)
                            .choose(rng)
                            .map(|direction| (i.saturating_as(), direction))
                    })
            })
            .collect();
        self.next_turn(&random_actions, rng)
    }

    pub fn next_turn(&mut self, actions: &[(u8, Direction)], rng: &mut impl Rng) -> bool {
        self.map.pre_update(&mut self.board, rng);
        let game_over = self.ruleset.execute(&mut self.board, actions);
        self.map.post_update(&mut self.board, rng);
        self.board.turn += 1;
        game_over
    }
}

#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Board {
    width: u8,
    height: u8,
    turn: u32,
    food: SmallBitBox,
    hazard: SmallBitBox,
    free: SmallBitBox,
    snakes: Vec<Snake>,
    id_map: Arc<[(u8, Arc<str>)]>,
}

impl From<&Request> for Board {
    fn from(value: &Request) -> Self {
        let width = value.board.width;
        let height = value.board.height;
        let fields = usize::from(width) * usize::from(height);
        let id_map = value
            .board
            .snakes
            .iter()
            .enumerate()
            .filter_map(|(i, snake)| {
                u8::try_from(i)
                    .inspect_err(|e| warn!("unable to convert id to u8: {e}"))
                    .ok()
                    .map(|i| (i, snake.id.clone()))
            })
            .collect::<Vec<_>>();
        let mut board = Self {
            width,
            height,
            turn: value.turn,
            food: SmallBitBox::new(false, fields),
            hazard: SmallBitBox::new(false, fields),
            free: SmallBitBox::new(true, fields),
            snakes: Vec::with_capacity(value.board.snakes.len()),
            id_map: id_map.into(),
        };

        for &food in &value.board.food {
            let index = board.coord_to_linear(food);
            board.food.set(index, true);
        }

        for &hazard in &value.board.hazards {
            let index = board.coord_to_linear(hazard);
            board.hazard.set(index, true);
        }

        let constrictor = value.game.ruleset.name.as_ref() == "constrictor";
        for (id, snake) in value.board.snakes.iter().enumerate() {
            for &tile in snake
                .body
                .iter()
                .take(snake.body.len() - usize::from(!constrictor))
            {
                let index = board.coord_to_linear(tile);
                board.free.set(index, false);
            }
            let snake = Snake {
                body: snake.body.iter().rev().copied().collect(),
                id: u8::try_from(id).unwrap_or(u8::MAX),
                health: snake.health,
            };
            board.snakes.push(snake);
        }

        board
    }
}

impl Display for Board {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "{}x{} @ {}", self.width, self.height, self.turn)?;

        for y in (0..self.height).rev() {
            for x in 0..self.width {
                let tile = Coord { x, y };
                if self.snakes.iter().any(|snake| snake.head() == tile) {
                    write!(f, "H")?;
                    continue;
                }
                if self.snakes.iter().any(|snake| snake.tail() == tile) {
                    write!(f, "T")?;
                    continue;
                }

                let index = self.coord_to_linear(tile);
                if !self.free[index] {
                    write!(f, "S")?;
                    continue;
                }
                if self.food[index] {
                    write!(f, "f")?;
                    continue;
                }
                if self.hazard[index] {
                    write!(f, "h")?;
                    continue;
                }

                write!(f, ".")?;
            }
            writeln!(f)?;
        }

        Ok(())
    }
}

impl Board {
    #[must_use]
    pub fn get_id(&self, str_id: &str) -> Option<u8> {
        self.id_map
            .iter()
            .find(|(_, str_id2)| *str_id == **str_id2)
            .map(|(id, _)| *id)
    }

    #[must_use]
    pub const fn turn(&self) -> u32 {
        self.turn
    }

    #[must_use]
    pub fn num_snakes(&self) -> usize {
        self.snakes.len()
    }

    #[must_use]
    pub fn alive(&self, id: u8) -> bool {
        self.id_to_index(id).is_some()
    }

    #[must_use]
    pub fn length(&self, id: u8) -> usize {
        let Some(index) = self.id_to_index(id) else {
            return 0;
        };
        self.snakes[index].body.len()
    }

    #[must_use]
    pub fn max_length(&self) -> usize {
        self.snakes
            .iter()
            .map(|snake| snake.body.len())
            .max()
            .unwrap_or(0)
    }

    #[must_use]
    pub fn is_food(&self, tile: Coord) -> bool {
        let index = self.coord_to_linear(tile);
        self.food[index]
    }

    #[must_use]
    pub fn is_hazard(&self, tile: Coord) -> bool {
        let index = self.coord_to_linear(tile);
        self.hazard[index]
    }

    #[must_use]
    pub fn is_free(&self, tile: Coord) -> bool {
        if !self.is_in_bounds(tile) {
            return false;
        }
        let index = self.coord_to_linear(tile);
        self.free[index]
    }

    pub fn valid_actions(&self, id: u8) -> impl Iterator<Item = Direction> + use<'_> {
        let index = self.id_to_index(id);
        index
            .into_iter()
            .flat_map(|index| self.valid_actions_index(index))
    }

    #[must_use]
    pub fn random_action(&self, id: u8, rng: &mut impl RngCore) -> Direction {
        let Some(index) = self.id_to_index(id) else {
            return Direction::Up;
        };
        self.valid_actions_index(index)
            .choose(rng)
            .unwrap_or(Direction::Up)
    }

    pub fn random_actions<'a, R: RngCore>(
        &self,
        rng: &'a mut R,
    ) -> impl Iterator<Item = (u8, Direction)> + use<'_, 'a, R> {
        (0..self.snakes.len()).map(|index| {
            (
                self.snakes[index].id,
                self.valid_actions_index(index)
                    .choose(rng)
                    .unwrap_or(Direction::Up),
            )
        })
    }

    fn id_to_index(&self, id: u8) -> Option<usize> {
        self.snakes.binary_search_by_key(&id, |snake| snake.id).ok()
    }

    fn valid_actions_index(&self, index: usize) -> impl Iterator<Item = Direction> + use<'_> {
        let head = self.snakes[index].head();
        enum_iterator::all::<Direction>()
            .filter(move |direction| self.is_free(head.wrapping_apply(*direction)))
    }

    fn update_free_map(&mut self) {
        // free tails
        for snake in &self.snakes {
            let tail = snake.tail();
            let pre_tail = snake.body[snake.body.len() - 2];
            if tail != pre_tail {
                let tail_index = self.coord_to_linear(tail);
                self.free.set(tail_index, true);
            }
        }
        // block heads
        for snake in &self.snakes {
            let head = snake.head();
            let head_index = self.coord_to_linear(head);
            self.free.set(head_index, false);
        }
    }

    fn get_unoccupied_points(
        &self,
        include_possible_moves: bool,
        include_hazards: bool,
    ) -> impl Iterator<Item = Coord> + use<'_> {
        let possible_moves: Vec<_> = if include_possible_moves {
            Vec::new()
        } else {
            self.snakes
                .iter()
                .flat_map(|snake| {
                    let head = snake.head();
                    enum_iterator::all::<Direction>()
                        .map(move |direction| head.wrapping_apply(direction))
                        .filter(|tile| self.is_in_bounds(*tile))
                        .map(|tile| self.coord_to_linear(tile))
                })
                .collect()
        };
        self.free
            .iter()
            .zip(self.food.iter())
            .zip(self.hazard.iter())
            .enumerate()
            .filter(move |(i, ((free, food), hazard))| {
                **free
                    && !**food
                    && (include_hazards || !**hazard)
                    && (include_possible_moves || !possible_moves.contains(i))
            })
            .map(|(i, _)| self.linear_to_coord(i))
    }

    fn coord_to_linear(&self, coord: Coord) -> usize {
        usize::from(coord.x) + usize::from(coord.y) * usize::from(self.width)
    }

    fn linear_to_coord(&self, linear: usize) -> Coord {
        let x = (linear % usize::from(self.width)).saturating_as();
        let y = (linear / usize::from(self.width)).saturating_as();
        Coord { x, y }
    }

    const fn is_in_bounds(&self, coord: Coord) -> bool {
        coord.x < self.width && coord.y < self.height
    }
}

#[derive(Debug, PartialEq, Eq, Clone)]
struct Snake {
    id: u8,
    health: u8,
    body: VecDeque<Coord>,
}

impl Snake {
    pub fn head(&self) -> Coord {
        self.body.back().copied().unwrap_or_else(|| {
            error!("Snake without a head: {self:?}");
            Coord { x: 0, y: 0 }
        })
    }

    pub fn tail(&self) -> Coord {
        self.body.front().copied().unwrap_or_else(|| {
            error!("Snake without a tail: {self:?}");
            Coord { x: 0, y: 0 }
        })
    }

    pub fn advance(&mut self, head: Coord) {
        self.body.push_back(head);
        self.body.pop_front();
    }

    pub fn feed(&mut self) {
        if let Some(tail) = self.body.front() {
            self.body.push_front(*tail);
            self.health = 100;
        }
    }
}
