arimaa-engine 0.1.0

use anyhow::Result;
use std::fmt::{self, Display, Formatter};
use std::str::FromStr;

use super::constants::{BOARD_HEIGHT, BOARD_WIDTH};
use super::game_state::GameState;
use super::piece_board::PieceBoard;
use super::{Piece, Square};

impl Display for GameState {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let move_number = self.get_move_number();
        let curr_player = if self.is_p1_turn_to_move() { "g" } else { "s" };
        let piece_board = &self.get_piece_board();
        writeln!(f, "{}{}", move_number, curr_player)?;

        writeln!(f, " +-----------------+")?;

        for row_idx in 0..BOARD_HEIGHT {
            write!(f, "{}|", BOARD_HEIGHT - row_idx)?;
            for col_idx in 0..BOARD_WIDTH {
                let idx = (row_idx * BOARD_WIDTH + col_idx) as u8;
                let square = Square::from_index(idx);
                let letter = if let Some(piece) = piece_board.get_piece_type_at_square(&square) {
                    let is_p1_piece = is_p1_piece(square.as_bit_board(), piece_board);
                    convert_piece_to_letter(&piece, is_p1_piece)
                } else if idx == 18 || idx == 21 || idx == 42 || idx == 45 {
                    "x".to_string()
                } else {
                    " ".to_string()
                };

                write!(f, " {}", letter)?;
            }
            writeln!(f, " |")?;
        }

        writeln!(f, " +-----------------+")?;
        writeln!(f, "   a b c d e f g h")?;

        Ok(())
    }
}

impl FromStr for GameState {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self> {
        // Compact AEI format: {optional_number}{player} [{64-char-board}]
        // e.g. "g [rredchrrrhrdcmrr                                MHHEDDCRRRRRRRRC]"
        //      "1g [rredchrrrhrdcmrr                                MHHEDDCRRRRRRRRC]"
        let compact_regex = regex::Regex::new(r"^\s*(?:(\d+))?([gswb])\s*\[([^\]]+)\]").unwrap();
        if let Some(caps) = compact_regex.captures(s) {
            let move_number: usize = caps.get(1).map_or(1, |m| m.as_str().parse().unwrap_or(1));
            let player = caps.get(2).unwrap().as_str();
            let p1_turn_to_move = player != "s" && player != "b";
            let pieces_str = caps.get(3).unwrap().as_str();

            let mut p1_pieces = 0u64;
            let mut elephants = 0u64;
            let mut camels = 0u64;
            let mut horses = 0u64;
            let mut dogs = 0u64;
            let mut cats = 0u64;
            let mut rabbits = 0u64;

            for (i, charr) in pieces_str.chars().enumerate().take(BOARD_WIDTH * BOARD_HEIGHT) {
                if let Some((piece, is_p1)) = convert_char_to_piece(charr) {
                    let row_idx = i / BOARD_WIDTH;
                    let col_idx = i % BOARD_WIDTH;
                    let idx = (row_idx * BOARD_WIDTH + col_idx) as u8;
                    let square = Square::from_index(idx);
                    let square_bit = square.as_bit_board();

                    match piece {
                        Piece::Elephant => elephants |= square_bit,
                        Piece::Camel => camels |= square_bit,
                        Piece::Horse => horses |= square_bit,
                        Piece::Dog => dogs |= square_bit,
                        Piece::Cat => cats |= square_bit,
                        Piece::Rabbit => rabbits |= square_bit,
                    }

                    if is_p1 {
                        p1_pieces |= square_bit;
                    }
                }
            }

            let piece_board = PieceBoard::new(p1_pieces, elephants, camels, horses, dogs, cats, rabbits);
            return Ok(GameState::new(p1_turn_to_move, move_number, piece_board));
        }

        // Visual board format with | separators:
        // e.g. "2g\n +-...-+\n8| r r r ... |\n7| ..."
        let lines: Vec<_> = s.split('|').enumerate().filter(|(i, _)| i % 2 == 1).map(|(_, s)| s).collect();
        let regex = regex::Regex::new(r"^\s*(\d+)([gswb])").unwrap();

        let (move_number, p1_turn_to_move) = regex.captures(s.split('|').next().unwrap()).map_or((2, true), |c| {
            (c.get(1).unwrap().as_str().parse().unwrap(), c.get(2).unwrap().as_str() != "s" && c.get(2).unwrap().as_str() != "b")
        });

        let mut p1_pieces = 0;
        let mut elephants = 0;
        let mut camels = 0;
        let mut horses = 0;
        let mut dogs = 0;
        let mut cats = 0;
        let mut rabbits = 0;

        for (row_idx, line) in lines.iter().enumerate() {
            for (col_idx, charr) in line.chars().enumerate().filter(|(i, _)| i % 2 == 1).map(|(_, s)| s).enumerate() {
                let idx = (row_idx * BOARD_WIDTH + col_idx) as u8;
                let square = Square::from_index(idx);
                if let Some((piece, is_p1)) = convert_char_to_piece(charr) {
                    let square_bit = square.as_bit_board();

                    match piece {
                        Piece::Elephant => elephants |= square_bit,
                        Piece::Camel => camels |= square_bit,
                        Piece::Horse => horses |= square_bit,
                        Piece::Dog => dogs |= square_bit,
                        Piece::Cat => cats |= square_bit,
                        Piece::Rabbit => rabbits |= square_bit,
                    }

                    if is_p1 {
                        p1_pieces |= square_bit;
                    }
                }
            }
        }

        let piece_board = PieceBoard::new(p1_pieces, elephants, camels, horses, dogs, cats, rabbits);

        Ok(GameState::new(p1_turn_to_move, move_number, piece_board))
    }
}

fn convert_char_to_piece(c: char) -> Option<(Piece, bool)> {
    let is_p1 = c.is_uppercase();

    let piece = match c {
        'E' | 'e' => Some(Piece::Elephant),
        'M' | 'm' => Some(Piece::Camel),
        'H' | 'h' => Some(Piece::Horse),
        'D' | 'd' => Some(Piece::Dog),
        'C' | 'c' => Some(Piece::Cat),
        'R' | 'r' => Some(Piece::Rabbit),
        _ => None,
    };

    piece.map(|p| (p, is_p1))
}

fn is_p1_piece(square_bit: u64, piece_board: &PieceBoard) -> bool {
    let p1_piece_mask = piece_board.get_player_piece_mask(true);
    (square_bit & p1_piece_mask) != 0
}

pub fn convert_piece_to_letter(piece: &Piece, is_p1: bool) -> String {
    let letter = match piece {
        Piece::Elephant => "E",
        Piece::Camel => "M",
        Piece::Horse => "H",
        Piece::Dog => "D",
        Piece::Cat => "C",
        Piece::Rabbit => "R",
    };

    if is_p1 { letter.to_string() } else { letter.to_lowercase() }
}

#[cfg(test)]
mod tests {
    use super::super::{Piece, take_actions};
    use super::*;

    #[must_use]
    fn place_gold_pieces(game_state: GameState) -> GameState {
        take_actions!(game_state => e1, d1, a1, h1, c1, f1, b1, g1)
    }

    #[must_use]
    fn place_silver_pieces(game_state: GameState) -> GameState {
        take_actions!(game_state => e8, d8, a8, h8, c8, f8, b8, g8)
    }

    fn initial_play_state() -> GameState {
        let game_state = GameState::initial();
        let game_state = place_gold_pieces(game_state);
        place_silver_pieces(game_state)
    }

    #[test]
    fn test_gamestate_fromstr() {
        let game_state: GameState = "
            +-----------------+
            8|   r   r r   r   |
            7| m   h     e   c |
            6|   r x r r x r   |
            5| h   d     c   d |
            4| E   H         M |
            3|   R x R R H R   |
            2| D   C     C   D |
            1|   R   R R   R   |
            +-----------------+
               a b c d e f g h"
            .parse()
            .unwrap();

        let piece_board = game_state.get_piece_board();
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Rabbit, true),
            0b__01011010__00000000__01011010__00000000__00000000__00000000__00000000__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Cat, true),
            0b__00000000__00100100__00000000__00000000__00000000__00000000__00000000__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Dog, true),
            0b__00000000__10000001__00000000__00000000__00000000__00000000__00000000__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Horse, true),
            0b__00000000__00000000__00100000__00000100__00000000__00000000__00000000__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Camel, true),
            0b__00000000__00000000__00000000__10000000__00000000__00000000__00000000__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Elephant, true),
            0b__00000000__00000000__00000000__00000001__00000000__00000000__00000000__00000000
        );

        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Rabbit, false),
            0b__00000000__00000000__00000000__00000000__00000000__01011010__00000000__01011010
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Cat, false),
            0b__00000000__00000000__00000000__00000000__00100000__00000000__10000000__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Dog, false),
            0b__00000000__00000000__00000000__00000000__10000100__00000000__00000000__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Horse, false),
            0b__00000000__00000000__00000000__00000000__00000001__00000000__00000100__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Camel, false),
            0b__00000000__00000000__00000000__00000000__00000000__00000000__00000001__00000000
        );
        assert_eq!(
            piece_board.get_bits_for_piece(Piece::Elephant, false),
            0b__00000000__00000000__00000000__00000000__00000000__00000000__00100000__00000000
        );
    }

    #[test]
    fn test_gamestate_fromstr_move_number_default() {
        let game_state: GameState = "
            +-----------------+
            8|   r   r r   r   |
            7| m   h     e   c |
            6|   r x r r x r   |
            5| h   d     c   d |
            4| E   H         M |
            3|   R x R R H R   |
            2| D   C     C   D |
            1|   R   R R   R   |
            +-----------------+
               a b c d e f g h"
            .parse()
            .unwrap();

        assert_eq!(game_state.get_move_number(), 2);
        assert!(game_state.is_p1_turn_to_move());
    }

    #[test]
    fn test_gamestate_fromstr_move_number() {
        let game_state: GameState = "
            5s
            +-----------------+
            8|   r   r r   r   |
            7| m   h     e   c |
            6|   r x r r x r   |
            5| h   d     c   d |
            4| E   H         M |
            3|   R x R R H R   |
            2| D   C     C   D |
            1|   R   R R   R   |
            +-----------------+
               a b c d e f g h"
            .parse()
            .unwrap();

        assert_eq!(game_state.get_move_number(), 5);
        assert!(!game_state.is_p1_turn_to_move());
    }

    #[test]
    fn test_gamestate_fromstr_player() {
        let game_state: GameState = "
            176b
            +-----------------+
            8|   r   r r   r   |
            7| m   h     e   c |
            6|   r x r r x r   |
            5| h   d     c   d |
            4| E   H         M |
            3|   R x R R H R   |
            2| D   C     C   D |
            1|   R   R R   R   |
            +-----------------+
               a b c d e f g h"
            .parse()
            .unwrap();

        assert_eq!(game_state.get_move_number(), 176);
        assert!(!game_state.is_p1_turn_to_move());
    }

    #[test]
    fn test_gamestate_fromstr_player_w() {
        let game_state: GameState = "
            13w
            +-----------------+
            8|   r   r r   r   |
            7| m   h     e   c |
            6|   r x r r x r   |
            5| h   d     c   d |
            4| E   H         M |
            3|   R x R R H R   |
            2| D   C     C   D |
            1|   R   R R   R   |
            +-----------------+
               a b c d e f g h"
            .parse()
            .unwrap();

        assert_eq!(game_state.get_move_number(), 13);
        assert!(game_state.is_p1_turn_to_move());
    }

    #[test]
    fn test_gamestate_to_str() {
        let expected_output = "2g
 +-----------------+
8| h c d m e d c h |
7| r r r r r r r r |
6|     x     x     |
5|                 |
4|                 |
3|     x     x     |
2| R R R R R R R R |
1| H C D M E D C H |
 +-----------------+
   a b c d e f g h
";

        assert_eq!(format!("{}", initial_play_state()), expected_output);
    }

    #[test]
    fn test_gamestate_from_str_and_to_str() {
        let orig_str = "14g
 +-----------------+
8|   r   r r   r   |
7| m   h     e   c |
6|   r x r r x r   |
5| h   d     c   d |
4| E   H         M |
3|   R x R R x R   |
2| D   C     C   D |
1|   R   R R   R   |
 +-----------------+
   a b c d e f g h
";

        let game_state: GameState = orig_str.parse().unwrap();
        let new_str = format!("{}", game_state);

        assert_eq!(new_str, orig_str);
    }

    #[test]
    fn test_gamestate_fromstr_compact_initial() {
        // Compact AEI format equivalent of the initial play state
        // Row 8→1: hcdmedch / rrrrrrrr / (empty x4) / RRRRRRRR / HCDMEDCH
        let expected_output = "2g
 +-----------------+
8| h c d m e d c h |
7| r r r r r r r r |
6|     x     x     |
5|                 |
4|                 |
3|     x     x     |
2| R R R R R R R R |
1| H C D M E D C H |
 +-----------------+
   a b c d e f g h
";

        let game_state: GameState = "2g [hcdmedchrrrrrrrr                                RRRRRRRRHCDMEDCH]"
            .parse()
            .unwrap();
        assert_eq!(format!("{}", game_state), expected_output);
    }

    #[test]
    fn test_gamestate_fromstr_compact_no_move_number() {
        // "g [...]" with no move number should default move_number to 1, gold to move
        let game_state: GameState = "g [rredchrrrhrdcmrr                                MHHEDDCRRRRRRRRC]"
            .parse()
            .unwrap();
        assert_eq!(game_state.get_move_number(), 1);
        assert!(game_state.is_p1_turn_to_move());
        // Verify a sample of the piece placements
        let piece_board = game_state.get_piece_board();
        // Gold elephant should be at d2 (row index 6, col index 3 → square index 51)
        assert_ne!(piece_board.get_bits_for_piece(Piece::Elephant, true), 0);
        // Silver elephant should be at c8 (row index 0, col index 2 → square index 2)
        assert_ne!(piece_board.get_bits_for_piece(Piece::Elephant, false), 0);
    }

    #[test]
    fn test_gamestate_fromstr_compact_silver_to_move() {
        let game_state: GameState = "5s [hcdmedchrrrrrrrr                                RRRRRRRRHCDMEDCH]"
            .parse()
            .unwrap();
        assert_eq!(game_state.get_move_number(), 5);
        assert!(!game_state.is_p1_turn_to_move());
    }
}