bitstackchess 0.1.1

//! Check / checkmate / stalemate detection.
//! Exposes `is_in_check`, `has_any_move`, and `game_result`.

use crate::board::{PieceMapping, Occupied};
use crate::core::Move10;
use crate::engine::move_generator::MoveGenerator;

/// Color = 0 (White) or 1 (Black)
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Color {
    White = 0,
    Black = 1,
}

#[derive(Debug, PartialEq, Eq)]
pub enum GameResult {
    Checkmate(Color),  // `<Color>` is the winning side
    Stalemate,
    Ongoing,
}

/// Stateless helper for check/checkmate/stalemate logic.
pub struct CheckmateLogic;

impl CheckmateLogic {
    /// Return true if `king_pid` is under attack by any opposing piece.
    pub fn is_in_check(mapping: &PieceMapping, occupied: Occupied, king_pid: u8) -> bool {
        // 1) Locate the king’s square:
        let king_sq = match mapping.piece_square[king_pid as usize] {
            Some(sq) => sq,
            None => return false, // King not on board => not in check
        };
        let king_color = if king_pid < 16 { 0 } else { 1 };
        let (kr, kf) = ((king_sq >> 3) as i8, (king_sq & 7) as i8);

        // 2) Pawn‐attack offsets:
        //    If king is White, check for Black pawns on (kr+1, kf±1).
        //    If king is Black, check for White pawns on (kr−1, kf±1).
        let pawn_dir: i8 = if king_color == 0 { -1 } else { 1 };
        for df in &[-1, 1] {
            let pr = kr + pawn_dir;
            let pf = kf + df;
            if (0..8).contains(&pr) && (0..8).contains(&pf) {
                let psq = ((pr as u8) << 3) | (pf as u8);
                if let Some(pid) = mapping.who_on_square(psq) {
                    let is_pawn = (king_color == 0 && (16..24).contains(&pid))
                        || (king_color == 1 && (0..8).contains(&pid));
                    if is_pawn {
                        return true;
                    }
                }
            }
        }

        // 3) Knight‐attack offsets:
        const KNIGHT_OFFSETS: &[(i8, i8)] = &[
            ( 2,  1), ( 2, -1), (-2,  1), (-2, -1),
            ( 1,  2), ( 1, -2), (-1,  2), (-1, -2),
        ];
        for &(dr, df) in KNIGHT_OFFSETS {
            let nr = kr + dr;
            let nf = kf + df;
            if (0..8).contains(&nr) && (0..8).contains(&nf) {
                let nsq = ((nr as u8) << 3) | (nf as u8);
                if let Some(pid) = mapping.who_on_square(nsq) {
                    // Opponent’s knight PIDs are 10,11 for White or 26,27 for Black
                    if king_color == 0 && (26..=27).contains(&pid) {
                        return true;
                    }
                    if king_color == 1 && (10..=11).contains(&pid) {
                        return true;
                    }
                }
            }
        }

        // 4) Sliding‐piece rays (orthogonal & diagonal).
        //    If we see an opposing rook/queen on an orthogonal ray, or bishop/queen on diagonal, return true.
        const ORTHO_DIRS: &[(i8, i8)] = &[( 1,  0), (-1,  0), ( 0,  1), ( 0, -1)];
        const DIAG_DIRS:  &[(i8, i8)] = &[( 1,  1), ( 1, -1), (-1,  1), (-1, -1)];
        // Orthogonal rays:
        for &(dr, df) in ORTHO_DIRS.iter() {
            let mut nr = kr + dr;
            let mut nf = kf + df;
            while (0..8).contains(&nr) && (0..8).contains(&nf) {
                let sq = ((nr as u8) << 3) | (nf as u8);
                if let Some(pid) = mapping.who_on_square(sq) {
                    if color_of(pid) != king_color {
                        // Opponent: check if it’s rook or queen
                        if (pid == 8 || pid == 9 || pid == 24 || pid == 25)
                            || (pid == 14 || pid == 30)
                        {
                            return true;
                        }
                    }
                    break; // blocked by any piece
                }
                nr += dr;
                nf += df;
            }
        }
        // Diagonal rays:
        for &(dr, df) in DIAG_DIRS.iter() {
            let mut nr = kr + dr;
            let mut nf = kf + df;
            while (0..8).contains(&nr) && (0..8).contains(&nf) {
                let sq = ((nr as u8) << 3) | (nf as u8);
                if let Some(pid) = mapping.who_on_square(sq) {
                    if color_of(pid) != king_color {
                        // Opponent: check if it’s bishop or queen
                        if (pid == 12 || pid == 13 || pid == 28 || pid == 29)
                            || (pid == 14 || pid == 30)
                        {
                            return true;
                        }
                    }
                    break;
                }
                nr += dr;
                nf += df;
            }
        }

        // 5) King adjacency (opposing king cannot be adjacent).
        const KING_OFFSETS: &[(i8, i8)] = &[
            ( 1,  0), ( 1,  1), ( 0,  1), (-1,  1),
            (-1,  0), (-1, -1), ( 0, -1), ( 1, -1),
        ];
        for &(dr, df) in KING_OFFSETS.iter() {
            let nr = kr + dr;
            let nf = kf + df;
            if (0..8).contains(&nr) && (0..8).contains(&nf) {
                let nsq = ((nr as u8) << 3) | (nf as u8);
                if let Some(pid) = mapping.who_on_square(nsq) {
                    if king_color == 0 && pid == 31 {
                        return true;
                    }
                    if king_color == 1 && pid == 15 {
                        return true;
                    }
                }
            }
        }

        false
    }

    /// Return true if `side` has any legal move (i.e. a pseudo‐legal move that does not leave
    /// its own king in check). Uses `MoveGenerator` to produce pseudo‐legal moves and filters.
    pub fn has_any_move(
        mapping: &PieceMapping,
        occupied: Occupied,
        captured_bits: u32,
        side: Color,
        planes: &crate::core::MovePlanes,
        ply: usize,
        en_passant_target: Option<u8>,
    ) -> bool {
        // 1) Generate all pseudo-legal:
        let side_u8 = match side {
            Color::White => 0,
            Color::Black => 1,
        };
        let pseudo = MoveGenerator::generate(mapping, occupied, captured_bits, en_passant_target, side_u8);

        // 2) Test each by making a temporary clone of mapping/occupied/captured_bits/ep_target:
        for mv in pseudo {
            // Clone current board state:
            let mut temp_mapping = mapping.clone();
            let mut temp_occupied = occupied;
            let mut temp_captured = captured_bits;
            let mut temp_ep = en_passant_target;
            let pid_within = mv.piece_id();
            let global_pid = 16 * side_u8 + pid_within;
            let src_sq = temp_mapping.piece_square[global_pid as usize].unwrap() as usize;
            let dst_sq = mv.dest() as usize;

            // Handle castling on temp:
            let is_castle = pid_within == 15
                && ((side_u8 == 0 && ((src_sq == 4 && dst_sq == 6) || (src_sq == 4 && dst_sq == 2)))
                    || (side_u8 == 1 && ((src_sq == 60 && dst_sq == 62) || (src_sq == 60 && dst_sq == 58))));
            if is_castle {
                if side_u8 == 0 && dst_sq == 6 {
                    crate::rules::castling::CastlingLogic::do_white_kingside(&mut temp_mapping, &mut temp_occupied);
                    temp_ep = None;
                } else if side_u8 == 0 && dst_sq == 2 {
                    crate::rules::castling::CastlingLogic::do_white_queenside(&mut temp_mapping, &mut temp_occupied);
                    temp_ep = None;
                } else if side_u8 == 1 && dst_sq == 62 {
                    crate::rules::castling::CastlingLogic::do_black_kingside(&mut temp_mapping, &mut temp_occupied);
                    temp_ep = None;
                } else if side_u8 == 1 && dst_sq == 58 {
                    crate::rules::castling::CastlingLogic::do_black_queenside(&mut temp_mapping, &mut temp_occupied);
                    temp_ep = None;
                }
            } else {
                // En-passant capture on temp:
                if (global_pid % 16) < 8 {
                    let (sr, sf) = ((src_sq >> 3) as i8, (src_sq & 7) as i8);
                    let (dr, df) = (((dst_sq >> 3) as i8), ((dst_sq & 7) as i8));
                    if (df - sf).abs() == 1 && (dr - sr).abs() == 1 {
                        if let Some(ep) = temp_ep {
                            if ep as usize == dst_sq {
                                let cap_sq = if side_u8 == 0 { dst_sq - 8 } else { dst_sq + 8 };
                                if let Some(opp) = temp_mapping.who_on_square(cap_sq as u8) {
                                    temp_mapping.remove_piece(opp);
                                    temp_occupied &= !(1u64 << cap_sq);
                                    temp_captured |= 1u32 << (opp as u32);
                                }
                            }
                        }
                    }
                }
                // Normal capture on temp:
                if let Some(opp) = temp_mapping.who_on_square(dst_sq as u8) {
                    temp_mapping.remove_piece(opp);
                    temp_occupied &= !(1u64 << dst_sq);
                    temp_captured |= 1u32 << (opp as u32);
                }
                // Move piece on temp:
                temp_mapping.move_piece(global_pid as u8, dst_sq as u8);
                temp_occupied &= !(1u64 << (src_sq as u64));
                temp_occupied |= 1u64 << dst_sq;
                // Recompute temp_ep (pawn double-step?):
                if (global_pid % 16) < 8 {
                    if let Some(new_ep) =
                        crate::rules::en_passant::EnPassantLogic::compute_ep_target(src_sq as u8, dst_sq as u8, side_u8)
                    {
                        temp_ep = Some(new_ep);
                    } else {
                        temp_ep = None;
                    }
                } else {
                    temp_ep = None;
                }
            }

            // 3) Check if own king is in check on temp:
            let king_pid = if side_u8 == 0 { 15 } else { 31 };
            if !Self::is_in_check(&temp_mapping, temp_occupied, king_pid) {
                return true;
            }
        }

        false
    }

    /// Return GameResult (Checkmate, Stalemate, or Ongoing) for current position.
    pub fn game_result(
        mapping: &PieceMapping,
        occupied: Occupied,
        side_to_move: Color,
        planes: &crate::core::MovePlanes,
        ply: usize,
        en_passant_target: Option<u8>,
    ) -> GameResult {
        let king_pid = if side_to_move == Color::White { 15 } else { 31 };
        let in_check = Self::is_in_check(mapping, occupied, king_pid);
        let has_move = Self::has_any_move(
            mapping,
            occupied,
            0, // captured_bits not needed for legality test
            side_to_move,
            planes,
            ply,
            en_passant_target,
        );

        if in_check && !has_move {
            // Checkmate → opponent wins
            let winner = if side_to_move == Color::White {
                Color::Black
            } else {
                Color::White
            };
            GameResult::Checkmate(winner)
        } else if !in_check && !has_move {
            GameResult::Stalemate
        } else {
            GameResult::Ongoing
        }
    }
}

/// Return side (0=White, 1=Black) from PID.
#[inline]
fn color_of(pid: u8) -> u8 {
    if pid < 16 { 0 } else { 1 }
}