use crate::{
    board::{Board, explode_board, squares::*},
    game::{Game, Team},
    table::{Bound, TTEntry, TTable},
};

/// Contains all corner squares
const CORNERS: Board = A1 | A8 | H1 | H8;
/// Contains all edge squares
const EDGES: Board = A2
    | A3
    | A4
    | A5
    | A6
    | A7
    | B1
    | B8
    | C1
    | C8
    | D1
    | D8
    | E1
    | E8
    | F1
    | F8
    | G1
    | G8
    | H2
    | H3
    | H4
    | H5
    | H6
    | H7;

#[derive(PartialEq, Eq, PartialOrd, Ord)]
/// Represents the _value_ of a move. Some moves at face value
/// better than others.
enum MoveRank {
    Corner(Board),
    Edge(Board),
    Other(Board),
}

impl From<Board> for MoveRank {
    fn from(value: Board) -> Self {
        // Do bitwise operations to check if we have a
        // corner or edge move.
        if value & CORNERS > 0 {
            Self::Corner(value)
        } else if value & EDGES > 0 {
            Self::Edge(value)
        } else {
            Self::Other(value)
        }
    }
}
impl MoveRank {
    /// Unwrap underlying move value out of rank structure
    fn into_inner(self) -> Board {
        match self {
            Self::Corner(m) => m,
            Self::Edge(m) => m,
            Self::Other(m) => m,
        }
    }
}

/// Using alpha-beta pruning and the minimax algorithm, determine the best move
/// for a game with a recursion depth of `depth`.
///
/// We use a very simple evaluation heuristic: (Black squares - White squares).
pub fn alphabeta(
    mut game: Game,
    depth: u8,
    mut alpha: i8,
    mut beta: i8,
    tt: &mut TTable,
) -> (Board, i8, u64) {
    let mut num_moves = 0;
    // if we reach our maximum recursion depth, return evaluation
    if depth == 0 {
        return (0, game.score().diff(), num_moves);
    }

    let moves = game.available();
    if moves == 0 {
        // if no move, skip and continue recursion
        // this seems to technically introduce a bias against move-chains
        // that include skips. I haven't found it to be a big deal in play.
        game.skip();
        return (0, alphabeta(game, depth - 1, alpha, beta, tt).1, num_moves);
    }

    // just initially assume that the best move is no move at all. This will
    // inevitably be corrected.
    let mut best_move: Board = 0;
    // we initially rank moves based on a couple basic heuristics:
    // - corner pieces are best
    // - edge pieces are great
    // - others considered last
    // This just allows us to prune the tree a bit more aggressively
    // since we're considering the "best" moves first.
    // We do this by mapping moves to ranked moves and then sorting.
    let mut moves = explode_board(moves).map(MoveRank::from).collect::<Vec<_>>();
    moves.sort_unstable();
    let mut moves = moves
        .into_iter()
        .map(MoveRank::into_inner)
        .collect::<Vec<_>>();

    // copy our existing alpha/beta for the sake of classifying bounds
    let original_alpha = alpha;
    let original_beta = beta;

    // the brilliance here is that even if we don't have a perfect value
    // computed already, the imperfect values still help us get to better values
    // quicker.
    match tt.get(game.hash) {
        Some(entry) if entry.depth >= depth => {
            match entry.bound {
                // if we know this is exact, trust it without question
                Bound::Exact => return (entry.best_move, entry.evaluation, num_moves),
                // if we have lower or upper bounds that are more precise than
                // our existing alpha and beta values, accept the ones found in
                // the cache.
                Bound::Lower => alpha = alpha.max(entry.evaluation),
                Bound::Upper => beta = beta.min(entry.evaluation),
            }
            // if we have collapsed the window between alpha and beta, just
            // accept the cached entry.
            if alpha >= beta {
                return (entry.best_move, entry.evaluation, num_moves);
            }

            // otherwise, if our best move is available, move it to the front
            if let Some(best_move_idx) = moves.iter().position(|m| *m == entry.best_move) {
                moves[..=best_move_idx].rotate_right(1);
            }
        }
        Some(entry) => {
            // otherwise, if our best move is available, move it to the front
            if let Some(best_move_idx) = moves.iter().position(|m| *m == entry.best_move) {
                moves[..=best_move_idx].rotate_right(1);
            }
        }
        None => {}
    }

    num_moves = moves.len() as u64;

    // I just establish a convention of maximizing for black and minimizing for white.
    // I'm not sure if that's conventional or not, but it's what I chose.
    match game.current_team {
        Team::Black => {
            for mv in moves {
                let mut g = game.clone();
                g.play(mv);
                // maximize for the evaluation of subsequent moves
                let (_, evaluation, num_moves_sub) = alphabeta(g, depth - 1, alpha, beta, tt);
                num_moves += num_moves_sub;
                // if our evaluated move is superior to the alpha, update
                // it.
                if evaluation > alpha {
                    alpha = evaluation;
                    best_move = mv;
                };
                // if our beta is less than alpha, prune the node.
                if beta <= alpha {
                    break;
                }
            }
            let bound = if alpha >= beta {
                Bound::Lower
            } else if alpha <= original_alpha {
                Bound::Upper
            } else {
                // i.e. alpha < beta || alpha < original_alpha
                Bound::Exact
            };
            tt.store(TTEntry {
                depth,
                evaluation: alpha,
                hash: game.hash,
                bound,
                best_move,
            });
            (best_move, alpha, num_moves)
        }
        Team::White => {
            for mv in moves {
                let mut g = game.clone();
                g.play(mv);
                // minimize for the evaluation of subsequent moves
                let (_, evaluation, num_moves_sub) = alphabeta(g, depth - 1, alpha, beta, tt);
                num_moves += num_moves_sub;
                // if our evaluated move produces lower eval than the beta,
                // update beta.
                if evaluation < beta {
                    beta = evaluation;
                    best_move = mv;
                };
                // if our beta is less than alpha, prune the node.
                if beta <= alpha {
                    break;
                }
            }
            let bound = if beta <= alpha {
                Bound::Upper
            } else if beta >= original_beta {
                Bound::Lower
            } else {
                Bound::Exact
            };
            tt.store(TTEntry {
                depth,
                evaluation: beta,
                hash: game.hash,
                bound,
                best_move,
            });
            (best_move, beta, num_moves)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::board::{BitBoard, Board, Score};
    use crate::game::{Game, Team};

    use rand::prelude::IndexedRandom;
    use rand::rngs::StdRng;
    use rand::{Rng, SeedableRng};

    /// An AI player that only makes random moves
    fn random_move(game: &Game, rng: &mut impl Rng) -> Board {
        let moves = explode_board(game.available()).collect::<Vec<_>>();
        *moves.choose(rng).unwrap()
    }

    fn game_after(moves: &[Board]) -> Game {
        let mut game = Game::default();
        for &mv in moves {
            game.safe_play(mv).expect("Move should be valid");
        }
        game
    }

    fn assert_ai_move_is_legal(game: &Game, depth: u8) -> Board {
        let available = game.available();
        let mut tt = TTable::with_mb(2);
        let best_move = alphabeta(game.clone(), depth, i8::MIN + 1, i8::MAX - 1, &mut tt).0;
        assert_ne!(best_move, 0, "AI should return a move when one exists");
        assert_eq!(
            best_move & available,
            best_move,
            "AI returned an illegal move"
        );
        best_move
    }

    #[test]
    // just a sanity check to ensure that my AI performs up to snuff with another popular engine
    fn opening() {
        let mut game = Game::default();
        let mut tt = TTable::with_mb(24);
        game.play(D3);
        let (best_move, _, _) = alphabeta(game.clone(), 14, i8::MIN + 1, i8::MAX - 1, &mut tt);
        assert_eq!(best_move, C3);
    }

    #[test]
    fn ai_returns_legal_moves_across_curated_positions() {
        let cases = vec![
            (game_after(&[]), 4),
            (game_after(&[D3]), 4),
            (game_after(&[F5, F6, E6, F4]), 4),
        ];

        for (game, depth) in cases {
            let available = game.available();
            if available == 0 {
                continue;
            }
            let mv = assert_ai_move_is_legal(&game, depth);
            assert_ne!(mv, 0);
        }
    }

    #[test]
    fn ai_prefers_forced_corner() {
        let board = BitBoard::from_jon("5bw//////").expect("Valid board");
        let game = Game::from_parts(Team::Black, board);
        assert_eq!(game.available(), H8);
        let mv = assert_ai_move_is_legal(&game, 3);
        assert_eq!(mv, H8);
    }

    #[test]
    fn ai_passes_when_no_moves_exist() {
        let board = BitBoard::from_jon("wwwwwwww/wwwwwwww/////").expect("Valid board");
        let mut tt = TTable::with_mb(2);
        let game = Game::from_parts(Team::Black, board);
        assert_eq!(game.available(), 0);
        let (mv, eval, _) = alphabeta(game.clone(), 4, i8::MIN + 1, i8::MAX - 1, &mut tt);
        assert_eq!(mv, 0);
        assert_eq!(eval, game.score().diff());
    }

    #[test]
    fn tt_exact_root_hit_eliminates_repeat_search() {
        let game = Game::default();
        let mut tt = TTable::with_mb(2);

        let (best_move, eval, first_considered) =
            alphabeta(game.clone(), 1, i8::MIN + 1, i8::MAX - 1, &mut tt);
        assert!(first_considered > 0);

        let (cached_move, cached_eval, second_considered) =
            alphabeta(game.clone(), 1, i8::MIN + 1, i8::MAX - 1, &mut tt);

        assert_eq!(cached_move, best_move);
        assert_eq!(cached_eval, eval);
        assert_eq!(second_considered, 0);
    }

    #[test]
    fn tt_lower_bound_hit_still_searches_with_wide_window() {
        let game = Game::default();
        let mut tt = TTable::with_mb(2);

        tt.store(TTEntry {
            bound: Bound::Lower,
            evaluation: 0,
            depth: 1,
            best_move: D3,
            hash: game.hash,
        });

        let (_, _, considered) = alphabeta(game.clone(), 1, i8::MIN + 1, i8::MAX - 1, &mut tt);
        assert!(considered > 0);
    }

    #[test]
    fn tt_upper_bound_hit_still_searches_with_wide_window() {
        let game = Game::default();
        let mut tt = TTable::with_mb(2);

        tt.store(TTEntry {
            bound: Bound::Upper,
            evaluation: 0,
            depth: 1,
            best_move: D3,
            hash: game.hash,
        });

        let (_, _, considered) = alphabeta(game.clone(), 1, i8::MIN + 1, i8::MAX - 1, &mut tt);
        assert!(considered > 0);
    }

    // I found that, despite the AI clobbering me, the AI could not
    // compete with itself very well. I'm honestly not quite sure why that is.
    #[test]
    fn ai_beats_random() {
        // just contains pairings of starting_team and seed value
        let cases = vec![
            (Team::Black, 1231293),
            (Team::White, 491823),
            (Team::White, 12931),
            (Team::Black, 982983713),
            (Team::Black, 123),
            (Team::White, 87132895),
        ];
        let mut tt = TTable::with_mb(2);

        for (team, seed) in cases {
            let mut rng = StdRng::seed_from_u64(seed);
            let mut game = Game::default();
            if team != Team::Black {
                let mv = random_move(&game, &mut rng);
                game.play(mv);
            }

            while !game.is_complete() {
                if game.available() == 0 {
                    game.skip();
                    continue;
                }
                let mv = if game.current_team == team {
                    alphabeta(game.clone(), 8, i8::MIN + 1, i8::MAX - 1, &mut tt).0
                } else {
                    random_move(&game, &mut rng)
                };
                assert_eq!(mv & game.available(), mv, "AI generated an illegal move");
                game.play(mv);
            }

            assert!(
                match (team, game.score()) {
                    (Team::Black, Score(b, w)) => b - w,
                    (Team::White, Score(b, w)) => w - b,
                } > 4,
                "game with seed {} and team {:?} failed to win by 4 points or more.",
                seed,
                team
            );
        }
    }

    #[test]
    fn move_ordering() {
        let mv = A1 | A8 | C3 | D5 | A4;
        let mut moves = explode_board(mv).map(MoveRank::from).collect::<Vec<_>>();
        moves.sort();
        let moves = moves
            .into_iter()
            .map(MoveRank::into_inner)
            .collect::<Vec<_>>();
        assert_eq!(moves, vec![A1, A8, A4, C3, D5]);
    }
}