use pyo3::exceptions::PyAssertionError;
use pyo3::prelude::*;

mod game;

use game::errors::InvalidAction;
use game::{Game, BOARD_SIZE};

#[pyclass(unsendable)]
struct PyBlokus(Game);

#[pymethods]
impl PyBlokus {
    #[new]
    pub fn new() -> Self {
        Self(Game::new())
    }

    pub fn reset(&mut self) {
        self.0 = Game::new();
    }

    #[getter(agents)]
    pub fn agents(&self) -> Vec<usize> {
        (0usize..self.0.num_agents).collect()
    }

    #[getter(agent_selection)]
    pub fn agent_selection(&self) -> usize {
        self.0.agent_selection
    }

    #[getter(num_actions)]
    pub fn num_actions(&self) -> usize {
        self.0.action_set.actions.len()
    }

    #[getter(terminations)]
    pub fn terminations(&self) -> [bool; 4] {
        self.0.terminations()
    }

    #[getter(truncations)]
    pub fn truncations(&self) -> [bool; 4] {
        self.0.terminations()
    }

    #[getter(rewards)]
    pub fn rewards(&self) -> Vec<i16> {
        self.0
            .rewards()
            .map_or_else(|| vec![0i16; self.0.num_agents], |x| x)
    }

    pub fn observe(&mut self, action_idx: usize) -> PyObservation {
        let obs = self.0.observe(action_idx);
        PyObservation {
            observation: obs.0,
            action_mask: obs.1.to_vec(),
        }
    }

    pub fn step(&mut self, action_idx: usize) -> Result<(), InvalidAction> {
        self.0.step(action_idx)?;
        Ok(())
    }

    pub fn render(&self) {
        self.0.render();
    }
}

#[pyclass(unsendable)]
struct PyObservation {
    observation: [[[bool; BOARD_SIZE]; BOARD_SIZE]; 4],
    action_mask: Vec<bool>,
}

#[pymethods]
impl PyObservation {
    #[getter(observation)]
    pub fn observation(&self) -> [[[bool; BOARD_SIZE]; BOARD_SIZE]; 4] {
        [
            self.observation[0],
            self.observation[1],
            self.observation[2],
            self.observation[3],
        ]
    }

    #[getter(action_mask)]
    pub fn action_mask(&self) -> Vec<bool> {
        self.action_mask.clone()
    }
}

impl From<InvalidAction> for PyErr {
    fn from(error: InvalidAction) -> Self {
        PyAssertionError::new_err(error.to_string())
    }
}

#[pymodule]
mod _blokus {
    #[pymodule_export]
    use super::PyBlokus;
    #[pymodule_export]
    use super::PyObservation;
}

#[cfg(test)]
mod tests {
    use crate::game::bitboard::{separating_bit_mask, Bitboard};
    use crate::game::Game;

    #[test]
    fn test_action_generation_valid() {
        let game = Game::new();
        assert_eq!(game.action_set.actions.len(), 30433);
        for a in game.action_set.actions {
            assert!(
                (a.bitboard & !separating_bit_mask()) == Bitboard::default(),
                "Invalid action!"
            );
        }
    }
    #[test]
    fn test_rotation_clock_valid() {
        let game = Game::new();
        for a in game.action_set.actions {
            let board_rot = a
                .bitboard
                .rotate_clock()
                .rotate_clock()
                .rotate_clock()
                .rotate_clock();
            assert!(board_rot == a.bitboard, "Invalid rotation!");
        }
    }
    #[test]
    fn test_rotation_anticlock_valid() {
        let game = Game::new();
        for a in game.action_set.actions {
            let board_rot = a
                .bitboard
                .rotate_anticlock()
                .rotate_anticlock()
                .rotate_anticlock()
                .rotate_anticlock();
            assert!(board_rot == a.bitboard, "Invalid rotation!");
        }
    }
    #[test]
    fn test_agent_selection() {
        let mut game = Game::new();
        assert_eq!(game.agent_selection, 0);
        let _ = game.step(400);
        assert_eq!(game.agent_selection, 1);
        let _ = game.step(400);
        assert_eq!(game.agent_selection, 2);
        let _ = game.step(400);
        assert_eq!(game.agent_selection, 3);
        let _ = game.step(0);
        assert_eq!(game.agent_selection, 0);
    }

    #[test]
    fn test_scoring() {
        let mut game = Game::new();
        game.step(400).unwrap();
        game.step(400).unwrap();
        game.step(400).unwrap();
        game.step(400).unwrap();
        game.observe(0);
        game.step(22).unwrap();
        game.observe(1);
        game.step(22).unwrap();
        game.observe(2);
        assert!(game.rewards().is_none());
        game.agents[0].done = true;
        game.agents[1].done = true;
        game.agents[2].done = true;
        game.agents[3].done = true;
        assert_eq!(game.rewards().unwrap()[0], -86);
        assert_eq!(game.rewards().unwrap()[1], -86);
        assert_eq!(game.rewards().unwrap()[2], -87);
        assert_eq!(game.rewards().unwrap()[3], -87);
    }
}
