#include "../lib/string/common.cpp"

// Constants used by the AI
const int ADV_AI_KING = 1, ADV_AI_KNIGHT = 2, ADV_AI_PAWN = 3, ADV_AI_QUEEN = 4, ADV_AI_ROOK = 5, ADV_AI_BISHOP = 6;
const int ADV_AI_WHITE = 8, ADV_AI_BLACK = 16;
const int ADV_AI_PIECE_MASK = 7, ADV_AI_COLOR_MASK = 24;
const int ADV_AI_EMPTY = 0;

// 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;
    }
}

// A helper class to store a move and its evaluated score.
class MoveScore {
    array<int>@ move;
    int score;
    MoveScore(array<int>@ m, int s) { @move = m; score = s; }

    // Implement opCmp for sorting. Returns > 0 if this score is greater.
    int opCmp(const MoveScore &in other) const {
        return score - other.score;
    }
}

class AdvancedAI {

    // --- AI Configuration ---
    int search_depth; 
    int OPENING_PLY_MAX;
    uint CANDIDATE_MOVE_COUNT;
    int CANDIDATE_SCORE_THRESHOLD;
    
    // --- Evaluation Data ---
    array<array<int>> pst_pawn;
    array<array<int>> pst_knight;
    array<array<int>> pst_bishop;
    array<array<int>> pst_rook;
    array<array<int>> pst_queen;
    array<array<int>> pst_king;
    array<int> piece_values;
    
    // --- Helper for Sorting ---
    private SimplifiedBoardState@ current_sbs_for_sorting;

    AdvancedAI() {
        // Initialize AI Configuration in constructor
        search_depth = 3;
        OPENING_PLY_MAX = 5;
        CANDIDATE_MOVE_COUNT = 4;
        CANDIDATE_SCORE_THRESHOLD = 50;

        // Standard piece values in centipawns
        piece_values.resize(7);
        piece_values[ADV_AI_EMPTY] = 0;
        piece_values[ADV_AI_PAWN] = 100;
        piece_values[ADV_AI_KNIGHT] = 320;
        piece_values[ADV_AI_BISHOP] = 330;
        piece_values[ADV_AI_ROOK] = 500;
        piece_values[ADV_AI_QUEEN] = 900;
        piece_values[ADV_AI_KING] = 20000;

        // Piece-Square Tables (PSQTs) - bonus points for piece placement
        const array<int> PAWN_PST_FLAT = {
            0,  0,  0,  0,  0,  0,  0,  0, 50, 50, 50, 50, 50, 50, 50, 50, 10, 10, 20, 30, 30, 20, 10, 10,
             5,  5, 10, 25, 25, 10,  5,  5,  0,  0,  0, 20, 20,  0,  0,  0,  5, -5,-10,  0,  0,-10, -5,  5,
             5, 10, 10,-20,-20, 10, 10,  5,  0,  0,  0,  0,  0,  0,  0,  0
        };
        const array<int> KNIGHT_PST_FLAT = {
            -50,-40,-30,-30,-30,-30,-40,-50,-40,-20,  0,  0,  0,  0,-20,-40,-30,  0, 10, 15, 15, 10,  0,-30,
            -30,  5, 15, 20, 20, 15,  5,-30,-30,  0, 15, 20, 20, 15,  0,-30,-30,  5, 10, 15, 15, 10,  5,-30,
            -40,-20,  0,  5,  5,  0,-20,-40,-50,-40,-30,-30,-30,-30,-40,-50
        };
        const array<int> BISHOP_PST_FLAT = {
            -20,-10,-10,-10,-10,-10,-10,-20,-10,  0,  0,  0,  0,  0,  0,-10,-10,  0,  5, 10, 10,  5,  0,-10,
            -10,  5,  5, 10, 10,  5,  5,-10,-10,  0, 10, 10, 10, 10,  0,-10,-10, 10, 10, 10, 10, 10, 10,-10,
            -10,  5,  0,  0,  0,  0,  5,-10,-20,-10,-10,-10,-10,-10,-10,-20
        };
        const array<int> ROOK_PST_FLAT = {
              0,  0,  0,  0,  0,  0,  0,  0,  5, 10, 10, 10, 10, 10, 10,  5, -5,  0,  0,  0,  0,  0,  0, -5,
             -5,  0,  0,  0,  0,  0,  0, -5, -5,  0,  0,  0,  0,  0,  0, -5, -5,  0,  0,  0,  0,  0,  0, -5,
             -5,  0,  0,  0,  0,  0,  0, -5,  0,  0,  0,  5,  5,  0,  0,  0
        };
        const array<int> QUEEN_PST_FLAT = {
            -20,-10,-10, -5, -5,-10,-10,-20,-10,  0,  0,  0,  0,  0,  0,-10,-10,  0,  5,  5,  5,  5,  0,-10,
             -5,  0,  5,  5,  5,  5,  0, -5,  0,  0,  5,  5,  5,  5,  0, -5,-10,  5,  5,  5,  5,  5,  0,-10,
            -10,  0,  5,  0,  0,  0,  0,-10,-20,-10,-10, -5, -5,-10,-10,-20
        };
        const array<int> KING_PST_FLAT = {
            -30,-40,-40,-50,-50,-40,-40,-30,-30,-40,-40,-50,-50,-40,-40,-30,-30,-40,-40,-50,-50,-40,-40,-30,
            -30,-40,-40,-50,-50,-40,-40,-30,-20,-30,-30,-40,-40,-30,-30,-20,-10,-20,-20,-20,-20,-20,-20,-10,
             20, 20,  0,  0,  0,  0, 20, 20, 20, 30, 10,  0,  0, 10, 30, 20
        };

        init_pst(pst_pawn, PAWN_PST_FLAT);
        init_pst(pst_knight, KNIGHT_PST_FLAT);
        init_pst(pst_bishop, BISHOP_PST_FLAT);
        init_pst(pst_rook, ROOK_PST_FLAT);
        init_pst(pst_queen, QUEEN_PST_FLAT);
        init_pst(pst_king, KING_PST_FLAT);
    }

    void init_pst(array<array<int>> &out pst, const array<int> &in flat_pst) {
        pst.resize(8);
        for(int r = 0; r < 8; r++) {
            pst[r].resize(8);
            for(int c = 0; c < 8; c++) {
                pst[r][c] = flat_pst[r * 8 + c];
            }
        }
    }

    string move_to_string(const array<int>@ move) {
        if (move is null || move.length() < 4) return "null";
        string s = "";
        s += string::chr(97 + move[1]);
        s += string::chr(56 - move[0]);
        s += string::chr(97 + move[3]);
        s += string::chr(56 - move[2]);
        return s;
    }

    array<int>@ get_move(SimplifiedBoardState@ sbs, replay_rand@ rrnd, int ply_count) {
        array<array<int>>@ moves = generate_moves(sbs, sbs.current_player, false);
        if (moves.length() == 0) return null;

        array<MoveScore@> scored_moves;
		//puts("--- AI Turn: " + (sbs.current_player == ADV_AI_WHITE ? "White" : "Black") + " ---");
		//puts("Searching " + moves.length() + " legal moves at depth " + search_depth);
        for (uint i = 0; i < moves.length(); i++) {
            array<int>@ current_move = moves[i];
            SimplifiedBoardState@ next_sbs = deep_copy_sbs(sbs);
            apply_move(next_sbs, current_move[0], current_move[1], current_move[2], current_move[3], 0);
            int score = -negamax(next_sbs, search_depth - 1, -999999, 999999);
            MoveScore@ ms = MoveScore(current_move, score);
            scored_moves.insertLast(ms);
			puts("Move: " + move_to_string(current_move) + " | Score: " + score);
        }

        // Sort moves by score, descending
        scored_moves.sortDesc();
		//puts("--- Sorted Moves ---");
		//for (uint i = 0; i < scored_moves.length(); i++) {
			// puts("#" + (i+1) + ": " + move_to_string(scored_moves[i].move) + " | Score: " + scored_moves[i].score);
		// }
        // --- Opening Move Randomness ---
        if (ply_count < OPENING_PLY_MAX && scored_moves.length() > 1) {
            array<array<int>@> candidate_moves;
            candidate_moves.insertLast(scored_moves[0].move);
            int best_score = scored_moves[0].score;
            
            // Find other top moves within the score threshold
            for (uint i = 1; i < scored_moves.length() && i < CANDIDATE_MOVE_COUNT; i++) {
                if (best_score - scored_moves[i].score < CANDIDATE_SCORE_THRESHOLD) {
                    candidate_moves.insertLast(scored_moves[i].move);
                } else {
                    break; // Moves are sorted, no need to check further.
                }
            }
            
            // If we have multiple good moves, pick one randomly
            if(candidate_moves.length() > 1) {
                int choice_index = rrnd.rand() % candidate_moves.length();
                //puts("AI choosing randomly from " + candidate_moves.length() + " opening moves. Picked #" + (choice_index + 1));
                return candidate_moves[choice_index];
            }
        }
        
        // Default behavior: return the best move
        //puts("AI picking best move: " + move_to_string(scored_moves[0].move) + " with score " + scored_moves[0].score);
        return scored_moves[0].move;
    }

    // --- Search Algorithm ---

    int negamax(SimplifiedBoardState@ sbs, int depth, int alpha, int beta) {
        if (depth == 0) {
            return quiescence(sbs, alpha, beta);
        }

        array<array<int>>@ moves = generate_moves(sbs, sbs.current_player, false);
        if (moves.length() == 0) {
            if (is_king_in_check(sbs, sbs.current_player)) {
                return -20000 - depth; // Checkmate
            }
            return 0; // Stalemate
        }
        
        @this.current_sbs_for_sorting = sbs;
        sort_moves(moves);
        @this.current_sbs_for_sorting = null;

        for (uint i = 0; i < moves.length(); i++) {
            array<int>@ m = moves[i];
            
            SimplifiedBoardState@ next_sbs = deep_copy_sbs(sbs);
            apply_move(next_sbs, m[0], m[1], m[2], m[3], 0);

            int score = -negamax(next_sbs, depth - 1, -beta, -alpha);

            if (score >= beta) {
                return beta; 
            }
            if (score > alpha) {
                alpha = score;
            }
        }
        return alpha;
    }
    
	int quiescence(SimplifiedBoardState@ sbs, int alpha, int beta) {
		int stand_pat = evaluate(sbs);
		if(stand_pat >= beta) return beta;
		if(alpha < stand_pat) alpha = stand_pat;

		array<array<int>>@ moves = generate_moves(sbs, sbs.current_player, true); 
		
		// This early exit is okay, but can be removed for clarity as the loop below handles it.
		if (moves.length() == 0) {
			return alpha; 
		}
		
		@this.current_sbs_for_sorting = sbs;
		sort_moves(moves);
		@this.current_sbs_for_sorting = null;

		for (uint i = 0; i < moves.length(); i++) {
			array<int>@ m = moves[i];
			
			SimplifiedBoardState@ next_sbs = deep_copy_sbs(sbs);
			apply_move(next_sbs, m[0], m[1], m[2], m[3], 0);
			int score = -quiescence(next_sbs, -beta, -alpha);

			if(score >= beta) {
				return beta;
			}
			if(score > alpha) {
				alpha = score;
			}
		}
		return alpha;
	}

    // --- Evaluation Function ---
	// Helper function to print the board for debugging


	// Helper function to print the board for debugging
// Helper function to print the board for debugging
	string board_to_string(const array<array<int>> &in board) {
		string result = "\n";
		for (int r = 0; r < 8; r++) {
			result += (8 - r) + " ";
			for (int c = 0; c < 8; c++) {
				int p = board[r][c];
				string piece_char = ".";
				int type = p & ADV_AI_PIECE_MASK;
				int color = p & ADV_AI_COLOR_MASK;
				
				// Determine the base character (lowercase for black)
				if (type == ADV_AI_PAWN)   piece_char = "p";
				if (type == ADV_AI_KNIGHT) piece_char = "n";
				if (type == ADV_AI_BISHOP) piece_char = "b";
				if (type == ADV_AI_ROOK)   piece_char = "r";
				if (type == ADV_AI_QUEEN)  piece_char = "q";
				if (type == ADV_AI_KING)   piece_char = "k";
				
				// --- THIS IS THE CORRECTED PART ---
				// If the piece is White, use the uppercase version.
				if (color == ADV_AI_WHITE) {
					if (type == ADV_AI_PAWN)   piece_char = "P";
					if (type == ADV_AI_KNIGHT) piece_char = "N";
					if (type == ADV_AI_BISHOP) piece_char = "B";
					if (type == ADV_AI_ROOK)   piece_char = "R";
					if (type == ADV_AI_QUEEN)  piece_char = "Q";
					if (type == ADV_AI_KING)   piece_char = "K";
				}
				
				result += piece_char + " ";
			}
			result += "\n";
		}
		result += "  a b c d e f g h\n";
		return result;
	}
	int evaluate(SimplifiedBoardState@ sbs) {
		int score = 0;
		// ... (the rest of your evaluation function as before)
		for (int r = 0; r < 8; r++) {
			for (int c = 0; c < 8; c++) {
				int piece = sbs.board[r][c];
				if (piece != ADV_AI_EMPTY) {
					int piece_type = piece & ADV_AI_PIECE_MASK;
					int color_mult = (piece & ADV_AI_COLOR_MASK) == ADV_AI_WHITE ? 1 : -1;
					
					score += piece_values[piece_type] * color_mult;

					if ((piece & ADV_AI_COLOR_MASK) == ADV_AI_WHITE) {
						if (piece_type == ADV_AI_PAWN) score += pst_pawn[r][c];
						if (piece_type == ADV_AI_KNIGHT) score += pst_knight[r][c];
						if (piece_type == ADV_AI_BISHOP) score += pst_bishop[r][c];
						if (piece_type == ADV_AI_ROOK) score += pst_rook[r][c];
						if (piece_type == ADV_AI_QUEEN) score += pst_queen[r][c];
						if (piece_type == ADV_AI_KING) score += pst_king[r][c];
					} else { // Black
						if (piece_type == ADV_AI_PAWN) score -= pst_pawn[7-r][c];
						if (piece_type == ADV_AI_KNIGHT) score -= pst_knight[7-r][c];
						if (piece_type == ADV_AI_BISHOP) score -= pst_bishop[7-r][c];
						if (piece_type == ADV_AI_ROOK) score -= pst_rook[7-r][c];
						if (piece_type == ADV_AI_QUEEN) score -= pst_queen[7-r][c];
						if (piece_type == ADV_AI_KING) score -= pst_king[7-r][c];
					}
				}
			}
		}
		
		int final_score = (sbs.current_player == ADV_AI_WHITE) ? score : -score;
		
		// --- START OF DEBUGGING LOGS ---
		// puts("--- Evaluating Board ---");
		// puts(board_to_string(sbs.board));
		// puts("Raw score (White's POV): " + score);
		// puts("Player to move: " + (sbs.current_player == ADV_AI_WHITE ? "White" : "Black"));
		// puts("Final score (Player's POV): " + final_score);
		// --- END OF DEBUGGING LOGS ---

		return final_score;
	}

    // --- Move Ordering ---

    void sort_moves(array<array<int>>@ moves) {
        if (moves.length() <= 1) return;

        for (uint i = 1; i < moves.length(); i++) {
            array<int>@ key = moves[i];
            int j = int(i) - 1;
            while (j >= 0 && compare_moves(key, moves[j])) {
                moves[j + 1] = moves[j];
                j = j - 1;
            }
            moves[j + 1] = key;
        }
    }
    
    bool compare_moves(const array<int>@ &in a, const array<int>@ &in b) const {
        return score_move_internal(a) > score_move_internal(b);
    }
    
    private int score_move_internal(const array<int>@ move) const {
        if (@current_sbs_for_sorting == null || @move == null) return 0;
        int r1 = move[0], c1 = move[1], r2 = move[2], c2 = move[3];
        int score = 0;
        int moving_piece = current_sbs_for_sorting.board[r1][c1];
        if (moving_piece == ADV_AI_EMPTY) return 0;
        int moving_piece_type = moving_piece & ADV_AI_PIECE_MASK;
        int captured_piece_type = current_sbs_for_sorting.board[r2][c2] & ADV_AI_PIECE_MASK;
        
        if (moving_piece_type == ADV_AI_PAWN && c1 != c2 && captured_piece_type == ADV_AI_EMPTY) {
            captured_piece_type = ADV_AI_PAWN;
        }

        if (captured_piece_type != ADV_AI_EMPTY) {
            score = 10 * piece_values[captured_piece_type] - piece_values[moving_piece_type];
        }
        if (moving_piece_type == ADV_AI_PAWN && (r2 == 0 || r2 == 7)) {
            score += piece_values[ADV_AI_QUEEN];
        }
        return score;
    }
    
    array<array<int>>@ generate_moves(SimplifiedBoardState@ sbs, int color, bool captures_only) {
        array<array<int>>@ pseudo_moves = generate_pseudo_legal_moves(sbs, color, captures_only);
        array<array<int>>@ legal_moves = {};

        for (uint i = 0; i < pseudo_moves.length(); i++) {
            array<int>@ m = pseudo_moves[i];
            
            SimplifiedBoardState@ next_sbs = deep_copy_sbs(sbs);
            apply_move(next_sbs, m[0], m[1], m[2], m[3], 0);
            
            if (!is_king_in_check(next_sbs, color)) {
                legal_moves.insertLast(m);
            }
        }
        return legal_moves;
    }

    private array<array<int>>@ generate_pseudo_legal_moves(SimplifiedBoardState@ sbs, int color, bool captures_only) {
        array<array<int>>@ moves = {};
        for (int r = 0; r < 8; r++) {
            for (int c = 0; c < 8; c++) {
                if ((sbs.board[r][c] & ADV_AI_COLOR_MASK) == color) {
                    int piece_type = sbs.board[r][c] & ADV_AI_PIECE_MASK;
                    switch(piece_type) {
                        case ADV_AI_PAWN:   add_pawn_moves(sbs, r, c, color, captures_only, moves); break;
                        case ADV_AI_KNIGHT: add_knight_moves(sbs, r, c, color, captures_only, moves); break;
                        case ADV_AI_BISHOP: add_sliding_moves(sbs, r, c, color, captures_only, moves, false, true); break;
                        case ADV_AI_ROOK:   add_sliding_moves(sbs, r, c, color, captures_only, moves, true, false); break;
                        case ADV_AI_QUEEN:  add_sliding_moves(sbs, r, c, color, captures_only, moves, true, true); break;
                        case ADV_AI_KING:   add_king_moves(sbs, r, c, color, captures_only, moves); break;
                    }
                }
            }
        }
        return moves;
    }

    private void add_pawn_moves(SimplifiedBoardState@ sbs, int r, int c, int color, bool captures_only, array<array<int>>@ moves) {
        int fwd = (color == ADV_AI_WHITE) ? -1 : 1;
        int start_row = (color == ADV_AI_WHITE) ? 6 : 1;

        if (!captures_only) {
            if (is_on_board(r + fwd, c) && sbs.board[r + fwd][c] == ADV_AI_EMPTY) {
                array<int>@ m1 = {r, c, r + fwd, c};
                moves.insertLast(m1);
                if (r == start_row && sbs.board[r + 2*fwd][c] == ADV_AI_EMPTY) {
                    array<int>@ m2 = {r, c, r + 2*fwd, c};
                    moves.insertLast(m2);
                }
            }
        }

        for (int dc = -1; dc <= 1; dc += 2) {
            int r2 = r + fwd;
            int c2 = c + dc;
            if (is_on_board(r2, c2)) {
                if ((sbs.board[r2][c2] & ADV_AI_COLOR_MASK) != color && sbs.board[r2][c2] != ADV_AI_EMPTY) {
                    array<int>@ m3 = {r, c, r2, c2};
                    moves.insertLast(m3);
                }
                if (r2 == sbs.en_passant_r && c2 == sbs.en_passant_c) {
                    array<int>@ m4 = {r, c, r2, c2};
                    moves.insertLast(m4);
                }
            }
        }
    }

    private void add_knight_moves(SimplifiedBoardState@ sbs, int r, int c, int color, bool captures_only, array<array<int>>@ moves) {
        const array<int> dr = {-2, -2, -1, -1, 1, 1, 2, 2};
        const array<int> dc = {-1, 1, -2, 2, -2, 2, -1, 1};

        for (uint i = 0; i < 8; i++) {
            int r2 = r + dr[i];
            int c2 = c + dc[i];
            if (is_on_board(r2, c2)) {
                int target_piece = sbs.board[r2][c2];
                if ((target_piece & ADV_AI_COLOR_MASK) != color) {
                    if (!captures_only || target_piece != ADV_AI_EMPTY) {
                        array<int>@ new_move = {r, c, r2, c2};
                        moves.insertLast(new_move);
                    }
                }
            }
        }
    }

    private void add_sliding_moves(SimplifiedBoardState@ sbs, int r, int c, int color, bool captures_only, array<array<int>>@ moves, bool straight, bool diagonal) {
        const array<int> dr = {-1, 1, 0, 0, -1, -1, 1, 1};
        const array<int> dc = {0, 0, -1, 1, -1, 1, -1, 1};
        int start_i = straight ? 0 : 4;
        int end_i = diagonal ? 8 : (straight ? 4 : 0);

        for (int i = start_i; i < end_i; i++) {
            for (int d = 1; d < 8; d++) {
                int r2 = r + dr[i] * d;
                int c2 = c + dc[i] * d;
                if (!is_on_board(r2, c2)) break;
                int target_piece = sbs.board[r2][c2];
                if (target_piece == ADV_AI_EMPTY) {
                    if (!captures_only) {
                        array<int>@ new_move = {r, c, r2, c2};
                        moves.insertLast(new_move);
                    }
                } else {
                    if ((target_piece & ADV_AI_COLOR_MASK) != color) {
                        array<int>@ new_move = {r, c, r2, c2};
                        moves.insertLast(new_move);
                    }
                    break;
                }
            }
        }
    }
    
    private void add_king_moves(SimplifiedBoardState@ sbs, int r, int c, int color, bool captures_only, array<array<int>>@ moves) {
        for (int dr = -1; dr <= 1; dr++) {
            for (int dc = -1; dc <= 1; dc++) {
                if (dr == 0 && dc == 0) continue;
                int r2 = r + dr;
                int c2 = c + dc;
                if (is_on_board(r2, c2)) {
                    int target_piece = sbs.board[r2][c2];
                    if ((target_piece & ADV_AI_COLOR_MASK) != color) {
                        if (!captures_only || target_piece != ADV_AI_EMPTY) {
                            array<int>@ new_move = {r, c, r2, c2};
                            moves.insertLast(new_move);
                        }
                    }
                }
            }
        }
        
        if (!captures_only && !is_king_in_check(sbs, color)) {
            if ((color == ADV_AI_WHITE && sbs.can_castle_w_k) || (color == ADV_AI_BLACK && sbs.can_castle_b_k)) {
                if (sbs.board[r][c+1] == ADV_AI_EMPTY && sbs.board[r][c+2] == ADV_AI_EMPTY) {
                    if (!is_square_attacked(sbs, r, c+1, color==ADV_AI_WHITE ? ADV_AI_BLACK : ADV_AI_WHITE) &&
                        !is_square_attacked(sbs, r, c+2, color==ADV_AI_WHITE ? ADV_AI_BLACK : ADV_AI_WHITE)) {
                        array<int>@ m_castle_k = {r, c, r, c+2};
                        moves.insertLast(m_castle_k);
                    }
                }
            }
            if ((color == ADV_AI_WHITE && sbs.can_castle_w_q) || (color == ADV_AI_BLACK && sbs.can_castle_b_q)) {
                if (sbs.board[r][c-1] == ADV_AI_EMPTY && sbs.board[r][c-2] == ADV_AI_EMPTY && sbs.board[r][c-3] == ADV_AI_EMPTY) {
                    if (!is_square_attacked(sbs, r, c-1, color==ADV_AI_WHITE ? ADV_AI_BLACK : ADV_AI_WHITE) &&
                        !is_square_attacked(sbs, r, c-2, color==ADV_AI_WHITE ? ADV_AI_BLACK : ADV_AI_WHITE)) {
                        array<int>@ m_castle_q = {r, c, r, c-2};
                        moves.insertLast(m_castle_q);
                    }
                }
            }
        }
    }

    SimplifiedBoardState@ deep_copy_sbs(SimplifiedBoardState@ sbs) {
        array<array<int>>@ new_board = {};
        new_board.resize(8);
        for(int r = 0; r < 8; r++) {
            new_board[r].resize(8);
            for(int c = 0; c < 8; c++) {
                new_board[r][c] = sbs.board[r][c];
            }
        }

        return SimplifiedBoardState(
            new_board, sbs.current_player,
            sbs.can_castle_w_k, sbs.can_castle_w_q, sbs.can_castle_b_k, sbs.can_castle_b_q,
            sbs.en_passant_r, sbs.en_passant_c,
            sbs.king_pos_w_r, sbs.king_pos_w_c, sbs.king_pos_b_r, sbs.king_pos_b_c
        );
    }
    
    void apply_move(SimplifiedBoardState@ sbs, int r1, int c1, int r2, int c2, int promotion_choice) {
        int piece = sbs.board[r1][c1];
        int piece_type = piece & ADV_AI_PIECE_MASK;
        int piece_color = piece & ADV_AI_COLOR_MASK;
        
        if (piece_type == ADV_AI_PAWN && c1 != c2 && sbs.board[r2][c2] == ADV_AI_EMPTY) {
            sbs.board[r1][c2] = ADV_AI_EMPTY;
        }

        sbs.board[r2][c2] = piece;
        sbs.board[r1][c1] = ADV_AI_EMPTY;

        if (piece_type == ADV_AI_PAWN && (r2 == 0 || r2 == 7)) {
            sbs.board[r2][c2] = ADV_AI_QUEEN | piece_color;
        }

        if (piece_type == ADV_AI_KING) {
            if (piece_color == ADV_AI_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; }
        }

        if (piece_type == ADV_AI_KING) {
            if (piece_color == ADV_AI_WHITE) { sbs.can_castle_w_k = false; sbs.can_castle_w_q = false; }
            else { sbs.can_castle_b_k = false; sbs.can_castle_b_q = false; }
        }
        if (piece_type == ADV_AI_ROOK) {
            if (piece_color == ADV_AI_WHITE) {
                if(r1 == 7 && c1 == 0) sbs.can_castle_w_q = false;
                if(r1 == 7 && c1 == 7) sbs.can_castle_w_k = false;
            } else {
                if(r1 == 0 && c1 == 0) sbs.can_castle_b_q = false;
                if(r1 == 0 && c1 == 7) sbs.can_castle_b_k = false;
            }
        }
        
        if (piece_type == ADV_AI_KING && abs(c1 - c2) == 2) {
            if (c2 > c1) {
                sbs.board[r1][5] = sbs.board[r1][7];
                sbs.board[r1][7] = ADV_AI_EMPTY;
            } else {
                sbs.board[r1][3] = sbs.board[r1][0];
                sbs.board[r1][0] = ADV_AI_EMPTY;
            }
        }

        sbs.en_passant_r = -1;
        sbs.en_passant_c = -1;
        if (piece_type == ADV_AI_PAWN && abs(r1 - r2) == 2) {
            sbs.en_passant_r = (r1 + r2) / 2;
            sbs.en_passant_c = c1;
        }

        sbs.current_player = (sbs.current_player == ADV_AI_WHITE) ? ADV_AI_BLACK : ADV_AI_WHITE;
    }

    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]&ADV_AI_COLOR_MASK)!=color || (sbs.board[r2][c2]&ADV_AI_COLOR_MASK)==color) return false;
        
        int piece_type = sbs.board[r1][c1] & ADV_AI_PIECE_MASK;
        if (!is_legal_shape(sbs, piece_type, r1, c1, r2, c2, color)) return false;

        SimplifiedBoardState@ next_sbs = deep_copy_sbs(sbs);
        apply_move(next_sbs, r1, c1, r2, c2, 0);
        bool in_check = is_king_in_check(next_sbs, color); 
        
        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 == ADV_AI_KING) {
            if (r1 == r2 && abs(dc) == 2) {
                return can_castle(sbs, color, dc > 0);
            }
            return abs(dr) <= 1 && abs(dc) <= 1;
        }
        if (type == ADV_AI_KNIGHT) return (abs(dr) == 1 && abs(dc) == 2) || (abs(dr) == 2 && abs(dc) == 1);
        if (type == ADV_AI_QUEEN) {
            if (abs(dr) == abs(dc) || dr == 0 || dc == 0) return is_path_clear(sbs, r1, c1, r2, c2);
        }
        if (type == ADV_AI_ROOK) {
            if (dr == 0 || dc == 0) return is_path_clear(sbs, r1, c1, r2, c2);
        }
        if (type == ADV_AI_BISHOP) {
            if (abs(dr) == abs(dc)) return is_path_clear(sbs, r1, c1, r2, c2);
        }
        if (type == ADV_AI_PAWN) {
            int fwd = (color == ADV_AI_WHITE) ? -1 : 1;
            if(dr == fwd && dc == 0 && sbs.board[r2][c2] == ADV_AI_EMPTY) return true;
            if(dr == 2*fwd && dc == 0 && (r1 == (color==ADV_AI_WHITE?6:1)) && sbs.board[r1+fwd][c1]==ADV_AI_EMPTY && sbs.board[r2][c2]==ADV_AI_EMPTY) return true;
            if(dr == fwd && abs(dc) == 1 && (sbs.board[r2][c2] & ADV_AI_COLOR_MASK) != color && sbs.board[r2][c2] != ADV_AI_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 == ADV_AI_WHITE) ? 7 : 0;
        int opponent_color = (color == ADV_AI_WHITE) ? ADV_AI_BLACK : ADV_AI_WHITE;
        if (is_king_in_check(sbs, color)) return false;
        if (kingside) {
            if ((color == ADV_AI_WHITE && !sbs.can_castle_w_k) || (color == ADV_AI_BLACK && !sbs.can_castle_b_k)) return false;
            if (sbs.board[r][5] != ADV_AI_EMPTY || sbs.board[r][6] != ADV_AI_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 == ADV_AI_WHITE && !sbs.can_castle_w_q) || (color == ADV_AI_BLACK && !sbs.can_castle_b_q)) return false;
            if (sbs.board[r][1] != ADV_AI_EMPTY || sbs.board[r][2] != ADV_AI_EMPTY || sbs.board[r][3] != ADV_AI_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] != ADV_AI_EMPTY) return false;
            r += dr; c += dc;
        }
        return true;
    }
    
    bool is_king_in_check(SimplifiedBoardState@ sbs, int color) {
        int r = (color == ADV_AI_WHITE) ? sbs.king_pos_w_r : sbs.king_pos_b_r;
        int c = (color == ADV_AI_WHITE) ? sbs.king_pos_w_c : sbs.king_pos_b_c;
        if(r < 0 || c < 0) return true;
        return is_square_attacked(sbs, r, c, (color == ADV_AI_WHITE) ? ADV_AI_BLACK : ADV_AI_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]&ADV_AI_COLOR_MASK)==by_color && is_legal_shape_for_attack(sbs, sbs.board[r1][c1]&ADV_AI_PIECE_MASK,r1,c1,r,c,by_color)) return true;
        return false;
    }

    bool is_legal_shape_for_attack(SimplifiedBoardState@ sbs, int type, int r1, int c1, int r2, int c2, int color) {
        int dr = r2 - r1, dc = c2 - c1;
        if (type == ADV_AI_PAWN) {
            int fwd = (color == ADV_AI_WHITE) ? -1 : 1;
            if (dr == fwd && abs(dc) == 1) return true;
            return false;
        }
        if (type == ADV_AI_KING) return abs(dr) <= 1 && abs(dc) <= 1;
        if (type == ADV_AI_KNIGHT) return (abs(dr) == 1 && abs(dc) == 2) || (abs(dr) == 2 && abs(dc) == 1);
        if (type == ADV_AI_QUEEN) return (abs(dr) == abs(dc) || dr == 0 || dc == 0) && is_path_clear(sbs, r1, c1, r2, c2);
        if (type == ADV_AI_ROOK) return (dr == 0 || dc == 0) && is_path_clear(sbs, r1, c1, r2, c2);
        if (type == ADV_AI_BISHOP) return abs(dr) == abs(dc) && is_path_clear(sbs, r1, c1, r2, c2);
        return false;
    }

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