#include "../lib/std.cpp"
#include "../lib/math/math.cpp"
#include "../lib/drawing/circle.cpp"
#include "../lib/easing/cubic.cpp" // Include for smooth animation
#include "ai_advanced.cpp" // <-- Include the new AI file
//#include "ai_simple.cpp" // <-- Include the new AI file

// --- Game Constants ---
// Constants are still needed for game state and piece identification in the main file.
const int EMPTY = 0, KING = 1, KNIGHT = 2, PAWN = 3, QUEEN = 4, ROOK = 5, BISHOP = 6;
const int WHITE = 8, BLACK = 16;
const int PIECE_MASK = 7, COLOR_MASK = 24;
const int PLAYING = 0, CHECKMATE_WHITE = 1, CHECKMATE_BLACK = 2, STALEMATE = 3;
const int AWAITING_PROMOTION = 4;

// Point values for sorting captured pieces
const array<int> PIECE_VALUES = {0, 100, 3, 1, 9, 5, 3}; 
// Defines the vertical order of captured piece groups
const array<int> PIECE_ORDER = {PAWN, KNIGHT, BISHOP, ROOK, QUEEN, KING};
// Defines the pieces available for pawn promotion.
const array<int> PROMOTION_PIECES = {QUEEN, ROOK, BISHOP, KNIGHT};
const int board_y_offset = 20;
// --- Helper class for replay-consistent randomness ---

// Dummy class for the replay encoders, required for replay-consistent randomness.
class replay_rand_dummy : enemy_base {
  void init(script@, scriptenemy@ self) {
    self.auto_physics(false);
  }
}

// --- Helper class for replay-consistent randomness ---
// This class uses entity positions to encode/decode a seed value, ensuring
// that the same seed is used for both live gameplay and replays. This is
// crucial for preventing replay desynchronization.
class replay_rand {
  scene@ g;
  array<uint> encoders;
  uint seed = 1;
  
  bool _is_initialised = false;
  bool _failed = false;

  int ecx = 0;
  int ecy = 0;
  uint frame_counter = 0;
  
  uint NUM_ENCODERS = 4;

  replay_rand() {
    @g = get_scene();
  }

  void step() {
    if (_failed) return;
    
    // On the first frame, create the encoder entities.
    if (frame_counter == 1) {
        if (!init_encoders()) {
            puts("replay_rand: FAILED to init encoders. Randomness will not be replay-consistent.");
            seed = timestamp_now();
            _is_initialised = true;
            _failed = true;
            return;
        }
    }

    if (is_replay()) {
        // In a replay, wait until frame 20 to decode the seed from replayed entity positions.
        if (frame_counter == 20) {
            if(decode_seed()) {
                _is_initialised = true;
            } else {
                puts("replay_rand: FAILED to decode seed. Randomness will not be replay-consistent.");
                seed = timestamp_now(); // Fallback
                _is_initialised = true;
                _failed = true;
            }
        }
    } else { // Live game
        // In a live game, we can generate the seed immediately.
        if (frame_counter == 1) {
            seed = timestamp_now() + get_time_us();
            _is_initialised = true;
            puts("replay_rand: Generated live seed: " + seed);
        }
        // At frame 10, move the entities to encode the seed for any future replay.
        if (frame_counter == 10) {
            encode_seed();
        }
    }

    frame_counter++;
  }

  uint rand() {
    // The AI's first move is delayed, so it's safe to assume _is_initialised will be true.
    if (!_is_initialised) {
       // This might be called on frame 0 in the editor, but not during gameplay logic.
       return 0;
    }
    return seed = (seed * 16807) % 2147483647;
  }
  
  private bool init_encoders() {
    controllable@ ec = @controller_controllable(0);
    if (@ec == null) {
      return false;
    }

    ecx = int(ec.x());
    ecy = int(ec.y());
    for (uint i = 0; i < NUM_ENCODERS; i++) {
      scriptenemy@ ent = create_scriptenemy(replay_rand_dummy());
      ent.x(ecx);
      ent.y(ecy);
      // THIS IS THE CORRECTED LINE
      g.add_entity(@ent.as_entity(), true); 
      encoders.insertLast(ent.id());
    }
    return true;
  }

  private void encode_seed() {
    uint val_to_encode = seed;
    puts("replay_rand: Encoding seed: " + val_to_encode);
    for (uint i = 0; i < NUM_ENCODERS; i++) {
      entity@ ent = @entity_by_id(encoders[i]);
      if (@ent == null) {
        puts("replay_rand: Encoding seed failed - entity missing.");
        _failed = true;
        return;
      }
      // Move entities to encode the seed value in their x-coordinates
      ent.x(ecx - 128 + ((val_to_encode >> (i * 8)) & 0xFF));
      ent.y(ecy + 100);
    }
  }

  private bool decode_seed() {
    uint decoded_val = 0;
    for (uint i = 0; i < NUM_ENCODERS; i++) {
      entity@ ent = @entity_by_id(encoders[i]);
      if (@ent == null) {
        puts("replay_rand: Seed reconstruction failed - entity missing.");
        return false;
      }
      int enc_x = int(round(ent.x())) + 128 - ecx;
      int enc_y = int(round(ent.y()));
      if (enc_y == ecy) {
        puts("replay_rand: Desync frames missing, seed reconstruction failed.");
        return false;
      }
      decoded_val |= enc_x << (i * 8);
    }
    seed = decoded_val;
    puts("replay_rand: Reconstructed seed: " + seed);
    return true;
  }
};

// Class to store sprite information, including origin for centering.
class PieceSprite {
    sprites@ spr;
    string name;
    float origin_x, origin_y;

    PieceSprite(string sprite_set, string name, float origin_x, float origin_y) {
        @this.spr = create_sprites();
        this.spr.add_sprite_set(sprite_set);
        this.name = name;
        this.origin_x = origin_x;
        this.origin_y = origin_y;
    }
}

// --- Class to handle piece animation ---
class PieceAnimation {
    int piece;
    int to_r, to_c;
    float from_x, from_y, to_x, to_y;
    float progress, duration;
    
    // Properties to handle rook animation during castling
    bool is_castling;
    int castle_rook_from_c, castle_rook_to_c;
    int castle_rook_piece;
    
    // This flag is now used to trigger the promotion state.
    bool is_promotion;

    PieceAnimation(int piece, int to_r, int to_c, float from_x, float from_y, float to_x, float to_y, float duration) {
        this.piece = piece;
        this.to_r = to_r; this.to_c = to_c;
        this.from_x = from_x; this.from_y = from_y;
        this.to_x = to_x; this.to_y = to_y;
        this.duration = duration;
        this.progress = 0;
        this.is_castling = false;
        this.is_promotion = false;
    }
}

// --- Main Game Logic Class ---
class KnightChess : enemy_base {
    array<array<int>> board;
    int current_player = WHITE;
    int game_state = PLAYING;
    int selected_r = -1, selected_c = -1;
    int king_pos_w_r, king_pos_w_c, king_pos_b_r, king_pos_b_c;

    int pre_move_r1 = -1, pre_move_c1 = -1;
    int pre_move_r2 = -1, pre_move_c2 = -1;

    int en_passant_c = -1;
    int en_passant_r = -1; 
    
    int last_move_r1 = -1, last_move_c1 = -1, last_move_r2 = -1, last_move_c2 = -1;

    bool white_king_in_check = false;
    bool black_king_in_check = false;

    // Stores the square of the pawn awaiting promotion choice.
    int promotion_r = -1, promotion_c = -1;

    bool can_castle_w_k, can_castle_w_q, can_castle_b_k, can_castle_b_q;

    dictionary captured_by_white;
    dictionary captured_by_black;

    scene@ g; scriptenemy@ self;
    canvas@ cvs; textfield@ txt;
    replay_rand@ rrnd;
    string player_character;
    
    float board_scale;
    float board_size_px;
    float tile_size;
    int last_r = -1, last_c = -1, last_mouse_st = 0;
    
    dictionary piece_sprites;
    PieceAnimation@ active_animation = null;
    AdvancedAI@ ai; // AI instance
	//SimpleAI@ ai; // AI instance
    int ply_count = 0; // Number of half-moves made in the game.

    KnightChess(replay_rand@ rrnd, string player_character, float board_scale) {
        @g = get_scene();
        @this.rrnd = rrnd;
        this.player_character = player_character;
        this.board_scale = board_scale;
        this.board_size_px = 400.0 * board_scale;
        this.tile_size = this.board_size_px / 8.0;
        @ai = AdvancedAI(); // Instantiate the AI
		//@ai = SimpleAI(); // Instantiate the AI
    }
    
    void init(script@, scriptenemy@ self) {
        @this.self = self;
        @cvs = create_canvas(false, self.layer(), 1);
        @txt = create_textfield();
        txt.set_font("sans_bold", 20);
        self.auto_physics(false);
        load_sprites();
        start_new_game();
    }

    void load_sprites() {
        @piece_sprites[''+(KING   | WHITE)] = PieceSprite("dustman", "idle", 0.5, 0.95);
        @piece_sprites[''+(QUEEN  | WHITE)] = PieceSprite("dustgirl", "idle", 0.5, 0.95);
        @piece_sprites[''+(ROOK   | WHITE)] = PieceSprite("props1", "furniture_9", 0.493, 0.884);
        @piece_sprites[''+(BISHOP | WHITE)] = PieceSprite("props6", "npc_2", 0.388, 0.940);
        @piece_sprites[''+(KNIGHT | WHITE)] = PieceSprite("dustworth", "idle", 0.5, 0.95);
        @piece_sprites[''+(PAWN   | WHITE)] = PieceSprite("dustkid", "idle", 0.5, 1.10);
        
        @piece_sprites[''+(KING   | BLACK)] = PieceSprite("dustwraith", "idle", 0.5, 0.95);
        @piece_sprites[''+(QUEEN  | BLACK)] = PieceSprite("leafsprite", "idle", 0.5, 0.95);
        @piece_sprites[''+(ROOK   | BLACK)] = PieceSprite("props2", "temple_5", 0.423, 0.900);
        @piece_sprites[''+(BISHOP | BLACK)] = PieceSprite("props6", "npc_1", 4.388, 0.940);
        @piece_sprites[''+(KNIGHT | BLACK)] = PieceSprite("slimeboss", "idle", 0.5, 0.95);
        @piece_sprites[''+(PAWN   | BLACK)] = PieceSprite("trashking", "idle", 0.5, 1.10);
    }

    void start_new_game() {
        board.resize(8);
        for(int i = 0; i < 8; i++) {
            board[i].resize(8);
            for(int j=0; j<8; j++) board[i][j] = EMPTY;
        }
        
        board[0][0] = ROOK | BLACK; board[0][7] = ROOK | BLACK;
        board[0][1] = KNIGHT | BLACK; board[0][6] = KNIGHT | BLACK;
        board[0][2] = BISHOP | BLACK; board[0][5] = BISHOP | BLACK;
        board[0][3] = QUEEN | BLACK; board[0][4] = KING | BLACK;
        for(int i = 0; i < 8; i++) board[1][i] = PAWN | BLACK;
        
        board[7][0] = ROOK | WHITE; board[7][7] = ROOK | WHITE;
        board[7][1] = KNIGHT | WHITE; board[7][6] = KNIGHT | WHITE;
        board[7][2] = BISHOP | WHITE; board[7][5] = BISHOP | WHITE;
        board[7][3] = QUEEN | WHITE; board[7][4] = KING | WHITE;
        for(int i = 0; i < 8; i++) board[6][i] = PAWN | WHITE;
        
        king_pos_b_r = 0; king_pos_b_c = 4;
        king_pos_w_r = 7; king_pos_w_c = 4;

        can_castle_w_k = true; can_castle_w_q = true;
        can_castle_b_k = true; can_castle_b_q = true;

        en_passant_c = -1; en_passant_r = -1;
        last_move_r1 = -1; last_move_c1 = -1;
        last_move_r2 = -1; last_move_c2 = -1;
        
        white_king_in_check = false;
        black_king_in_check = false;
        
        pre_move_r1 = -1; pre_move_c1 = -1;
        pre_move_r2 = -1; pre_move_c2 = -1;

        captured_by_white.deleteAll();
        captured_by_black.deleteAll();
        ply_count = 0;
    }
    
    void step() {
        if (game_state == AWAITING_PROMOTION) {
            handle_promotion_input();
            return;
        }

        if (game_state != PLAYING) return;

        if (@active_animation != null) {
            active_animation.progress += DT;
            if (active_animation.progress >= active_animation.duration) {
                board[active_animation.to_r][active_animation.to_c] = active_animation.piece;
                
                if (active_animation.is_castling) {
                    board[active_animation.to_r][active_animation.castle_rook_to_c] = active_animation.castle_rook_piece;
                }
                
                bool was_promotion = active_animation.is_promotion;
                int promo_r = active_animation.to_r, promo_c = active_animation.to_c;
                int promo_color = active_animation.piece & COLOR_MASK;
                
                @active_animation = null;

                if (was_promotion) {
                    if (promo_color == WHITE) {
                        game_state = AWAITING_PROMOTION;
                        promotion_r = promo_r;
                        promotion_c = promo_c;
                    } else {
                        // AI always promotes to queen for simplicity
                        board[promo_r][promo_c] = QUEEN | BLACK;
                        current_player = WHITE;
                        white_king_in_check = is_king_in_check(WHITE);
                        black_king_in_check = is_king_in_check(BLACK);
                        check_for_game_end();
                    }
                } else {
                    current_player = (current_player == WHITE) ? BLACK : WHITE;
                    white_king_in_check = is_king_in_check(WHITE);
                    black_king_in_check = is_king_in_check(BLACK);
                    check_for_game_end();
                    if (current_player == BLACK) self.state_timer(30);
                }
            }
        }
        
        if (current_player == BLACK || @active_animation != null) {
            handle_pre_move_input();
        }

        if (@active_animation == null && game_state == PLAYING) {
            if (current_player == WHITE) {
                if (pre_move_r2 != -1) {
                    if (is_valid_move(pre_move_r1, pre_move_c1, pre_move_r2, pre_move_c2, WHITE)) {
                        make_move(pre_move_r1, pre_move_c1, pre_move_r2, pre_move_c2);
                    }
                    pre_move_r1 = -1; pre_move_c1 = -1;
                    pre_move_r2 = -1; pre_move_c2 = -1;
                } 
                else if (pre_move_r1 != -1) {
                    selected_r = pre_move_r1;
                    selected_c = pre_move_c1;
                    pre_move_r1 = -1;
                    pre_move_c1 = -1;
                }
                
                handle_regular_move_input();

            } else { 
                self.state_timer(self.state_timer() - 1);
                if (self.state_timer() <= 0) {
                    handle_ai_turn();
                }
            }
        }
    }
    
	void handle_promotion_input() {
		int player_index = self.player_index(); if (player_index == -1) return;
		float mx = g.mouse_x_world(player_index, self.layer()), my = g.mouse_y_world(player_index, self.layer());
		int mst = g.mouse_state(player_index);
		bool click = (mst & 4) == 0 && (last_mouse_st & 4) != 0;

		if (click) {
			float lft = self.x() - board_size_px/2.0;
			float top = self.y() - board_size_px/2.0 - board_y_offset;
			
			float ui_y = top - tile_size * 0.6;
			float ui_x = lft + promotion_c * tile_size + tile_size / 2;
			float spacing = tile_size * 0.8;
			float clickable_radius = spacing / 2.0;
			float total_width = (PROMOTION_PIECES.length() - 1) * spacing;
			float start_x = ui_x - total_width / 2.0;

			for (uint i = 0; i < PROMOTION_PIECES.length(); i++) {
				float piece_x = start_x + i * spacing;
				if (mx >= piece_x - clickable_radius && mx <= piece_x + clickable_radius &&
					my >= ui_y - clickable_radius && my <= ui_y + clickable_radius) {
					
					board[promotion_r][promotion_c] = PROMOTION_PIECES[i] | WHITE;
					game_state = PLAYING;
					promotion_r = -1;
					promotion_c = -1;
					current_player = BLACK;
					
					white_king_in_check = is_king_in_check(WHITE);
					black_king_in_check = is_king_in_check(BLACK);
					check_for_game_end();
					self.state_timer(30);
					break;
				}
			}
		}
		last_mouse_st = mst;
	}

    void handle_regular_move_input() {
        if(game_state != PLAYING) return;
        int player_index = self.player_index(); if (player_index == -1) return;
        float lft = self.x() - board_size_px/2.0, top = self.y() - board_size_px/2.0 - board_y_offset;
        float mx = g.mouse_x_world(player_index, self.layer()), my = g.mouse_y_world(player_index, self.layer());
        int r = int(floor((my-top)/tile_size)), c = int(floor((mx-lft)/tile_size));
        int mst = g.mouse_state(player_index);
        bool click = (mst & 4) == 0 && (last_mouse_st & 4) != 0;
        
        if (click && is_on_board(r, c)) {
            if (selected_r == -1) {
                if ((board[r][c] & COLOR_MASK) == WHITE) { selected_r = r; selected_c = c; }
            } else {
                if (is_valid_move(selected_r, selected_c, r, c, WHITE)) {
                    make_move(selected_r, selected_c, r, c);
                }
                selected_r = -1; selected_c = -1;
            }
        }
        last_mouse_st = mst;
    }

    void handle_pre_move_input() {
        if(game_state != PLAYING) return;
        int player_index = self.player_index(); if (player_index == -1) return;
        float lft = self.x() - board_size_px/2.0, top = self.y() - board_size_px/2.0 - board_y_offset;
        float mx = g.mouse_x_world(player_index, self.layer()), my = g.mouse_y_world(player_index, self.layer());
        int r = int(floor((my-top)/tile_size)), c = int(floor((mx-lft)/tile_size));
        int mst = g.mouse_state(player_index);
        bool click = (mst & 4) == 0 && (last_mouse_st & 4) != 0;
        
        if (click && is_on_board(r, c)) {
            // Determine the logical piece at the clicked square, considering any active animation.
            int piece_at_rc = board[r][c];
            if (@active_animation != null) {
                if (r == active_animation.to_r && c == active_animation.to_c) {
                    piece_at_rc = active_animation.piece;
                } else if (active_animation.is_castling && r == active_animation.to_r && c == active_animation.castle_rook_to_c) {
                    piece_at_rc = active_animation.castle_rook_piece;
                }
            }
            
            if (pre_move_r1 == -1) {
                if ((piece_at_rc & COLOR_MASK) == WHITE) {
                    pre_move_r1 = r; pre_move_c1 = c;
                }
            } else {
                if (r == pre_move_r1 && c == pre_move_c1) {
                    pre_move_r1 = -1; pre_move_c1 = -1; pre_move_r2 = -1; pre_move_c2 = -1;
                } else if ((piece_at_rc & COLOR_MASK) == WHITE) {
                    pre_move_r1 = r; pre_move_c1 = c;
                    pre_move_r2 = -1; pre_move_c2 = -1;
                } else {
                    // Temporarily fast-forward board state to validate premove
                    bool premove_is_valid = false;
                    int original_anim_dest_piece = EMPTY;
                    int original_rook_dest_piece = EMPTY;
                    if (@active_animation != null) {
                        original_anim_dest_piece = board[active_animation.to_r][active_animation.to_c];
                        board[active_animation.to_r][active_animation.to_c] = active_animation.piece;
                        if (active_animation.is_castling) {
                            original_rook_dest_piece = board[active_animation.to_r][active_animation.castle_rook_to_c];
                            board[active_animation.to_r][active_animation.castle_rook_to_c] = active_animation.castle_rook_piece;
                        }
                    }
                    premove_is_valid = is_valid_move(pre_move_r1, pre_move_c1, r, c, WHITE);
                    if (@active_animation != null) {
                        board[active_animation.to_r][active_animation.to_c] = original_anim_dest_piece;
                        if (active_animation.is_castling) {
                            board[active_animation.to_r][active_animation.castle_rook_to_c] = original_rook_dest_piece;
                        }
                    }
                    if (premove_is_valid) {
                        pre_move_r2 = r; pre_move_c2 = c;
                    } else {
                        pre_move_r1 = -1; pre_move_c1 = -1;
                        pre_move_r2 = -1; pre_move_c2 = -1;
                    }
                }
            }
        }
        last_mouse_st = mst;
    }

    void handle_ai_turn() {
        if (@active_animation != null) return;
        
        // Package the current state
        SimplifiedBoardState@ sbs = SimplifiedBoardState(
            board, current_player,
            can_castle_w_k, can_castle_w_q, can_castle_b_k, can_castle_b_q,
            en_passant_r, en_passant_c,
            king_pos_w_r, king_pos_w_c, king_pos_b_r, king_pos_b_c
        );

        // Get move from AI
        array<int>@ move = ai.get_move(sbs, rrnd, ply_count);

        if (move !is null && move.length() == 4) {
            make_move(move[0], move[1], move[2], move[3]);
        } else {
            // No legal moves found by AI, end the game.
            if (is_king_in_check(BLACK)) game_state = CHECKMATE_BLACK;
            else game_state = STALEMATE;
        }
    }

    void make_move(int r1, int c1, int r2, int c2) {
        if (@active_animation != null) return;
        
        ply_count++;
        int piece_to_move = board[r1][c1];
        int piece_type = piece_to_move & PIECE_MASK;
        int piece_color = piece_to_move & COLOR_MASK;

        last_move_r1 = r1; last_move_c1 = c1; last_move_r2 = r2; last_move_c2 = c2;
        if (piece_type == PAWN && c1 != c2 && board[r2][c2] == EMPTY) {
            int captured_pawn_r = r1, captured_pawn_c = c2;
            int ep_captured_piece = board[captured_pawn_r][captured_pawn_c];
            if(ep_captured_piece != EMPTY) {
                int captured_type = ep_captured_piece & PIECE_MASK;
                dictionary@ captured_dict = (ep_captured_piece & COLOR_MASK) == WHITE ? @captured_by_black : @captured_by_white;
                int count = 0;
                if(captured_dict.get(''+captured_type, count)) { captured_dict.set(''+captured_type, count + 1); }
                else { captured_dict.set(''+captured_type, 1); }
                board[captured_pawn_r][captured_pawn_c] = EMPTY;
            }
        }
        en_passant_c = -1; en_passant_r = -1;
        if (piece_type == KING) {
            if (piece_color == WHITE) { can_castle_w_k = false; can_castle_w_q = false; } else { can_castle_b_k = false; can_castle_b_q = false; }
        }
        if (piece_type == ROOK) {
            if (piece_color == WHITE) {
                if (r1 == 7 && c1 == 0) can_castle_w_q = false;
                if (r1 == 7 && c1 == 7) can_castle_w_k = false;
            } else {
                if (r1 == 0 && c1 == 0) can_castle_b_q = false;
                if (r1 == 0 && c1 == 7) can_castle_b_k = false;
            }
        }
        if (piece_type == PAWN && abs(r1 - r2) == 2) { en_passant_r = (r1 + r2) / 2; en_passant_c = c1; }
        int captured_piece = board[r2][c2];
        if (captured_piece != EMPTY) {
            int captured_type = captured_piece & PIECE_MASK;
            dictionary@ captured_dict = (captured_piece & COLOR_MASK) == WHITE ? @captured_by_black : @captured_by_white;
            int count = 0;
            if(captured_dict.get(''+captured_type, count)) { captured_dict.set(''+captured_type, count + 1); }
            else { captured_dict.set(''+captured_type, 1); }
        }
        if (piece_type == KING) {
            if (piece_color == WHITE) { king_pos_w_r = r2; king_pos_w_c = c2; } else { king_pos_b_r = r2; king_pos_b_c = c2; }
        }
        float lft = self.x()-board_size_px/2.0, top = self.y()-board_size_px/2.0 - board_y_offset;
        float from_x = lft + c1*tile_size + tile_size/2, from_y = top + r1*tile_size + tile_size/2;
        float to_x = lft + c2*tile_size + tile_size/2, to_y = top + r2*tile_size + tile_size/2;
        @active_animation = PieceAnimation(piece_to_move, r2, c2, from_x, from_y, to_x, to_y, 0.4);
        if (piece_type == PAWN && (r2 == 0 || r2 == 7)) { active_animation.is_promotion = true; }
        board[r1][c1] = EMPTY; board[r2][c2] = EMPTY;
        if (piece_type == KING && abs(c1 - c2) == 2) {
            active_animation.is_castling = true;
            if (c2 > c1) { active_animation.castle_rook_from_c = 7; active_animation.castle_rook_to_c = 5; }
            else { active_animation.castle_rook_from_c = 0; active_animation.castle_rook_to_c = 3; }
            active_animation.castle_rook_piece = board[r1][active_animation.castle_rook_from_c];
            board[r1][active_animation.castle_rook_from_c] = EMPTY;
        }
    }
    
    // The rest of the functions are just copied from the AI file, as they are now internal to the main game logic.
    bool is_valid_move(int r1, int c1, int r2, int c2, int color) {
        SimplifiedBoardState@ sbs = SimplifiedBoardState(board, color, can_castle_w_k, can_castle_w_q, can_castle_b_k, can_castle_b_q, en_passant_r, en_passant_c, king_pos_w_r, king_pos_w_c, king_pos_b_r, king_pos_b_c);
        return ai.is_valid_move(sbs, r1, c1, r2, c2, color);
    }

    bool is_king_in_check(int color) {
        SimplifiedBoardState@ sbs = SimplifiedBoardState(board, color, can_castle_w_k, can_castle_w_q, can_castle_b_k, can_castle_b_q, en_passant_r, en_passant_c, king_pos_w_r, king_pos_w_c, king_pos_b_r, king_pos_b_c);
        return ai.is_king_in_check(sbs, color);
    }

    void check_for_game_end() {
        bool can_move = false;
        for (int r1 = 0; r1 < 8; r1++) for (int c1 = 0; c1 < 8; c1++)
            if ((board[r1][c1] & COLOR_MASK) == current_player)
                for (int r2 = 0; r2 < 8; r2++) for (int c2 = 0; c2 < 8; c2++)
                    if (is_valid_move(r1, c1, r2, c2, current_player))
                        { can_move = true; r1=8;c1=8;r2=8;c2=8; }
        if (can_move) return;

        if (is_king_in_check(current_player)) {
            if (current_player == WHITE) { // Player is checkmated
                game_state = CHECKMATE_WHITE;
                
                // Reduce finesse by one rank
                g.combo_break_count(1); // S -> A

                // End the level
                g.end_level(self.x(), self.y());
            } else { // AI (BLACK) is checkmated
                game_state = CHECKMATE_BLACK;
                g.end_level(self.x(), self.y());
            }
        } else { 
            game_state = STALEMATE; 
			g.combo_break_count(1); // S -> A

			// End the level
			g.end_level(self.x(), self.y());
        }
    }
    
    bool is_on_board(int r, int c) { return r >= 0 && r < 8 && c >= 0 && c < 8; }

    void draw(float subframe) {
        float lft = self.x()-board_size_px/2.0, top = self.y()-board_size_px/2.0 - board_y_offset;
        for(int r=0;r<8;r++) for(int c=0;c<8;c++)
            g.draw_rectangle_world(self.layer(),0,lft+c*tile_size,top+r*tile_size,lft+(c+1)*tile_size,top+(r+1)*tile_size,0,((r+c)%2==0)?0xFFE3C16F:0xFFB88B4A);
        
        if(selected_r!=-1) g.draw_rectangle_world(self.layer(),1,lft+selected_c*tile_size,top+selected_r*tile_size,lft+(selected_c+1)*tile_size,top+(selected_r+1)*tile_size,0,0x5500FF00);
        
        if(last_move_r1 != -1) {
            uint highlight_color_light = 0x66FFF89E, highlight_color_dark = 0x33FFF89E;
            uint start_square_color = ((last_move_r1 + last_move_c1) % 2 == 0) ? highlight_color_light : highlight_color_dark;
            g.draw_rectangle_world(self.layer(), 1, lft+last_move_c1*tile_size, top+last_move_r1*tile_size, lft+(last_move_c1+1)*tile_size, top+(last_move_r1+1)*tile_size, 0, start_square_color);
            uint end_square_color = ((last_move_r2 + last_move_c2) % 2 == 0) ? highlight_color_light : highlight_color_dark;
            g.draw_rectangle_world(self.layer(), 1, lft+last_move_c2*tile_size, top+last_move_r2*tile_size, lft+(last_move_c2+1)*tile_size, top+(last_move_r2+1)*tile_size, 0, end_square_color);
        }

        if (white_king_in_check) g.draw_rectangle_world(self.layer(), 1, lft+king_pos_w_c*tile_size, top+king_pos_w_r*tile_size, lft+(king_pos_w_c+1)*tile_size, top+(king_pos_w_r+1)*tile_size, 0, 0x77FF0000);
        if (black_king_in_check) g.draw_rectangle_world(self.layer(), 1, lft+king_pos_b_c*tile_size, top+king_pos_b_r*tile_size, lft+(king_pos_b_c+1)*tile_size, top+(king_pos_b_r+1)*tile_size, 0, 0x77FF0000);
        if(pre_move_r1 != -1) g.draw_rectangle_world(self.layer(), 1, lft+pre_move_c1*tile_size, top+pre_move_r1*tile_size, lft+(pre_move_c1+1)*tile_size, top+(pre_move_r1+1)*tile_size, 0, 0x5500FF00);
        if(pre_move_r2 != -1) g.draw_rectangle_world(self.layer(), 1, lft+pre_move_c2*tile_size, top+pre_move_r2*tile_size, lft+(pre_move_c2+1)*tile_size, top+(pre_move_r2+1)*tile_size, 0, 0x5500FF00);

        for(int r=0;r<8;r++) for(int c=0;c<8;c++) {
            if (@active_animation != null && active_animation.from_x == lft+c*tile_size+tile_size/2 && active_animation.from_y == top+r*tile_size+tile_size/2) continue;
            if (@active_animation != null && active_animation.is_castling && r == active_animation.to_r && c == active_animation.castle_rook_from_c) continue;
            int piece = board[r][c]; if(piece==EMPTY) continue;
            draw_piece(piece, lft+c*tile_size+tile_size/2, top+r*tile_size+tile_size/2);
        }

        if (@active_animation != null) {
            float t = ease_in_out_cubic(clamp01(active_animation.progress / active_animation.duration));
            float anim_x = lerp(active_animation.from_x, active_animation.to_x, t);
            float anim_y = lerp(active_animation.from_y, active_animation.to_y, t);
            draw_piece(active_animation.piece, anim_x, anim_y);
            if (active_animation.is_castling) {
                float rook_from_x = lft + active_animation.castle_rook_from_c * tile_size + tile_size / 2;
                float rook_to_x = lft + active_animation.castle_rook_to_c * tile_size + tile_size / 2;
                float rook_anim_x = lerp(rook_from_x, rook_to_x, t);
                draw_piece(active_animation.castle_rook_piece, rook_anim_x, active_animation.to_y);
            }
        }
        
        draw_captured_pieces(WHITE, lft, top);
        draw_captured_pieces(BLACK, lft, top);
        
        if (game_state == AWAITING_PROMOTION) draw_promotion_options();

        // START of new mouse tracker code
        if (is_replay()) {
            int player_index = self.player_index();
            if (player_index != -1) {
                float mx = g.mouse_x_world(player_index, self.layer());
                float my = g.mouse_y_world(player_index, self.layer());
                
                // Draw a small red square to represent the mouse cursor
                g.draw_rectangle_world(
                    self.layer(), 10, // Draw on top
                    mx - 5, my - 5,
                    mx + 5, my + 5,
                    0, 0x33FF0000
                );
            }
        }
        // END of new mouse tracker code

        string status;
        switch(game_state){case PLAYING:status=(current_player==WHITE?"Your Turn (White)":"AI's Turn (Black)");break;case CHECKMATE_WHITE:status="Checkmate! Black Wins.";break;case CHECKMATE_BLACK:status="Checkmate! You Win!";break;case STALEMATE:status="Stalemate!";break; case AWAITING_PROMOTION: status="Choose a promotion piece"; break;}
        if(game_state==PLAYING && is_king_in_check(current_player)) status += " (Check)";
        txt.align_horizontal(0);txt.align_vertical(-1);txt.colour(0xFFFFFFFF);
        txt.text(status);
        //txt.draw_world(self.layer(),10,self.x(),self.y()-board_size_px/2.0-40,1,1,0) - board_y_offset;
    }
    
	void draw_promotion_options() {
		float lft = self.x()-board_size_px/2.0, top = self.y()-board_size_px/2.0 - board_y_offset;
		float ui_x = lft + promotion_c * tile_size + tile_size / 2;
		float ui_y = top - tile_size * 0.6; 
		float spacing = tile_size * 0.8;
		float piece_size_multiplier = 0.75;
		float total_width = (PROMOTION_PIECES.length() - 1) * spacing;
		float start_x = ui_x - total_width / 2.0;
		float panel_padding_x = tile_size * 0.4;
		float panel_padding_y = tile_size * 0.5;
		g.draw_rectangle_world(self.layer(), 1, start_x - panel_padding_x, ui_y - panel_padding_y, start_x + total_width + panel_padding_x, ui_y + panel_padding_y, 0, 0x99222222);
		for (uint i = 0; i < PROMOTION_PIECES.length(); i++) {
			float piece_x = start_x + i * spacing;
			draw_piece(PROMOTION_PIECES[i] | WHITE, piece_x, ui_y, piece_size_multiplier);
		}
	}

    void draw_captured_pieces(int capturer_color, float lft, float top) {
        dictionary@ counts = (capturer_color == WHITE) ? @captured_by_white : @captured_by_black;
        int opponent_color = (capturer_color == WHITE) ? BLACK : WHITE;
        float captured_start_x = lft + board_size_px + tile_size * 0.2 + 40;
        float size_multiplier = 0.5, horizontal_spacing = tile_size * size_multiplier * 0.7, vertical_spacing = tile_size * size_multiplier;
        float current_y_offset = 0;
        float start_y, y_dir;
        if (capturer_color == WHITE) { start_y = top + board_size_px - vertical_spacing/2; y_dir = -1; }
        else { start_y = top + vertical_spacing/2; y_dir = 1; }
        for (uint i = 0; i < PIECE_ORDER.length(); i++) {
            int piece_type = PIECE_ORDER[i];
            int count = 0;
            if (counts.get(''+piece_type, count) && count > 0) {
                float group_y = start_y + current_y_offset * y_dir;
                for (int j = 0; j < count; j++) {
                    float piece_x = captured_start_x + j * horizontal_spacing;
                    draw_piece(piece_type | opponent_color, piece_x, group_y, size_multiplier);
                }
                current_y_offset += vertical_spacing;
            }
        }
    }
    
    void draw_piece(int piece, float x, float y, float size_multiplier = 1.0) {
        PieceSprite@ piece_sprite = cast<PieceSprite@>(piece_sprites[''+piece]); if(@piece_sprite == null) return;
        sprites@ spr = piece_sprite.spr; string spr_name = piece_sprite.name;
        float scale; uint palette; int piece_type = piece & PIECE_MASK;
        switch (piece_type) {
            case KING: case KNIGHT: case QUEEN: case BISHOP: scale = 0.4 * board_scale * size_multiplier; palette = ((piece & COLOR_MASK) == WHITE ? 2 : 1); break;
            case PAWN: scale = 0.35 * board_scale * size_multiplier; palette = ((piece & COLOR_MASK) == WHITE ? 2 : 1); break;
            case ROOK: scale = 0.45 * board_scale * size_multiplier; palette = ((piece & COLOR_MASK) == WHITE ? 0 : 1); break;
        }
        rectangle@ rect = spr.get_sprite_rect(spr_name, 0);
        float ox = (rect.left() + rect.get_width() * piece_sprite.origin_x);
        float oy = (rect.top() + rect.get_height() * piece_sprite.origin_y) - 42;
        float draw_x = x - ox * scale; float draw_y = y - oy * scale;
        if(spr_name!="") spr.draw_world(self.layer(),2,spr_name,0,palette,draw_x,draw_y,0,(piece&COLOR_MASK)==WHITE?scale:-scale,scale,0xFFFFFFFF);
    }
}

// --- Main Script Class ---
class script {
    replay_rand rrnd;
    [text] float board_scale = 2.0;
    void step(int) { rrnd.step(); }
    void spawn_player(message@ msg) {
        string player_char = msg.get_string("character", "dustman");
        scriptenemy@ game_entity = create_scriptenemy(KnightChess(@rrnd, player_char, board_scale));
        game_entity.x(msg.get_float("x"));
        game_entity.y(msg.get_float("y"));
        msg.set_entity("player", @game_entity.as_entity());
    }
}