src/position.rs

use crate::types::*;
use crate::bitboard::{Bitboard, EMPTY, get_bit, pop_lsb};

#[derive(Clone)]
pub struct Position {
    pub pieces: [Bitboard; 7],
    pub colors: [Bitboard; 2],
    pub side_to_move: Color,
    pub castling_rights: u8,
    pub ep_square: Square,
    pub halfmove_clock: u32,
    pub fullmove_number: u32,
    pub history: Vec<UndoState>,
}

#[derive(Clone)]
pub struct UndoState {
    pub castling_rights: u8,
    pub ep_square: Square,
    pub halfmove_clock: u32,
    pub captured: PieceType,
}

impl Position {
    pub fn new() -> Self {
        let mut pos = Position {
            pieces: [EMPTY; 7],
            colors: [EMPTY; 2],
            side_to_move: Color::White,
            castling_rights: 0,
            ep_square: Square::None,
            halfmove_clock: 0,
            fullmove_number: 1,
            history: Vec::new(),
        };
        pos.set_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1");
        pos
    }

    pub fn get_piece_at(&self, sq: Square) -> PieceType {
        for p in 1..=6 {
            if get_bit(self.pieces[p], sq) {
                return unsafe { std::mem::transmute(p as u8) };
            }
        }
        PieceType::None
    }

    pub fn set_fen(&mut self, fen: &str) {
        let tokens: Vec<&str> = fen.split_whitespace().collect();
        if tokens.is_empty() { return; }

        for p in &mut self.pieces { *p = EMPTY; }
        for c in &mut self.colors { *c = EMPTY; }

        let mut rank = 7;
        let mut file = 0;
        
        for c in tokens[0].chars() {
            if c == '/' {
                rank -= 1;
                file = 0;
            } else if let Some(digit) = c.to_digit(10) {
                file += digit as u8;
            } else {
                let sq = Square::from_u8(rank * 8 + file);
                let color = if c.is_uppercase() { Color::White } else { Color::Black };
                let p_type = match c.to_ascii_lowercase() {
                    'p' => PieceType::Pawn,
                    'n' => PieceType::Knight,
                    'b' => PieceType::Bishop,
                    'r' => PieceType::Rook,
                    'q' => PieceType::Queen,
                    'k' => PieceType::King,
                    _ => PieceType::None,
                };
                
                self.pieces[p_type as usize] |= 1 << (sq as u8);
                self.colors[color as usize] |= 1 << (sq as u8);
                file += 1;
            }
        }

        self.side_to_move = if tokens.get(1).copied().unwrap_or("w") == "w" { Color::White } else { Color::Black };

        self.castling_rights = 0;
        if let Some(&c_str) = tokens.get(2) {
            if c_str != "-" {
                for c in c_str.chars() {
                    match c {
                        'K' => self.castling_rights |= CastlingRights::WHITE_OO,
                        'Q' => self.castling_rights |= CastlingRights::WHITE_OOO,
                        'k' => self.castling_rights |= CastlingRights::BLACK_OO,
                        'q' => self.castling_rights |= CastlingRights::BLACK_OOO,
                        _ => {}
                    }
                }
            }
        }

        self.ep_square = Square::None;
        if let Some(&ep_str) = tokens.get(3) {
            if ep_str != "-" && ep_str.len() == 2 {
                let chars: Vec<char> = ep_str.chars().collect();
                let f = (chars[0] as u8) - b'a';
                let r = (chars[1] as u8) - b'1';
                self.ep_square = Square::from_u8(r * 8 + f);
            }
        }

        self.halfmove_clock = tokens.get(4).and_then(|s| s.parse().ok()).unwrap_or(0);
        self.fullmove_number = tokens.get(5).and_then(|s| s.parse().ok()).unwrap_or(1);
    }

    pub fn do_move(&mut self, m: Move) -> bool {
        let from = m.from();
        let to = m.to();
        let promo = m.promotion();
        let us = self.side_to_move;
        let them = us.flip();
        
        let p_type = self.get_piece_at(from);
        let captured = self.get_piece_at(to);

        self.history.push(UndoState {
            castling_rights: self.castling_rights,
            ep_square: self.ep_square,
            halfmove_clock: self.halfmove_clock,
            captured,
        });

        let from_bb = 1 << (from as u8);
        let to_bb = 1 << (to as u8);

        self.pieces[p_type as usize] ^= from_bb;
        self.colors[us as usize] ^= from_bb;

        if captured != PieceType::None {
            self.pieces[captured as usize] ^= to_bb;
            self.colors[them as usize] ^= to_bb;
            self.halfmove_clock = 0;
        } else {
            self.halfmove_clock += 1;
        }

        if p_type == PieceType::Pawn {
            self.halfmove_clock = 0;
            if promo != PieceType::None {
                self.pieces[promo as usize] |= to_bb;
            } else {
                self.pieces[p_type as usize] |= to_bb;
            }
            
            if to == self.ep_square {
                let ep_cap_sq = Square::from_u8(if us == Color::White { (to as u8) - 8 } else { (to as u8) + 8 });
                self.pieces[PieceType::Pawn as usize] ^= 1 << (ep_cap_sq as u8);
                self.colors[them as usize] ^= 1 << (ep_cap_sq as u8);
            }
            
            if (from as i8 - to as i8).abs() == 16 {
                self.ep_square = Square::from_u8(if us == Color::White { (from as u8) + 8 } else { (from as u8) - 8 });
            } else {
                self.ep_square = Square::None;
            }
        } else {
            self.pieces[p_type as usize] |= to_bb;
            self.ep_square = Square::None;
        }

        self.colors[us as usize] |= to_bb;

        if p_type == PieceType::King {
            if us == Color::White {
                if from == Square::E1 && to == Square::G1 {
                    self.pieces[PieceType::Rook as usize] ^= (1 << (Square::H1 as u8)) | (1 << (Square::F1 as u8));
                    self.colors[Color::White as usize] ^= (1 << (Square::H1 as u8)) | (1 << (Square::F1 as u8));
                } else if from == Square::E1 && to == Square::C1 {
                    self.pieces[PieceType::Rook as usize] ^= (1 << (Square::A1 as u8)) | (1 << (Square::D1 as u8));
                    self.colors[Color::White as usize] ^= (1 << (Square::A1 as u8)) | (1 << (Square::D1 as u8));
                }
                self.castling_rights &= !CastlingRights::ANY_WHITE;
            } else {
                if from == Square::E8 && to == Square::G8 {
                    self.pieces[PieceType::Rook as usize] ^= (1 << (Square::H8 as u8)) | (1 << (Square::F8 as u8));
                    self.colors[Color::Black as usize] ^= (1 << (Square::H8 as u8)) | (1 << (Square::F8 as u8));
                } else if from == Square::E8 && to == Square::C8 {
                    self.pieces[PieceType::Rook as usize] ^= (1 << (Square::A8 as u8)) | (1 << (Square::D8 as u8));
                    self.colors[Color::Black as usize] ^= (1 << (Square::A8 as u8)) | (1 << (Square::D8 as u8));
                }
                self.castling_rights &= !CastlingRights::ANY_BLACK;
            }
        }

        if from == Square::A1 || to == Square::A1 { self.castling_rights &= !CastlingRights::WHITE_OOO; }
        if from == Square::H1 || to == Square::H1 { self.castling_rights &= !CastlingRights::WHITE_OO; }
        if from == Square::A8 || to == Square::A8 { self.castling_rights &= !CastlingRights::BLACK_OOO; }
        if from == Square::H8 || to == Square::H8 { self.castling_rights &= !CastlingRights::BLACK_OO; }

        if us == Color::Black {
            self.fullmove_number += 1;
        }
        self.side_to_move = them;

        let king_sq = pop_lsb(&mut (self.pieces[PieceType::King as usize] & self.colors[us as usize]));
        if self.is_square_attacked(Square::from_u8(king_sq), us) {
            self.undo_move(m);
            return false;
        }

        true
    }

    pub fn undo_move(&mut self, m: Move) {
        let state = self.history.pop().unwrap();
        self.side_to_move = self.side_to_move.flip();
        let us = self.side_to_move;
        let them = us.flip();

        let from = m.from();
        let to = m.to();
        let promo = m.promotion();
        
        let p_type = if promo != PieceType::None { PieceType::Pawn } else { self.get_piece_at(to) };
        
        let from_bb = 1 << (from as u8);
        let to_bb = 1 << (to as u8);

        self.pieces[p_type as usize] |= from_bb;
        self.colors[us as usize] |= from_bb;

        if promo != PieceType::None {
            self.pieces[promo as usize] ^= to_bb;
        } else {
            self.pieces[p_type as usize] ^= to_bb;
        }
        self.colors[us as usize] ^= to_bb;

        if state.captured != PieceType::None {
            self.pieces[state.captured as usize] |= to_bb;
            self.colors[them as usize] |= to_bb;
        }

        if p_type == PieceType::Pawn && to == state.ep_square {
            let ep_cap_sq = Square::from_u8(if us == Color::White { (to as u8) - 8 } else { (to as u8) + 8 });
            self.pieces[PieceType::Pawn as usize] |= 1 << (ep_cap_sq as u8);
            self.colors[them as usize] |= 1 << (ep_cap_sq as u8);
        }

        if p_type == PieceType::King {
            if us == Color::White {
                if from == Square::E1 && to == Square::G1 {
                    self.pieces[PieceType::Rook as usize] ^= (1 << (Square::H1 as u8)) | (1 << (Square::F1 as u8));
                    self.colors[Color::White as usize] ^= (1 << (Square::H1 as u8)) | (1 << (Square::F1 as u8));
                } else if from == Square::E1 && to == Square::C1 {
                    self.pieces[PieceType::Rook as usize] ^= (1 << (Square::A1 as u8)) | (1 << (Square::D1 as u8));
                    self.colors[Color::White as usize] ^= (1 << (Square::A1 as u8)) | (1 << (Square::D1 as u8));
                }
            } else {
                if from == Square::E8 && to == Square::G8 {
                    self.pieces[PieceType::Rook as usize] ^= (1 << (Square::H8 as u8)) | (1 << (Square::F8 as u8));
                    self.colors[Color::Black as usize] ^= (1 << (Square::H8 as u8)) | (1 << (Square::F8 as u8));
                } else if from == Square::E8 && to == Square::C8 {
                    self.pieces[PieceType::Rook as usize] ^= (1 << (Square::A8 as u8)) | (1 << (Square::D8 as u8));
                    self.colors[Color::Black as usize] ^= (1 << (Square::A8 as u8)) | (1 << (Square::D8 as u8));
                }
            }
        }

        self.castling_rights = state.castling_rights;
        self.ep_square = state.ep_square;
        self.halfmove_clock = state.halfmove_clock;
        if us == Color::Black {
            self.fullmove_number -= 1;
        }
    }

    pub fn is_square_attacked(&self, sq: Square, side_by: Color) -> bool {
        let occ = self.colors[0] | self.colors[1];
        
        let knight_attacks = crate::movegen::knight_mask(sq);
        if (knight_attacks & self.pieces[PieceType::Knight as usize] & self.colors[side_by as usize]) != 0 { return true; }

        let king_attacks = crate::movegen::king_mask(sq);
        if (king_attacks & self.pieces[PieceType::King as usize] & self.colors[side_by as usize]) != 0 { return true; }

        let pawn_attacks = crate::movegen::pawn_attack_mask(sq, side_by.flip());
        if (pawn_attacks & self.pieces[PieceType::Pawn as usize] & self.colors[side_by as usize]) != 0 { return true; }

        if (crate::movegen::sliding_attacks(sq, occ, true) & (self.pieces[PieceType::Bishop as usize] | self.pieces[PieceType::Queen as usize]) & self.colors[side_by as usize]) != 0 { return true; }
        if (crate::movegen::sliding_attacks(sq, occ, false) & (self.pieces[PieceType::Rook as usize] | self.pieces[PieceType::Queen as usize]) & self.colors[side_by as usize]) != 0 { return true; }

        false
    }
}