Chess engine in zig
const std = @import("std");

const Str = []const u8;

const Chess = @import("Chess.zig");
const Attacks = @import("Attacks.zig").Attacks;
const Search = @import("Search").Search;

const zobrist = @import("zobrist.zig");

/// SIZE is row and column size of the chessboard.
pub const SIZE = 8;

pub const SquareType = u6;
pub const Square = enum(SquareType) {
    // zig fmt: off
    a8, b8, c8, d8, e8, f8, g8, h8,
    a7, b7, c7, d7, e7, f7, g7, h7,
    a6, b6, c6, d6, e6, f6, g6, h6,
    a5, b5, c5, d5, e5, f5, g5, h5,
    a4, b4, c4, d4, e4, f4, g4, h4,
    a3, b3, c3, d3, e3, f3, g3, h3,
    a2, b2, c2, d2, e2, f2, g2, h2,
    a1, b1, c1, d1, e1, f1, g1, h1,
    // zig fmt: on

    pub fn rank(self: @This()) u3 {
        return @as(u3, @intCast(@intFromEnum(self) / SIZE));
    }

    pub fn file(self: @This()) u3 {
        return @as(u3, @intCast(@intFromEnum(self) % SIZE));
    }
};

pub const BoardType = u64;
pub const BitBoard = packed struct(BoardType) {
    // zig fmt: off
    a8:bool=false, b8:bool=false, c8:bool=false, d8:bool=false, e8:bool=false, f8:bool=false, g8:bool=false, h8:bool=false,
    a7:bool=false, b7:bool=false, c7:bool=false, d7:bool=false, e7:bool=false, f7:bool=false, g7:bool=false, h7:bool=false,
    a6:bool=false, b6:bool=false, c6:bool=false, d6:bool=false, e6:bool=false, f6:bool=false, g6:bool=false, h6:bool=false,
    a5:bool=false, b5:bool=false, c5:bool=false, d5:bool=false, e5:bool=false, f5:bool=false, g5:bool=false, h5:bool=false,
    a4:bool=false, b4:bool=false, c4:bool=false, d4:bool=false, e4:bool=false, f4:bool=false, g4:bool=false, h4:bool=false,
    a3:bool=false, b3:bool=false, c3:bool=false, d3:bool=false, e3:bool=false, f3:bool=false, g3:bool=false, h3:bool=false,
    a2:bool=false, b2:bool=false, c2:bool=false, d2:bool=false, e2:bool=false, f2:bool=false, g2:bool=false, h2:bool=false,
    a1:bool=false, b1:bool=false, c1:bool=false, d1:bool=false, e1:bool=false, f1:bool=false, g1:bool=false, h1:bool=false,
    // zig fmt: on

    fn set(self: *@This(), s: Square) void {
        var ret: BoardType = @as(BoardType, @bitCast(self.*));
        ret |= @as(BoardType, 1) << @intFromEnum(s);
        self.* = @as(BitBoard, @bitCast(ret));
    }

    pub fn setSlice(self: *@This(), squares: []const Square) void {
        var ret = @as(BoardType, @bitCast(self.*));
        for (squares) |square| {
            ret |= @as(BoardType, 1) << @intFromEnum(square);
        }
        self.* = @as(BitBoard, @bitCast(ret));
    }

    pub fn isSet(self: @This(), s: Square) bool {
        var ret = @as(BoardType, @bitCast(self));
        return (ret & @as(BoardType, 1) << @intFromEnum(s) != 0);
    }

    fn pop(self: *@This(), s: Square) void {
        var ret = @as(BoardType, @bitCast(self.*));
        ret ^= @as(BoardType, 1) << @intFromEnum(s);
        self.* = @as(BitBoard, @bitCast(ret));
    }

    fn unSet(self: *@This(), s: Square) void {
        var ret = @as(BoardType, @bitCast(self.*));
        ret &= ~(@as(BoardType, 1) << @intFromEnum(s));
        self.* = @as(BitBoard, @bitCast(ret));
    }

    pub fn show(self: @This()) void {
        std.debug.print("\n", .{});
        std.debug.print("{s:>18}\n", .{"0 1 2 3 4 5 6 7"});
        std.debug.print("{s:>18}\n", .{"---------------"});
        var rank: usize = 0;
        while (rank < SIZE) : (rank += 1) {
            var file: usize = 0;
            while (file < SIZE) : (file += 1) {
                if (file == 0) std.debug.print("{d}| ", .{rank});

                const square = @as(SquareType, @intCast(rank * SIZE + file));
                const mask: usize = if (@as(BoardType, @bitCast(self)) & (@as(BoardType, 1) << square) != 0) 1 else 0;
                std.debug.print("{d} ", .{mask});
            }
            std.debug.print("|{d}\n", .{SIZE - rank});
        }
        std.debug.print("{s:>18}\n", .{"---------------"});
        std.debug.print("{s:>18}\n", .{"a b c d e f g h"});

        std.debug.print("\nBitBoard: {d}\n", .{@as(BoardType, @bitCast(self))});
    }
};

const Moves = enum {
    all,
    captures,
};

pub const MovePrio = struct {
    move: BitMove,
    score: isize = undefined,

    pub fn moreThan(context: void, a: @This(), b: @This()) bool {
        _ = context;
        if (a.score > b.score) return true;
        return false;
    }
};

pub const MoveList = std.BoundedArray(MovePrio, 128);

pub const BitMoveType = u24;
pub const BitMove = packed struct(BitMoveType) {
    source: Square,
    target: Square,
    piece: Chess.PE,
    prom: Chess.PE = .none,
    capture: bool = false,
    double: bool = false,
    enpassant: bool = false,
    castling: bool = false,

    pub fn show(self: @This()) void {
        // std.debug.print("source: {any}\n", .{self.source});
        // std.debug.print("target: {any}\n", .{self.target});
        // std.debug.print("piece: {any}\n", .{self.piece});
        // std.debug.print("promote: {any}\n", .{self.prom});
        // std.debug.print("capture: {any}\n", .{self.capture});
        // std.debug.print("double: {any}\n", .{self.double});
        // std.debug.print("enpassant: {any}\n", .{self.enpassant});
        // std.debug.print("castling: {any}\n", .{self.castling});
        std.debug.print("{s} {s}{s}\t", .{
            @tagName(self.piece),
            @tagName(self.source),
            @tagName(self.target),
        });
    }
};

pub const CastlingType = u4;
const Castling = packed struct(CastlingType) {
    WK: bool = false,
    WQ: bool = false,
    BK: bool = false,
    BQ: bool = false,

    //                           castling   move     in      in
    //                              right update     binary  decimal
    // king & rooks didn't move:     1111 & 1111  =  1111    15
    //        white king  moved:     1111 & 1100  =  1100    12
    //  white king's rook moved:     1111 & 1110  =  1110    14
    // white queen's rook moved:     1111 & 1101  =  1101    13
    //
    //         black king moved:     1111 & 0011  =  1011    3
    //  black king's rook moved:     1111 & 1011  =  1011    11
    // black queen's rook moved:     1111 & 0111  =  0111    7

    // zig fmt: off
    const CASTLING_RIGHTS: [64]CastlingType= .{
         7, 15, 15, 15,  3, 15, 15, 11,
        15, 15, 15, 15, 15, 15, 15, 15,
        15, 15, 15, 15, 15, 15, 15, 15,
        15, 15, 15, 15, 15, 15, 15, 15,
        15, 15, 15, 15, 15, 15, 15, 15,
        15, 15, 15, 15, 15, 15, 15, 15,
        15, 15, 15, 15, 15, 15, 15, 15,
        13, 15, 15, 15, 12, 15, 15, 14
    };
    // zig fmt: on

    fn update(self: *@This(), square: Square) void {
        self.* = @as(Castling, @bitCast(@as(CastlingType, @bitCast(self.*)) &
            CASTLING_RIGHTS[@intFromEnum(square)]));
    }
};

pub const GameState = struct {
    bitboards: [12]BitBoard = blk: {
        var bbs: [12]BitBoard = undefined;
        for (&bbs) |*bb| {
            bb.* = BitBoard{};
        }
        break :blk bbs;
    },
    // TODO: try with separate both occupancy
    occupancies: [@typeInfo(Chess.Colors).Enum.fields.len]BitBoard = blk: {
        var bbs: [2]BitBoard = undefined;
        for (&bbs) |*bb| {
            bb.* = BitBoard{};
        }
        break :blk bbs;
    },
    attacks: Attacks,
    side: Chess.Colors = .white,
    enpassant: ?Square = null,
    castling: Castling = .{},
    fifty: u7 = 0,
    hash: u64 = undefined,
    history: HashHistory = undefined,
    allocator: std.mem.Allocator,

    const HashHistoryContext = struct {
        pub const hash = hashFn;
        pub const eql = eqlFn;
        fn hashFn(self: @This(), key: u64) u32 {
            _ = self;
            return @as(u32, @truncate(key));
        }
        fn eqlFn(self: @This(), a: u64, b: u64, b_index: usize) bool {
            _ = self;
            _ = b_index;
            return a == b;
        }
    };
    const HashHistory = std.ArrayHashMap(u64, void, HashHistoryContext, false);

    fn reset(self: *@This()) void {
        for (&self.bitboards) |*bb| {
            bb.* = .{};
        }
        for (&self.occupancies) |*bb| {
            bb.* = .{};
        }
        self.castling = .{};
    }

    pub fn init(allocator: std.mem.Allocator, FEN: ?Str) !@This() {
        var gs = GameState{ .allocator = allocator, .attacks = Attacks.init() };
        gs.history = HashHistory.init(allocator);
        zobrist.init();
        if (FEN != null)
            try gs.parseFEN(FEN.?)
        else
            try gs.parseFEN("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1");
        return gs;
    }

    pub fn deinit(self: *@This()) void {
        self.history.deinit();
    }

    pub fn backup(self: *const @This(), new_gs: *@This()) void {
        for (self.bitboards, 0..) |bb, idx| {
            new_gs.bitboards[idx] = bb;
        }
        for (self.occupancies, 0..) |occ, idx| {
            new_gs.occupancies[idx] = occ;
        }
        new_gs.side = self.side;
        new_gs.enpassant = self.enpassant;
        new_gs.castling = self.castling;
        new_gs.hash = self.hash;
        new_gs.fifty = self.fifty;
    }

    pub fn restore(self: *@This(), bck: *const @This()) void {
        for (bck.bitboards, 0..) |bb, idx| {
            self.bitboards[idx] = bb;
        }
        for (bck.occupancies, 0..) |occ, idx| {
            self.occupancies[idx] = occ;
        }
        self.side = bck.side;
        self.enpassant = bck.enpassant;
        self.castling = bck.castling;
        self.hash = bck.hash;
        self.fifty = bck.fifty;
    }

    pub fn occupBoth(self: @This()) BoardType {
        return @as(BoardType, @bitCast(self.occupancies[@intFromEnum(Chess.Colors.white)])) |
            @as(BoardType, @bitCast(self.occupancies[@intFromEnum(Chess.Colors.black)]));
    }

    pub fn show(self: @This()) void {
        std.debug.print("\n", .{});
        var rank: usize = 0;
        while (rank < SIZE) : (rank += 1) {
            var file: usize = 0;
            FILE: while (file < SIZE) : (file += 1) {
                if (file == 0) std.debug.print("{d}| ", .{SIZE - rank});

                const square = @as(SquareType, @intCast(rank * SIZE + file));

                for (self.bitboards, 0..) |bb, idx| {
                    if (@as(BoardType, @bitCast(bb)) & (@as(BoardType, 1) << square) != 0) {
                        std.debug.print("{c} ", .{@as(Chess.PE, @enumFromInt(idx)).char()});
                        continue :FILE;
                    }
                }
                std.debug.print(". ", .{});
            }
            std.debug.print("\n", .{});
        }
        std.debug.print("{s:>18}\n", .{"---------------"});
        std.debug.print("{s:>18}\n", .{"a b c d e f g h"});

        std.debug.print("Side to move: {any}\n", .{self.side});
        std.debug.print("Enpassant squares: {any}\n", .{self.enpassant});
        std.debug.print("Castling: {any}\n", .{self.castling});
        std.debug.print("Hash: {any}\n", .{self.hash});
    }

    pub fn parseFEN(self: *@This(), in: []const u8) !void {
        self.reset();

        // std.debug.print("in: {s}\n", .{in});

        // parse ranks
        var rank: SquareType = 0;
        var file: SquareType = 0;
        var idx: usize = 0;
        while (idx < in.len) : ({
            idx += 1;
            if (in[idx - 1] != '/') file += 1;
        }) {
            // std.debug.print("{c} {d}\n", .{ in[idx], rank * SIZE + file });
            const ch = in[idx];
            switch (ch) {
                '/' => {
                    rank += 1;
                    file = 0;
                },
                '0'...'9' => {
                    file += try std.fmt.parseUnsigned(SquareType, in[idx .. idx + 1], 10) - 1;
                },
                'P', 'N', 'B', 'R', 'Q', 'K' => { // uppercase
                    const square = @as(Square, @enumFromInt(rank * SIZE + file));
                    self.occupancies[@intFromEnum(Chess.Colors.white)].set(square);
                    switch (ch) {
                        'P' => self.bitboards[@intFromEnum(Chess.PE.P)].set(square),
                        'N' => self.bitboards[@intFromEnum(Chess.PE.N)].set(square),
                        'B' => self.bitboards[@intFromEnum(Chess.PE.B)].set(square),
                        'R' => self.bitboards[@intFromEnum(Chess.PE.R)].set(square),
                        'Q' => self.bitboards[@intFromEnum(Chess.PE.Q)].set(square),
                        'K' => self.bitboards[@intFromEnum(Chess.PE.K)].set(square),
                        else => unreachable,
                    }
                },
                'p', 'n', 'b', 'r', 'q', 'k' => { // lowercase
                    const square = @as(Square, @enumFromInt(rank * SIZE + file));
                    self.occupancies[@intFromEnum(Chess.Colors.black)].set(square);
                    switch (ch) {
                        'p' => self.bitboards[@intFromEnum(Chess.PE.p)].set(square),
                        'n' => self.bitboards[@intFromEnum(Chess.PE.n)].set(square),
                        'b' => self.bitboards[@intFromEnum(Chess.PE.b)].set(square),
                        'r' => self.bitboards[@intFromEnum(Chess.PE.r)].set(square),
                        'q' => self.bitboards[@intFromEnum(Chess.PE.q)].set(square),
                        'k' => self.bitboards[@intFromEnum(Chess.PE.k)].set(square),
                        else => unreachable,
                    }
                },
                ' ' => break,
                else => unreachable,
            }
        }

        // parse rest
        // std.debug.print("rest: {s}\n", .{in[idx..]});
        var parts = std.mem.tokenize(u8, in[idx..], " ");

        // side
        const side = parts.next().?[0];
        // std.debug.print("side: {c}\n", .{side});
        switch (side) {
            'w' => {
                self.side = .white;
            },
            'b' => {
                self.side = .black;
            },
            else => unreachable,
        }

        // castling
        const castling = parts.next().?;
        // std.debug.print("castling: {s}\n", .{castling});
        for (castling) |ch| {
            switch (ch) {
                'K' => self.castling.WK = true,
                'Q' => self.castling.WQ = true,
                'k' => self.castling.BK = true,
                'q' => self.castling.BQ = true,
                '-' => break,
                else => unreachable,
            }
        }

        // enpassant
        const enpassant = parts.next().?;
        // std.debug.print("enpassant: {s}\n", .{enpassant});
        if (enpassant[0] != '-') {
            const f = enpassant[0] - 'a';
            const r = SIZE - @as(
                SquareType,
                @intCast(try std.fmt.parseUnsigned(u3, enpassant[1..], 10)),
            );
            self.enpassant = @as(Square, @enumFromInt(r * SIZE + f));
        } else self.enpassant = null;

        const fifty = parts.next().?;
        self.fifty = std.fmt.parseUnsigned(u7, fifty, 10) catch 0;

        zobrist.updateHash(self);

        return;
    }

    pub fn isSquareAttacked(self: *const @This(), square: SquareType, color: Chess.Colors) bool {
        switch (color) {
            .white => {
                // attacked by white pawns
                const pa = self.attacks.pawn_attacks[@intFromEnum(Chess.Colors.black)][square];
                if (pa & @as(BoardType, @bitCast(self.bitboards[@intFromEnum(Chess.PE.P)])) != 0) return true;
            },
            .black => {
                const pa = self.attacks.pawn_attacks[@intFromEnum(Chess.Colors.white)][square];
                if (pa & @as(BoardType, @bitCast(self.bitboards[@intFromEnum(Chess.PE.p)])) != 0) return true;
            },
        }

        // attacked by knight
        const na = self.attacks.knight_attacks[square];
        const nbb = switch (color) {
            .white => self.bitboards[@intFromEnum(Chess.PE.N)],
            .black => self.bitboards[@intFromEnum(Chess.PE.n)],
        };
        if (na & @as(BoardType, @bitCast(nbb)) != 0) return true;

        // attacked by king
        const ka = self.attacks.king_attacks[square];
        const kbb = switch (color) {
            .white => self.bitboards[@intFromEnum(Chess.PE.K)],
            .black => self.bitboards[@intFromEnum(Chess.PE.k)],
        };
        if (ka & @as(BoardType, @bitCast(kbb)) != 0) return true;

        const occup_both = self.occupBoth();

        // attacked by bishop
        const ba = self.attacks.getBishopAttacks(square, occup_both);
        const bbb = switch (color) {
            .white => self.bitboards[@intFromEnum(Chess.PE.B)],
            .black => self.bitboards[@intFromEnum(Chess.PE.b)],
        };
        if (ba.* & @as(BoardType, @bitCast(bbb)) != 0) return true;

        // attacked by rook
        const ra = self.attacks.getRookAttacks(square, occup_both);
        const rbb = switch (color) {
            .white => self.bitboards[@intFromEnum(Chess.PE.R)],
            .black => self.bitboards[@intFromEnum(Chess.PE.r)],
        };
        if (ra.* & @as(BoardType, @bitCast(rbb)) != 0) return true;

        // attacked by queen
        const qbb = switch (color) {
            .white => self.bitboards[@intFromEnum(Chess.PE.Q)],
            .black => self.bitboards[@intFromEnum(Chess.PE.q)],
        };
        if ((ba.* | ra.*) & @as(BoardType, @bitCast(qbb)) != 0) return true;

        return false;
    }

    pub fn generateMoves(self: *const @This(), ml: *MoveList) !void {
        const pieces = switch (self.side) {
            .white => &[_]Chess.PE{ .P, .N, .B, .R, .Q, .K },
            .black => &[_]Chess.PE{ .p, .n, .b, .r, .q, .k },
        };

        for (pieces) |piece| {
            const board = @as(BoardType, @bitCast(self.bitboards[@intFromEnum(piece)]));

            switch (piece) {
                .P, .p => try self.pawnMoves(ml, board, piece),
                .K, .k => {
                    try self.castlingMoves(ml, piece);
                    try self.genMoves(ml, board, piece);
                },
                .none => unreachable,
                else => try self.genMoves(ml, board, piece),
            }
        }
    }

    fn pawnMoves(self: *const @This(), ml: *MoveList, board_in: BoardType, piece: Chess.PE) !void {
        var source_square: SquareType = undefined;
        var target_square: SquareType = undefined;
        var attacks: BoardType = undefined;

        const promote_options = switch (self.side) {
            .white => [_]Chess.PE{ .Q, .R, .B, .N },
            .black => [_]Chess.PE{ .q, .r, .b, .n },
        };

        var board = board_in;
        while (board != 0) {
            source_square = @as(SquareType, @intCast(@ctz(board)));

            // early exit on underflow/overflow
            switch (self.side) {
                .white => {
                    const ts = @subWithOverflow(source_square, 8);
                    if (ts[1] == 1) continue;
                    target_square = ts[0];
                },
                .black => {
                    const ts = @addWithOverflow(source_square, 8);
                    if (ts[1] == 1) continue;
                    target_square = ts[0];
                },
            }

            const promote_condition = switch (self.side) {
                .white => source_square >= @intFromEnum(Square.a7) and
                    source_square <= @intFromEnum(Square.h7),
                .black => source_square >= @intFromEnum(Square.a2) and
                    source_square <= @intFromEnum(Square.h2),
            };

            // generate quiet moves
            if (self.occupBoth() & (@as(BoardType, 1) << target_square) == 0) {
                const double_step_condition = switch (self.side) {
                    .white => source_square >= @intFromEnum(Square.a2) and
                        source_square <= @intFromEnum(Square.h2) and
                        self.occupBoth() & (@as(BoardType, 1) << (target_square - 8)) == 0,
                    .black => source_square >= @intFromEnum(Square.a7) and
                        source_square <= @intFromEnum(Square.h7) and
                        self.occupBoth() & (@as(BoardType, 1) << (target_square + 8)) == 0,
                };

                // promotion
                if (promote_condition) {
                    for (promote_options) |prom| {
                        const move = BitMove{
                            .source = @as(Square, @enumFromInt(source_square)),
                            .target = @as(Square, @enumFromInt(target_square)),
                            .piece = piece,
                            .prom = prom,
                        };
                        try ml.append(MovePrio{ .move = move });
                    }
                } else {
                    {
                        // single pawn move
                        const move = BitMove{
                            .source = @as(Square, @enumFromInt(source_square)),
                            .target = @as(Square, @enumFromInt(target_square)),
                            .piece = piece,
                        };
                        try ml.append(.{ .move = move });
                    }

                    // double pawn move
                    if (double_step_condition) {
                        const double_target = switch (self.side) {
                            .white => target_square - 8,
                            .black => target_square + 8,
                        };
                        const move = BitMove{
                            .source = @as(Square, @enumFromInt(source_square)),
                            .target = @as(Square, @enumFromInt(double_target)),
                            .piece = piece,
                            .double = true,
                        };
                        try ml.append(.{ .move = move });
                    }
                }
            }

            // generate attacks
            attacks =
                self.attacks.pawn_attacks[@intFromEnum(self.side)][source_square] &
                @as(BoardType, @bitCast(self.occupancies[@intFromEnum(self.side.enemy())]));
            while (attacks != 0) {
                target_square = @as(SquareType, @intCast(@ctz(attacks)));

                if (promote_condition) {
                    for (promote_options) |prom| {
                        const move = BitMove{
                            .source = @as(Square, @enumFromInt(source_square)),
                            .target = @as(Square, @enumFromInt(target_square)),
                            .piece = piece,
                            .prom = prom,
                            .capture = true,
                        };
                        try ml.append(.{ .move = move });
                    }
                } else {
                    const move = BitMove{
                        .source = @as(Square, @enumFromInt(source_square)),
                        .target = @as(Square, @enumFromInt(target_square)),
                        .piece = piece,
                        .capture = true,
                    };
                    try ml.append(.{ .move = move });
                }

                // pop processed attack bit
                attacks ^= @as(BoardType, 1) << target_square;
            }

            // generate enpassant captures
            if (self.enpassant != null) {
                const enpassant_attacks = self.attacks.pawn_attacks[@intFromEnum(self.side)][source_square] &
                    @as(BoardType, 1) << @intFromEnum(self.enpassant.?);

                if (enpassant_attacks != 0) {
                    const target_enpassant: SquareType = @as(u6, @intCast(@ctz(enpassant_attacks)));
                    const move = BitMove{
                        .source = @as(Square, @enumFromInt(source_square)),
                        .target = @as(Square, @enumFromInt(target_enpassant)),
                        .piece = piece,
                        .capture = true,
                        .enpassant = true,
                    };
                    try ml.append(.{ .move = move });
                }
            }

            // pop processed board bit
            board ^= @as(BoardType, 1) << source_square;
        }
    }

    fn castlingMoves(self: *const @This(), ml: *MoveList, piece: Chess.PE) !void {
        { // king-side castling
            const king_side = switch (self.side) {
                .white => [_]Square{ .f1, .g1 },
                .black => [_]Square{ .f8, .g8 },
            };
            const attacked_king = switch (self.side) {
                .white => [_]Square{ .e1, .f1 },
                .black => [_]Square{ .e8, .f8 },
            };

            if ((self.side == .white and self.castling.WK) or
                (self.side == .black and self.castling.BK))
            {
                const OK = blk: {
                    // make sure king side is emptry
                    for (king_side) |square| {
                        if (self.occupBoth() & @as(BoardType, 1) << @intFromEnum(square) != 0) break :blk false;
                    }

                    // check that castling squares are not attacked
                    for (attacked_king) |square| {
                        if (self.isSquareAttacked(@intFromEnum(square), self.side.enemy())) break :blk false;
                    }

                    break :blk true;
                };

                if (OK) {
                    const move = BitMove{
                        .source = attacked_king[0],
                        .target = king_side[1],
                        .piece = piece,
                        .castling = true,
                    };
                    try ml.append(.{ .move = move });
                }
            }
        }

        { // queen-side castling
            const queen_side = switch (self.side) {
                .white => [_]Square{ .d1, .c1, .b1 },
                .black => [_]Square{ .d8, .c8, .b8 },
            };
            const attacked_queen = switch (self.side) {
                .white => [_]Square{ .e1, .d1 },
                .black => [_]Square{ .e8, .d8 },
            };

            if ((self.side == .white and self.castling.WQ) or
                (self.side == .black and self.castling.BQ))
            {
                const OK = blk: {
                    // make sure queen side is emptry
                    for (queen_side) |square| {
                        if (self.occupBoth() & @as(BoardType, 1) << @intFromEnum(square) != 0) break :blk false;
                    }

                    // check that castling squares are not attacked
                    for (attacked_queen) |square| {
                        if (self.isSquareAttacked(@intFromEnum(square), self.side.enemy())) break :blk false;
                    }

                    break :blk true;
                };

                if (OK) {
                    const move = BitMove{
                        .source = attacked_queen[0],
                        .target = queen_side[1],
                        .piece = piece,
                        .castling = true,
                    };
                    try ml.append(.{ .move = move });
                }
            }
        }
    }

    fn genMoves(self: *const @This(), ml: *MoveList, board_in: BoardType, piece: Chess.PE) !void {
        var source_square: SquareType = undefined;
        var target_square: SquareType = undefined;
        var attacks: BoardType = undefined;

        var board = board_in;
        while (board != 0) {
            source_square = @as(SquareType, @intCast(@ctz(board)));

            const attack_mask = switch (piece) {
                .N, .n => self.attacks.knight_attacks[source_square],
                .K, .k => self.attacks.king_attacks[source_square],
                .B, .b => self.attacks.getBishopAttacks(source_square, self.occupBoth()).*,
                .R, .r => self.attacks.getRookAttacks(source_square, self.occupBoth()).*,
                .Q, .q => self.attacks.getQueenAttacks(source_square, self.occupBoth()),
                .P, .p, .none => unreachable,
            };

            attacks = attack_mask &
                ~@as(BoardType, @bitCast(self.occupancies[@intFromEnum(self.side)]));

            while (attacks != 0) {
                target_square = @as(SquareType, @intCast(@ctz(attacks)));

                if (self.occupancies[@intFromEnum(self.side.enemy())].isSet(@as(Square, @enumFromInt(target_square)))) {
                    const move = BitMove{
                        .source = @as(Square, @enumFromInt(source_square)),
                        .target = @as(Square, @enumFromInt(target_square)),
                        .piece = piece,
                        .capture = true,
                    };
                    try ml.append(.{ .move = move });
                } else {
                    const move = BitMove{
                        .source = @as(Square, @enumFromInt(source_square)),
                        .target = @as(Square, @enumFromInt(target_square)),
                        .piece = piece,
                    };
                    try ml.append(.{ .move = move });
                }

                attacks ^= @as(BoardType, 1) << target_square;
            }

            board ^= @as(BoardType, 1) << source_square;
        }
    }

    pub fn makeMove(self: *@This(), move: BitMove, flag: Moves) bool {
        switch (flag) {
            .captures => {
                if (move.capture) {
                    return self.makeMove(move, .all);
                }

                // the move is not a capture
                return false;
            },
            .all => {
                // move piece
                self.bitboards[@intFromEnum(move.piece)].pop(move.source);
                self.bitboards[@intFromEnum(move.piece)].set(move.target);

                // Zobrist: update piece hash
                self.hash ^= zobrist.piece_hashes[@intFromEnum(move.piece)][@intFromEnum(move.source)];
                self.hash ^= zobrist.piece_hashes[@intFromEnum(move.piece)][@intFromEnum(move.target)];

                if (move.piece == .P or move.piece == .p) {
                    self.fifty = 0;
                } else {
                    self.fifty += 1;
                }

                // handle occupancy arrays
                // TODO: perf test against full update used below!
                // {
                //     // remove source target, one of them is no-op
                //     self.occupancies[@enumToInt(Chess.Colors.white)].unSet(move.source);
                //     self.occupancies[@enumToInt(Chess.Colors.black)].unSet(move.source);

                //     // set target for the proper side
                //     switch (self.side) {
                //         .white => {
                //             self.occupancies[@enumToInt(Chess.Colors.white)].set(move.target);
                //             self.occupancies[@enumToInt(Chess.Colors.black)].unSet(move.target);
                //         },
                //         .black => {
                //             self.occupancies[@enumToInt(Chess.Colors.black)].set(move.target);
                //             self.occupancies[@enumToInt(Chess.Colors.white)].unSet(move.target);
                //         },
                //     }
                // }

                { // handling captures
                    if (move.capture) {
                        self.fifty = 0;

                        const pieces = switch (self.side) {
                            .white => [_]Chess.PE{ .p, .n, .b, .r, .q, .k },
                            .black => [_]Chess.PE{ .P, .N, .B, .R, .Q, .K },
                        };

                        // loop over enemy bitboards, when target occupancy found pop it
                        for (pieces) |piece| {
                            if (self.bitboards[@intFromEnum(piece)].isSet(move.target)) {
                                self.bitboards[@intFromEnum(piece)].pop(move.target);

                                // Zobrist: update capture hash
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(piece)][@intFromEnum(move.target)];
                                break;
                            }
                        }
                    }
                }

                { // pawn stuff
                    // handling promotions
                    if (move.prom != .none) {
                        // erase pawn from target square (already moved)
                        self.bitboards[@intFromEnum(move.piece)].pop(move.target);
                        // add promoted item to target peace
                        self.bitboards[@intFromEnum(move.prom)].set(move.target);

                        // Zobrist
                        self.hash ^= zobrist.piece_hashes[@intFromEnum(move.piece)][@intFromEnum(move.target)];
                        self.hash ^= zobrist.piece_hashes[@intFromEnum(move.prom)][@intFromEnum(move.target)];
                    }

                    // handling enpassant captures
                    if (move.enpassant) {
                        switch (self.side) {
                            .white => {
                                self.bitboards[@intFromEnum(Chess.PE.p)]
                                    .pop(@as(Square, @enumFromInt(@intFromEnum(move.target) + 8)));
                                // Zobrist
                                self.hash ^=
                                    zobrist.piece_hashes[@intFromEnum(Chess.PE.p)][@intFromEnum(move.target) + 8];
                            },
                            .black => {
                                self.bitboards[@intFromEnum(Chess.PE.P)]
                                    .pop(@as(Square, @enumFromInt(@intFromEnum(move.target) - 8)));
                                // Zobrist
                                self.hash ^=
                                    zobrist.piece_hashes[@intFromEnum(Chess.PE.P)][@intFromEnum(move.target) - 8];
                            },
                        }
                    }
                    // Zobrist: remove enpassant hash
                    if (self.enpassant != null)
                        self.hash ^= zobrist.enpassant_hashes[@intFromEnum(self.enpassant.?)];

                    // reset enpassant
                    self.enpassant = null;

                    // handle double pawn step
                    if (move.double) {
                        const enpassant = switch (self.side) {
                            .white => @as(Square, @enumFromInt(@intFromEnum(move.target) + 8)),
                            .black => @as(Square, @enumFromInt(@intFromEnum(move.target) - 8)),
                        };
                        self.enpassant = enpassant;
                        // Zobrist: add enpassant hash
                        self.hash ^= zobrist.enpassant_hashes[@intFromEnum(enpassant)];
                    }
                }

                { // handle castling
                    if (move.castling) {
                        // king part is already handled, take care of the rooks
                        switch (move.target) {
                            .g1 => {
                                self.bitboards[@intFromEnum(Chess.PE.R)].pop(Square.h1);
                                self.bitboards[@intFromEnum(Chess.PE.R)].set(Square.f1);
                                // self.occupancies[@enumToInt(Chess.Colors.white)].pop(Square.h1);
                                // self.occupancies[@enumToInt(Chess.Colors.white)].set(Square.f1);

                                // Zobrist
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(Chess.PE.R)][@intFromEnum(Square.h1)];
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(Chess.PE.R)][@intFromEnum(Square.f1)];
                            },
                            .c1 => {
                                self.bitboards[@intFromEnum(Chess.PE.R)].pop(Square.a1);
                                self.bitboards[@intFromEnum(Chess.PE.R)].set(Square.d1);
                                // self.occupancies[@enumToInt(Chess.Colors.white)].pop(Square.a1);
                                // self.occupancies[@enumToInt(Chess.Colors.white)].set(Square.d1);

                                // Zobrist
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(Chess.PE.R)][@intFromEnum(Square.a1)];
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(Chess.PE.R)][@intFromEnum(Square.d1)];
                            },
                            .g8 => {
                                self.bitboards[@intFromEnum(Chess.PE.r)].pop(Square.h8);
                                self.bitboards[@intFromEnum(Chess.PE.r)].set(Square.f8);
                                // self.occupancies[@enumToInt(Chess.Colors.black)].pop(Square.h1);
                                // self.occupancies[@enumToInt(Chess.Colors.black)].set(Square.f1);

                                // Zobrist
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(Chess.PE.r)][@intFromEnum(Square.h8)];
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(Chess.PE.r)][@intFromEnum(Square.f8)];
                            },
                            .c8 => {
                                self.bitboards[@intFromEnum(Chess.PE.r)].pop(Square.a8);
                                self.bitboards[@intFromEnum(Chess.PE.r)].set(Square.d8);
                                // self.occupancies[@enumToInt(Chess.Colors.black)].pop(Square.a8);
                                // self.occupancies[@enumToInt(Chess.Colors.black)].set(Square.d8);

                                // Zobrist
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(Chess.PE.r)][@intFromEnum(Square.a8)];
                                self.hash ^= zobrist.piece_hashes[@intFromEnum(Chess.PE.r)][@intFromEnum(Square.d8)];
                            },
                            else => unreachable,
                        }
                    }

                    // Zobrist - castling rights - remove all
                    self.hash ^= zobrist.castle_hashes[@as(CastlingType, @bitCast(self.castling))];

                    // update castling rights
                    self.castling.update(move.source);
                    self.castling.update(move.target); // have to handle target too!

                    // Zobrist - castling rights - reset all
                    self.hash ^= zobrist.castle_hashes[@as(CastlingType, @bitCast(self.castling))];
                }

                { // update occupancy boards
                    var white: BoardType = 0;
                    var black: BoardType = 0;
                    for ([_]Chess.PE{ .P, .N, .B, .R, .Q, .K }) |piece| {
                        white |= @as(BoardType, @bitCast(self.bitboards[@intFromEnum(piece)]));
                    }
                    for ([_]Chess.PE{ .p, .n, .b, .r, .q, .k }) |piece| {
                        black |= @as(BoardType, @bitCast(self.bitboards[@intFromEnum(piece)]));
                    }

                    self.occupancies[@intFromEnum(Chess.Colors.white)] = @as(BitBoard, @bitCast(white));
                    self.occupancies[@intFromEnum(Chess.Colors.black)] = @as(BitBoard, @bitCast(black));
                }

                { // make sure the move is valid (king is not in check)
                    if (self.inCheck()) return false;
                }

                // FLIP SIDE
                self.side = self.side.enemy();

                // update Zobrist hash
                self.hash ^= zobrist.side_hash;

                return true;
            },
        }
    }

    pub fn inCheck(self: *const @This()) bool {
        const king_square = switch (self.side) {
            .white => @as(Square, @enumFromInt(@ctz(
                @as(BoardType, @bitCast(self.bitboards[@intFromEnum(Chess.PE.K)])),
            ))),
            .black => @as(Square, @enumFromInt(@ctz(
                @as(BoardType, @bitCast(self.bitboards[@intFromEnum(Chess.PE.k)])),
            ))),
        };

        if (self.isSquareAttacked(@intFromEnum(king_square), self.side.enemy())) return true;
        return false;
    }
};

test "isSquareAttacked" {
    var game = try GameState.init(
        std.testing.allocator,
        "r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq - 0 1",
    );
    defer game.deinit();

    var got: BoardType = 0;
    for (std.enums.values(Square)) |square| {
        if (game.isSquareAttacked(@intFromEnum(square), .white)) {
            got |= @as(BoardType, 1) << @intFromEnum(square);
        }
    }

    try std.testing.expectEqual(@as(BoardType, 9149624999898064896), got);
}

test "generateMoves" {
    const tricky_position = "r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq c6 0 1";
    var gs = try GameState.init(std.testing.allocator, tricky_position);
    defer gs.deinit();

    var ml = try MoveList.init(0);
    try gs.generateMoves(&ml);

    try std.testing.expectEqual(@as(usize, 49), ml.len);
}

test "backup and restore" {
    const tricky_position = "r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq c6 0 1";
    var gs = try GameState.init(std.testing.allocator, tricky_position);
    defer gs.deinit();

    var backup: GameState = undefined;
    gs.backup(&backup);

    gs.bitboards[@intFromEnum(Chess.PE.P)].b8 = true;
    gs.enpassant = null;
    gs.castling = .{};
    gs.side = .black;

    gs.restore(&backup);

    try std.testing.expectEqual(false, gs.bitboards[@intFromEnum(Chess.PE.P)].h8);
    try std.testing.expectEqual(Square.c6, gs.enpassant.?);
    try std.testing.expectEqual(@as(u4, 0b1111), @as(u4, @bitCast(gs.castling)));
}