#include "bitboard.h"
#include <stdio.h>

uint64_t PAWN_ATTACKS[2][64];
uint64_t KNIGHT_ATTACKS[64];
uint64_t KING_ATTACKS[64];

void create_knight_attack_cache(void) {
    for (int sq = 0; sq < 64; sq++) {
        uint64_t b = 1ULL << sq;
        uint64_t mask = 0ULL;

        mask |= (b & ~FILE_A)  << 15;
        mask |= (b & ~FILE_H)  << 17;
        mask |= (b & ~(FILE_A | FILE_B)) << 6;
        mask |= (b & ~(FILE_G | FILE_H)) << 10;
        mask |= (b & ~FILE_A)  >> 17;
        mask |= (b & ~FILE_H)  >> 15;
        mask |= (b & ~(FILE_A | FILE_B)) >> 10;
        mask |= (b & ~(FILE_G | FILE_H)) >> 6;

        KNIGHT_ATTACKS[sq] = mask;
    }
}

void create_pawn_attack_cache(void) {
    for (int sq = 0; sq < 64; sq++) {
        uint64_t b = 1ULL << sq;

        // White: NE (+9), NW (+7), but never from rank 8
        uint64_t w  = ((b & ~FILE_H & ~RANK_8) << 9) | ((b & ~FILE_A & ~RANK_8) << 7);

        // Black: SE (-7), SW (-9), but never from rank 1
        uint64_t bl = ((b & ~FILE_H & ~RANK_1) >> 7) | ((b & ~FILE_A & ~RANK_1) >> 9);

        PAWN_ATTACKS[WHITE][sq] = w;
        PAWN_ATTACKS[BLACK][sq] = bl;
    }
}

void create_king_attack_cache(void) {
    for (int sq = 0; sq < 64; sq++) {
        uint64_t b = 1ULL << sq;
        uint64_t mask = 0ULL;

        // North / South
        mask |= (b & ~RANK_8) << 8;
        mask |= (b & ~RANK_1) >> 8;

        // East / West
        mask |= (b & ~FILE_H) << 1;
        mask |= (b & ~FILE_A) >> 1;

        // Diagonals
        mask |= (b & ~FILE_A & ~RANK_8) << 7;
        mask |= (b & ~FILE_H & ~RANK_8) << 9;
        mask |= (b & ~FILE_A & ~RANK_1) >> 9;
        mask |= (b & ~FILE_H & ~RANK_1) >> 7;

        KING_ATTACKS[sq] = mask;
    }
}

int first_set_index(uint64_t bb) {
    for (int i = 0; i < 64; ++i) {
        if ((bb >> i) & 1ULL) return i;
    }
    return -1;
}

int pop_lsb_index(uint64_t *bb) {
    if (*bb == 0) return -1;
    int idx = first_set_index(*bb);
    // Clears bit.
    *bb &= ~(1ULL << idx);
    return idx;
}

uint64_t ray_attacks(int start_square, int delta_file, int delta_rank, uint64_t occupied)
{
    uint64_t attacks = 0;

    int start_file = start_square % 8;
    int start_rank = start_square / 8;

    int next_file = start_file + delta_file;
    int next_rank = start_rank + delta_rank;

    while (next_file >= 0 && next_file < 8 && next_rank >= 0 && next_rank < 8) {
        int next_square = next_rank * 8 + next_file;
        uint64_t next_bit = 1ULL << next_square;

        // can attack this square
        attacks |= next_bit;

        // blocker: include it, then stop
        if (occupied & next_bit)
            break;

        next_file += delta_file;
        next_rank += delta_rank;
    }

    return attacks;
}

uint64_t rook_attacks(int square, uint64_t occupied) {
    return  ray_attacks(square, +1,  0, occupied)   // east
          | ray_attacks(square, -1,  0, occupied)   // west
          | ray_attacks(square,  0, +1, occupied)   // north
          | ray_attacks(square,  0, -1, occupied);  // south
}

uint64_t bishop_attacks(int square, uint64_t occupied) {
    return  ray_attacks(square, +1, +1, occupied)   // NE
          | ray_attacks(square, -1, +1, occupied)   // NW
          | ray_attacks(square, +1, -1, occupied)   // SE
          | ray_attacks(square, -1, -1, occupied);  // SW
}

uint64_t queen_attacks(int square, uint64_t occupied) {
    // Simply combine both types of attacks
    return rook_attacks(square, occupied) | bishop_attacks(square, occupied);
}

void push_move(struct Move *out, int *count, int from, int to, uint8_t piece, uint8_t promo, uint8_t flags)
{
    out[*count] = (struct Move){
        .from  = (uint16_t)from,
        .to    = (uint16_t)to,
        .piece = piece,
        .promo = promo,    // 0 for non-promo;
        .flags = flags
    };
    (*count)++;
}

/**
    PAWN struct Movement Psuedo Moves


    Separate the white and black pawn logic. Although the logic is similar,
    by separating the logic, we can isolate the different cases and have clearer parent functions.
    These pieces are also the only ones that move in opposite ways, meaning we only need
    to separate out seemingly related code for this case. Additionally, pawns have the
    most complicated movement patterns and behaviors in the game. Separating out the logic should
    hopefully make debugging easier.
*/
void gen_white_pawn_quiet_pushes(const struct Board *board, struct Move *out, int *count) {
    const uint64_t occ   = board->occ[BOTH];
    const uint64_t empty = ~occ;

    const uint64_t pawns = board->pieces[P];

    // One-step pushes: shift pawns up by 8 onto empty squares
    uint64_t one_step = (pawns << 8) & empty;

    // Exclude promotions here (dest on rank 8) — handled by a separate promo-push function
    uint64_t one_step_no_promo = one_step & ~RANK_8;

    // Two-step pushes: those one-step pawns that landed on rank 3 can go one more if still empty
    // (means they originally stood on rank 2)
    uint64_t two_step = ((one_step & RANK_3) << 8) & empty;

    uint64_t bb = one_step_no_promo;
    while (bb) {
        int to = pop_lsb_index(&bb);
        int from = to - 8;
        push_move(out, count, from, to, P, 0, MF_NONE);
    }

    bb = two_step;
    while (bb) {
        int to = pop_lsb_index(&bb);
        int from = to - 16;
        push_move(out, count, from, to, P, 0, MF_DOUBLE_PUSH);
    }
}

void gen_black_pawn_quiet_pushes(const struct Board *board, struct Move *out, int *count) {
    const uint64_t occ   = board->occ[BOTH];
    const uint64_t empty = ~occ;

    const uint64_t pawns = board->pieces[p];

    // One-step pushes: shift pawns down by 8 onto empty squares
    uint64_t one_step = (pawns >> 8) & empty;

    // Exclude promotions here (dest on rank 1) — handle in a promo-push function
    uint64_t one_step_no_promo = one_step & ~RANK_1;

    uint64_t two_step = ((one_step & RANK_6) >> 8) & empty;

    uint64_t bb = one_step_no_promo;
    while (bb) {
        int to   = pop_lsb_index(&bb);
        int from = to + 8;
        push_move(out, count,from, to, p, 0, MF_NONE);
    }

    bb = two_step;
    while (bb) {
        int to   = pop_lsb_index(&bb);
        int from = to + 16;
        push_move(out, count,from, to, p, 0, MF_DOUBLE_PUSH);
    }
}

// We will only allow pawns to promote to queen. Technically, they should be allowed
// to promote to any of the following: Rook, Bishop, Knight, Queen.
void gen_white_pawn_push_promotions(const struct Board *board, struct Move *out, int *count) {
    const uint64_t occ   = board->occ[BOTH];
    const uint64_t empty = ~occ;
    const uint64_t pawns = board->pieces[P];

    // destinations on rank 8 reachable by a single push
    uint64_t promos = ((pawns << 8) & empty) & RANK_8;

    while (promos) {
        int to   = pop_lsb_index(&promos);
        int from = to - 8;
        push_move(out, count, from, to, P, Q, MF_PROMO);
        push_move(out, count, from, to, P, R, MF_PROMO);
        push_move(out, count, from, to, P, B, MF_PROMO);
        push_move(out, count, from, to, P, N, MF_PROMO);
    }
}

void gen_black_pawn_push_promotions(const struct Board *board, struct Move *out, int *count) {
    const uint64_t occ   = board->occ[BOTH];
    const uint64_t empty = ~occ;
    const uint64_t pawns = board->pieces[p];

    // destinations on rank 1 reachable by a single push
    uint64_t promos = ((pawns >> 8) & empty) & RANK_1;

    while (promos) {
        int to   = pop_lsb_index(&promos);
        int from = to + 8;
        push_move(out, count, from, to, p, q, MF_PROMO);
        push_move(out, count, from, to, p, r, MF_PROMO);
        push_move(out, count, from, to, p, b, MF_PROMO);
        push_move(out, count, from, to, p, n, MF_PROMO);
    }
}

void gen_white_pawn_capture_promotions(const struct Board *board, struct Move *out, int *count) {
    const uint64_t opp = board->occ[BLACK];
    // left capture (from white view): +7, mask off file A
    uint64_t left  = ((board->pieces[P] & ~FILE_A) << 7) & opp & RANK_8;
    while (left) {
        int to = pop_lsb_index(&left);
        int from = to - 7;
        push_move(out, count, from, to, P, Q, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, P, R, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, P, B, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, P, N, MF_CAPTURE | MF_PROMO);
    }
    // right capture: +9, mask off file H
    uint64_t right = ((board->pieces[P] & ~FILE_H) << 9) & opp & RANK_8;
    while (right) {
        int to = pop_lsb_index(&right);
        int from = to - 9;
        push_move(out, count, from, to, P, Q, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, P, R, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, P, B, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, P, N, MF_CAPTURE | MF_PROMO);
    }
}

void gen_black_pawn_capture_promotions(const struct Board *board, struct Move *out, int *count) {
    const uint64_t opp = board->occ[WHITE];
    // from black view, “left” is -7 (mask off file H before shifting)
    uint64_t left  = ((board->pieces[p] & ~FILE_H) >> 7) & opp & RANK_1;
    while (left) {
        int to = pop_lsb_index(&left);
        int from = to + 7;
        push_move(out, count, from, to, p, q, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, p, r, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, p, b, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, p, n, MF_CAPTURE | MF_PROMO);
    }
    // “right” is -9 (mask off file A)
    uint64_t right = ((board->pieces[p] & ~FILE_A) >> 9) & opp & RANK_1;
    while (right) {
        int to = pop_lsb_index(&right);
        int from = to + 9;
        push_move(out, count, from, to, p, q, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, p, r, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, p, b, MF_CAPTURE | MF_PROMO);
        push_move(out, count, from, to, p, n, MF_CAPTURE | MF_PROMO);
    }
}

void gen_white_pawn_captures(const struct Board *board, struct Move *out, int *count) {
    const uint64_t pawns  = board->pieces[P];
    const uint64_t opp    = board->occ[BLACK];

    // Normal captures (exclude promotion rank)
    uint64_t left_caps  = ((pawns & ~FILE_A) << 7) & opp & ~RANK_8;
    uint64_t right_caps = ((pawns & ~FILE_H) << 9) & opp & ~RANK_8;

    // Emit normal captures
    while (left_caps) {
        int to   = pop_lsb_index(&left_caps);
        int from = to - 7;
        push_move(out, count,from, to, P, 0, MF_CAPTURE);
    }
    while (right_caps) {
        int to   = pop_lsb_index(&right_caps);
        int from = to - 9;
        push_move(out, count,from, to, P, 0, MF_CAPTURE);
    }

    // En passant (destination is ep_square)
    if (board->ep_square >= 0) {
        uint64_t ep = 1ULL << board->ep_square;

        // From-squares that can capture onto ep (reverse of +7/+9)
        uint64_t ep_from = (((ep & ~FILE_H) >> 7) | ((ep & ~FILE_A) >> 9)) & pawns;

        while (ep_from) {
            int from = pop_lsb_index(&ep_from);
            int to   = board->ep_square;
            // EP never promotes, so no promo piece; still a capture
            push_move(out, count,from, to, P, 0, MF_CAPTURE | MF_ENPASSANT);
        }
    }
}

void gen_black_pawn_captures(const struct Board *board, struct Move *out, int *count) {
    const uint64_t pawns  = board->pieces[p];
    const uint64_t opp    = board->occ[WHITE];

    // Normal captures (exclude promotion rank)
    uint64_t left_caps  = ((pawns & ~FILE_H) >> 7) & opp & ~RANK_1; // from black POV, "left" is -7
    uint64_t right_caps = ((pawns & ~FILE_A) >> 9) & opp & ~RANK_1; // "right" is -9

    while (left_caps) {
        int to   = pop_lsb_index(&left_caps);
        int from = to + 7;
        push_move(out, count,from, to, p, 0, MF_CAPTURE);
    }
    while (right_caps) {
        int to   = pop_lsb_index(&right_caps);
        int from = to + 9;
        push_move(out, count,from, to, p, 0, MF_CAPTURE);
    }

    // En passant
    if (board->ep_square >= 0) {
        uint64_t ep = 1ULL << board->ep_square;

        // From-squares that can capture onto ep (reverse of -7/-9)
        uint64_t ep_from = (((ep & ~FILE_A) << 7) | ((ep & ~FILE_H) << 9)) & pawns;

        while (ep_from) {
            int from = pop_lsb_index(&ep_from);
            int to   = board->ep_square;
            push_move(out, count,from, to, p, 0, MF_CAPTURE | MF_ENPASSANT);
        }
    }
}

/**


    Pieces that are not PAWNs


*/
void gen_knight_moves(const struct Board *board, struct Move *out, int *count, bool captures_only) {
    enum Color side = board->side_to_move;
    uint64_t own = board->occ[side];
    uint64_t opp = board->occ[side ^ 1];
    uint8_t pid = (side == WHITE) ? N : n;
    uint64_t bb = (side == WHITE) ? board->pieces[N] : board->pieces[n];

    while (bb) {
        int from = pop_lsb_index(&bb);
        uint64_t mask = KNIGHT_ATTACKS[from] & ~own;
        uint64_t caps = mask & opp;

        if (!captures_only) {
            uint64_t quiet = mask & ~opp;
            while (quiet) {
                int to = pop_lsb_index(&quiet);
                push_move(out, count, from, to, pid, 0, MF_NONE);
            }
        }
        while (caps) {
            int to = pop_lsb_index(&caps);
            push_move(out, count, from, to, pid, 0, MF_CAPTURE);
        }
    }
}

void gen_bishop_moves(const struct Board *board, struct Move *out, int *count, bool captures_only) {
    enum Color side = board->side_to_move;
    uint64_t own = board->occ[side];
    uint64_t opp = board->occ[side ^ 1];
    uint64_t occ = board->occ[BOTH];
    uint8_t pid = (side == WHITE) ? B : b;
    uint64_t bb = (side == WHITE) ? board->pieces[B] : board->pieces[b];

    while (bb) {
        int from = pop_lsb_index(&bb);
        uint64_t mask = bishop_attacks(from, occ) & ~own;
        uint64_t caps = mask & opp;

        if (!captures_only) {
            uint64_t quiet = mask & ~opp;
            while (quiet) {
                int to = pop_lsb_index(&quiet);
                push_move(out, count, from, to, pid, 0, MF_NONE);
            }
        }
        while (caps) {
            int to = pop_lsb_index(&caps);
            push_move(out, count, from, to, pid, 0, MF_CAPTURE);
        }
    }
}

void gen_rook_moves(const struct Board *board, struct Move *out, int *count, bool captures_only) {
    enum Color side = board->side_to_move;
    uint64_t own = board->occ[side];
    uint64_t opp = board->occ[side ^ 1];
    uint64_t occ = board->occ[BOTH];
    uint8_t pid = (side == WHITE) ? R : r;
    uint64_t bb = (side == WHITE) ? board->pieces[R] : board->pieces[r];

    while (bb) {
        int from = pop_lsb_index(&bb);
        uint64_t mask = rook_attacks(from, occ) & ~own;
        uint64_t caps = mask & opp;

        if (!captures_only) {
            uint64_t quiet = mask & ~opp;
            while (quiet) {
                int to = pop_lsb_index(&quiet);
                push_move(out, count,from, to, pid, 0, MF_NONE);
            }
        }
        while (caps) {
            int to = pop_lsb_index(&caps);
            push_move(out, count,from, to, pid, 0, MF_CAPTURE);
        }
    }
}

void gen_queen_moves(const struct Board *board, struct Move *out, int *count, bool captures_only) {
    enum Color side = board->side_to_move;
    uint64_t own = board->occ[side];
    uint64_t opp = board->occ[side ^ 1];
    uint64_t occ = board->occ[BOTH];
    uint8_t pid = (side == WHITE) ? Q : q;
    uint64_t bb = (side == WHITE) ? board->pieces[Q] : board->pieces[q];

    while (bb) {
        int from = pop_lsb_index(&bb);
        uint64_t mask = (rook_attacks(from, occ) | bishop_attacks(from, occ)) & ~own;
        uint64_t caps = mask & opp;

        if (!captures_only) {
            uint64_t quiet = mask & ~opp;
            while (quiet) {
                int to = pop_lsb_index(&quiet);
                push_move(out, count,from, to, pid, 0, MF_NONE);
            }
        }
        while (caps) {
            int to = pop_lsb_index(&caps);
            push_move(out, count,from, to, pid, 0, MF_CAPTURE);
        }
    }
}

void gen_king_moves(const struct Board *board, struct Move *out, int *count, bool captures_only) {
    enum Color side = board->side_to_move;
    uint64_t own = board->occ[side];
    uint64_t opp = board->occ[side ^ 1];
    uint8_t pid = (side == WHITE) ? K : k;
    uint64_t kk = (side == WHITE) ? board->pieces[K] : board->pieces[k];

    if (!kk) return;
    int from = first_set_index(kk);
    uint64_t mask = KING_ATTACKS[from] & ~own;
    uint64_t caps = mask & opp;

    if (!captures_only) {
        uint64_t quiet = mask & ~opp;
        while (quiet) {
            int to = pop_lsb_index(&quiet);
            push_move(out, count,from, to, pid, 0, MF_NONE);
        }
    }
    while (caps) {
        int to = pop_lsb_index(&caps);
        push_move(out, count,from, to, pid, 0, MF_CAPTURE);
    }

    if (!captures_only) {
        uint64_t occ = board->occ[BOTH];
        if (side == WHITE) {
            if ((board->castling_rights & CASTLE_WK) && !(occ & WK_EMPTY_MASK)) {
                push_move(out, count,E1, WK_TO, K, 0, MF_CASTLE);
            }
            if ((board->castling_rights & CASTLE_WQ) && !(occ & WQ_EMPTY_MASK)) {
                push_move(out, count,E1, WQ_TO, K, 0, MF_CASTLE);
            }
        } else {
            if ((board->castling_rights & CASTLE_BK) && !(occ & BK_EMPTY_MASK)) {
                push_move(out, count,E8, BK_TO, k, 0, MF_CASTLE);
            }
            if ((board->castling_rights & CASTLE_BQ) && !(occ & BQ_EMPTY_MASK)) {
                push_move(out, count,E8, BQ_TO, k, 0, MF_CASTLE);
            }
        }
    }
}

int gen_pseudo_moves(const struct Board *board, struct Move *moves, bool captures_only) {
    int count = 0;

    if (board->side_to_move == WHITE) {
        if (!captures_only) {
            gen_white_pawn_quiet_pushes(board, moves, &count);
            gen_white_pawn_push_promotions(board, moves, &count);
        }
        gen_white_pawn_captures(board, moves, &count);
        gen_white_pawn_capture_promotions(board, moves, &count);
    } else {
        if (!captures_only) {
            gen_black_pawn_quiet_pushes(board, moves, &count);
            gen_black_pawn_push_promotions(board, moves, &count);
        }
        gen_black_pawn_captures(board, moves, &count);
        gen_black_pawn_capture_promotions(board, moves, &count);
    }

    gen_knight_moves(board, moves, &count, captures_only);
    gen_king_moves(board, moves, &count, captures_only);
    gen_rook_moves(board, moves, &count, captures_only);
    gen_bishop_moves(board, moves, &count, captures_only);
    gen_queen_moves(board, moves, &count, captures_only);

    return count;
}

uint64_t attackers_to(const struct Board *board, int sq, enum Color by) {
    uint64_t occ = board->occ[BOTH];

    uint64_t pawns = (by == WHITE) ? board->pieces[P] : board->pieces[p];
    uint64_t knights = (by == WHITE) ? board->pieces[N] : board->pieces[n];
    uint64_t bishops = (by == WHITE) ? board->pieces[B] : board->pieces[b];
    uint64_t rooks = (by == WHITE) ? board->pieces[R] : board->pieces[r];
    uint64_t queens = (by == WHITE) ? board->pieces[Q] : board->pieces[q];
    uint64_t kings = (by == WHITE) ? board->pieces[K] : board->pieces[k];

    uint64_t attacks = 0;

    attacks |= PAWN_ATTACKS[by ^ 1][sq] & pawns;
    attacks |= KNIGHT_ATTACKS[sq] & knights;
    attacks |= KING_ATTACKS[sq]   & kings;

    // Squares than can see sq given the occupancy of both sides but intersect
    // with the specified color pieces.
    attacks |= bishop_attacks(sq, occ) & (bishops | queens);
    attacks |= rook_attacks(sq, occ) & (rooks | queens);

    return attacks;
}

bool in_check(const struct Board *board, enum Color side) {
    uint64_t king_bb = (side == WHITE) ? board->pieces[K] : board->pieces[k];
    
    // Do we need this check here?
    // This would be an illegal position.
    if (!king_bb) return false;

    int king_sq = pop_lsb_index(&king_bb);

    // Is the king square in the attack map of opposite side?
    return attackers_to(board, king_sq, side ^ 1) != 0;
}

bool square_attacked(const struct Board *board, int sq, enum Color by) {
    // Should we have another way to calculate this and have
    // earlier exits in the code? For example, we could iterate
    // each piece type and return boolean sooner and skip the
    // ray scanning.
    return attackers_to(board, sq, by) != 0;
}

int piece_at(const struct Board *board, int sq) {
    uint64_t mask = 1ULL << sq;
    for (int piece = 0; piece < 12; ++piece)
        if (board->pieces[piece] & mask) return piece;
    return -1;
}

void rebuild_occ(struct Board *board) {
    uint64_t white=0;
    uint64_t black=0;

    for (int p = P; p <= K; ++p) white |= board->pieces[p];
    for (int p = p; p <= k; ++p) black |= board->pieces[p];

    board->occ[WHITE] = white;
    board->occ[BLACK] = black;
    board->occ[BOTH]  = white | black;
}

uint8_t castlerights_without(uint8_t castlerights, uint8_t mask) {
    // Update castle rights by removing mask.
    return castlerights & ~mask;
}

uint8_t update_castlerights_after_move(
    uint8_t cr, int moved_pid, int from, int to, int cap_pid
) {
    // King moved: lose both
    if (moved_pid == K) cr &= ~(CASTLE_WK | CASTLE_WQ);
    if (moved_pid == k) cr &= ~(CASTLE_BK | CASTLE_BQ);

    // Rook moved from its home square
    if (moved_pid == R) {
        if (from == H1) cr &= ~CASTLE_WK;
        if (from == A1) cr &= ~CASTLE_WQ;
    } else if (moved_pid == r) {
        if (from == H8) cr &= ~CASTLE_BK;
        if (from == A8) cr &= ~CASTLE_BQ;
    }

    // Captured a rook on its home square.
    // Obviously if a rook has moved and is captured,
    // the prior check accounts for that.
    if (cap_pid == R) {
        if (to == H1) cr &= ~CASTLE_WK;
        if (to == A1) cr &= ~CASTLE_WQ;
    } else if (cap_pid == r) {
        if (to == H8) cr &= ~CASTLE_BK;
        if (to == A8) cr &= ~CASTLE_BQ;
    }
    return cr;
}

void put_piece(struct Board *board, int pid, int sq) {
    board->pieces[pid] |= 1ULL << sq;
}

void del_piece(struct Board *board, int pid, int sq) {
    board->pieces[pid] &= ~(1ULL << sq);
}

// Create a copy of the board to simulate move and test state.
bool apply_move_on_copy(struct Board *in, struct Board *out, struct Move m) {
    *out = *in;
    const enum Color us = in->side_to_move;
    const enum Color opp = us ^ 1;

    // Reset enpassant.
    out->ep_square = -1;

    // Handle capture / EP capture
    int captured_pid = -1;
    if (m.flags & MF_ENPASSANT) {
        int cap_sq = (us == WHITE) ? (m.to - 8) : (m.to + 8);
        captured_pid = piece_at(out, cap_sq);
        if (captured_pid < 0) return false;
        del_piece(out, captured_pid, cap_sq);
    } else if (m.flags & MF_CAPTURE) {
        captured_pid = piece_at(out, m.to);
        if (captured_pid < 0) {      
            return false;
        };
        del_piece(out, captured_pid, m.to);
    }

    // Move (with promotion)
    del_piece(out, m.piece, m.from);
    // move.promo is the promotion piece.
    int placed = (m.flags & MF_PROMO) ? m.promo : m.piece;
    put_piece(out, placed, m.to);

    // Castling: move rook
    if (m.flags & MF_CASTLE) {
        if (us == WHITE) {
            if (m.to == WK_TO) { 
                del_piece(out, R, H1);
                put_piece(out, R, F1);
            }
            else {
                del_piece(out, R, A1);
                put_piece(out, R, D1);
            }
        } else {
            if (m.to == BK_TO) {
                del_piece(out, r, H8);
                put_piece(out, r, F8);
            }
            else {
                del_piece(out, r, A8);
                put_piece(out, r, D8);
            }
        }
    }

    // Double pawn push -> set EP square
    if (m.flags & MF_DOUBLE_PUSH) out->ep_square = (us == WHITE) ? (m.to - 8) : (m.to + 8);

    // Update castling rights
    out->castling_rights = update_castlerights_after_move(
        in->castling_rights, m.piece, m.from, m.to, captured_pid
    );

    // Side to move flips
    out->side_to_move = opp;

    rebuild_occ(out);
    return true;
}

bool castle_path_safe(const struct Board *board, enum Color side, int king_to) {
    enum Color opp = side ^ 1;
    if (side == WHITE) {
        if (king_to == WK_TO) {
            // E1-F1-G1
            if (square_attacked(board, E1, opp)) return false;
            if (square_attacked(board, F1, opp)) return false;
        } else {
            // E1-D1-C1
            if (square_attacked(board, E1, opp)) return false;
            if (square_attacked(board, D1, opp)) return false;
        }
    } else {
        if (king_to == BK_TO) {
            // E8-F8-G8
            if (square_attacked(board, E8, opp)) return false;
            if (square_attacked(board, F8, opp)) return false;
        } else {
            // E8-D8-C8
            if (square_attacked(board, E8, opp)) return false;
            if (square_attacked(board, D8, opp)) return false;
        }
    }
    return true;
}

bool move_is_legal(struct Board *board, struct Move move) {
    enum Color us  = board->side_to_move;
    enum Color opp = us ^ 1;

    // Castle-through-check rule
    if (move.flags & MF_CASTLE) {
        if (!castle_path_safe(board, us, move.to)) {
            return false;
        }
    }

    struct Board after;
    if (!apply_move_on_copy(board, &after, move)) {
        return false;
    }


    // Our king must not be in check after the move
    uint64_t king_bb = (us == WHITE) ? after.pieces[K] : after.pieces[k];

    if (!king_bb) return false;

    int king_sq = first_set_index(king_bb);
    return !square_attacked(&after, king_sq, opp);
}

int get_legal_moves(struct Board *board, struct Move *out) {
    struct Move tmp[256];
    int n = gen_pseudo_moves(board, tmp, false);
    
    int count = 0;
    for (int i = 0; i < n; ++i)
        if (move_is_legal(board, tmp[i]))
            out[count++] = tmp[i];
    return count;
}

uint64_t perft(struct Board *board, int depth) {
    if (depth == 0) return 1;

    struct Move moves[256];
    int n = get_legal_moves(board, moves);

    if (depth == 1) return (uint64_t) n;
    
    uint64_t nodes = 0;
    for (int i = 0; i < n; ++i) {
        struct Board after;
        if (!apply_move_on_copy(board, &after, moves[i])) {
            // Shouldn't happen with legal moves, but be defensive
            continue;
        }
        nodes += perft(&after, depth - 1);
    }
    return nodes;
}

void print_board(const struct Board *board) {
    const char PIECE_CH[12] = {
        'P','N','B','R','Q','K',
        'p','n','b','r','q','k'
    };

    char grid[64];
    for (int i = 0; i < 64; ++i) grid[i] = '.';

    for (int p = 0; p < 12; ++p) {
        uint64_t bb = board->pieces[p];
        while (bb) {
            int sq = pop_lsb_index(&bb);
            grid[sq] = PIECE_CH[p];
        }
    }

    for (int r = 7; r >= 0; --r) {
        for (int f = 0; f < 8; ++f) {
            putchar(grid[r * 8 + f]);
            if (f < 7) putchar(' ');
        }
        putchar('\n');
    }
}