#include "../lib/std.cpp"
#include "../lib/math/math.cpp"

// Helper class to pass the game state to the AI.
class SimplifiedBoardState {
    array<array<int>> board;
    int current_player;
    bool can_castle_w_k, can_castle_w_q, can_castle_b_k, can_castle_b_q;
    int en_passant_r, en_passant_c;
    int king_pos_w_r, king_pos_w_c, king_pos_b_r, king_pos_b_c;

    SimplifiedBoardState(
        const array<array<int>> &in board, int current_player,
        bool cwk, bool cwq, bool cbk, bool cbq,
        int epr, int epc,
        int kwr, int kwc, int kbr, int kbc
    ) {
        this.board = board;
        this.current_player = current_player;
        this.can_castle_w_k = cwk; this.can_castle_w_q = cwq;
        this.can_castle_b_k = cbk; this.can_castle_b_q = cbq;
        this.en_passant_r = epr; this.en_passant_c = epc;
        this.king_pos_w_r = kwr; this.king_pos_w_c = kwc;
        this.king_pos_b_r = kbr; this.king_pos_b_c = kbc;
    }
}

class SimpleAI {
    // --- All move validation logic is now part of the AI ---

    bool is_valid_move(SimplifiedBoardState@ sbs, int r1, int c1, int r2, int c2, int color) {
        if (!is_on_board(r1,c1) || !is_on_board(r2,c2) || (sbs.board[r1][c1]&COLOR_MASK)!=color || (sbs.board[r2][c2]&COLOR_MASK)==color) return false;
        
        int piece_type = sbs.board[r1][c1] & PIECE_MASK;
        if (!is_legal_shape(sbs, piece_type, r1, c1, r2, c2, color)) return false;

        // Simulate the move to check for checks
        int p_dest = sbs.board[r2][c2]; sbs.board[r2][c2] = sbs.board[r1][c1]; sbs.board[r1][c1] = EMPTY;
        int kr_orig = (color==WHITE) ? sbs.king_pos_w_r : sbs.king_pos_b_r;
        int kc_orig = (color==WHITE) ? sbs.king_pos_w_c : sbs.king_pos_b_c;

        if (piece_type == KING) {
            if(color == WHITE) { sbs.king_pos_w_r = r2; sbs.king_pos_w_c = c2; }
            else { sbs.king_pos_b_r = r2; sbs.king_pos_b_c = c2; }
        }
        
        bool in_check = is_king_in_check(sbs, color);
        
        // Revert the simulation
        sbs.board[r1][c1] = sbs.board[r2][c2]; sbs.board[r2][c2] = p_dest;
        if(color == WHITE) { sbs.king_pos_w_r = kr_orig; sbs.king_pos_w_c = kc_orig; }
        else { sbs.king_pos_b_r = kr_orig; sbs.king_pos_b_c = kc_orig; }
        
        return !in_check;
    }
    
    bool is_legal_shape(SimplifiedBoardState@ sbs, int type, int r1, int c1, int r2, int c2, int color) {
        int dr = r2 - r1, dc = c2 - c1;
        if (type == KING) {
            if (r1 == r2 && abs(dc) == 2) {
                return can_castle(sbs, color, dc > 0);
            }
            return abs(dr) <= 1 && abs(dc) <= 1;
        }
        if (type == KNIGHT) return (abs(dr) == 1 && abs(dc) == 2) || (abs(dr) == 2 && abs(dc) == 1);
        if (type == QUEEN) {
            if (abs(dr) == abs(dc) || dr == 0 || dc == 0) return is_path_clear(sbs, r1, c1, r2, c2);
        }
        if (type == ROOK) {
            if (dr == 0 || dc == 0) return is_path_clear(sbs, r1, c1, r2, c2);
        }
        if (type == BISHOP) {
            if (abs(dr) == abs(dc)) return is_path_clear(sbs, r1, c1, r2, c2);
        }
        if (type == PAWN) {
            int fwd = (color == WHITE) ? -1 : 1;
            if(dr == fwd && dc == 0 && sbs.board[r2][c2] == EMPTY) return true;
            if(dr == 2*fwd && dc == 0 && (r1 == (color==WHITE?6:1)) && sbs.board[r1+fwd][c1]==EMPTY && sbs.board[r2][c2]==EMPTY) return true;
            if(dr == fwd && abs(dc) == 1 && (sbs.board[r2][c2] & COLOR_MASK) != color && sbs.board[r2][c2] != EMPTY) return true;
            if(dr == fwd && abs(dc) == 1 && r2 == sbs.en_passant_r && c2 == sbs.en_passant_c) return true;
        }
        return false;
    }

    bool can_castle(SimplifiedBoardState@ sbs, int color, bool kingside) {
        int r = (color == WHITE) ? 7 : 0;
        int opponent_color = (color == WHITE) ? BLACK : WHITE;
        if (is_king_in_check(sbs, color)) return false;
        if (kingside) {
            if ((color == WHITE && !sbs.can_castle_w_k) || (color == BLACK && !sbs.can_castle_b_k)) return false;
            if (sbs.board[r][5] != EMPTY || sbs.board[r][6] != EMPTY) return false;
            if (is_square_attacked(sbs, r, 5, opponent_color) || is_square_attacked(sbs, r, 6, opponent_color)) return false;
        } else {
            if ((color == WHITE && !sbs.can_castle_w_q) || (color == BLACK && !sbs.can_castle_b_q)) return false;
            if (sbs.board[r][1] != EMPTY || sbs.board[r][2] != EMPTY || sbs.board[r][3] != EMPTY) return false;
            if (is_square_attacked(sbs, r, 2, opponent_color) || is_square_attacked(sbs, r, 3, opponent_color)) return false;
        }
        return true;
    }

    bool is_path_clear(SimplifiedBoardState@ sbs, int r1, int c1, int r2, int c2) {
        int dr = (r2 > r1) ? 1 : ((r2 < r1) ? -1 : 0);
        int dc = (c2 > c1) ? 1 : ((c2 < c1) ? -1 : 0);
        int r = r1 + dr; int c = c1 + dc;
        while (r != r2 || c != c2) {
            if (sbs.board[r][c] != EMPTY) return false;
            r += dr; c += dc;
        }
        return true;
    }
    
    bool is_king_in_check(SimplifiedBoardState@ sbs, int color) {
        int r = (color == WHITE) ? sbs.king_pos_w_r : sbs.king_pos_b_r;
        int c = (color == WHITE) ? sbs.king_pos_w_c : sbs.king_pos_b_c;
        return is_square_attacked(sbs, r, c, (color == WHITE) ? BLACK : WHITE);
    }

    bool is_square_attacked(SimplifiedBoardState@ sbs, int r, int c, int by_color) {
        for (int r1=0; r1<8; r1++) for (int c1=0; c1<8; c1++)
            if ((sbs.board[r1][c1]&COLOR_MASK)==by_color && is_legal_shape(sbs, sbs.board[r1][c1]&PIECE_MASK,r1,c1,r,c,by_color)) return true;
        return false;
    }

    bool is_on_board(int r, int c) { return r >= 0 && r < 8 && c >= 0 && c < 8; }

    // --- Main AI entry point ---

    array<int>@ get_move(SimplifiedBoardState@ sbs, replay_rand@ rrnd) {
        array<int> moves;
        for (int r1=0;r1<8;r1++) {
            for (int c1=0;c1<8;c1++) {
                if((sbs.board[r1][c1] & COLOR_MASK) == sbs.current_player) {
                    for(int r2=0;r2<8;r2++) {
                        for(int c2=0;c2<8;c2++) {
                            if(is_valid_move(sbs, r1, c1, r2, c2, sbs.current_player)) {
                                moves.insertLast(r1); moves.insertLast(c1);
                                moves.insertLast(r2); moves.insertLast(c2);
                            }
                        }
                    }
                }
            }
        }
        
        if (moves.length() > 0) {
            int i = (rrnd.rand() % (moves.length()/4))*4;
            array<int>@ chosen_move = {moves[i], moves[i+1], moves[i+2], moves[i+3]};
            return chosen_move;
        }
        
        return null; // No legal moves
    }
}