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

// --- Game Constants ---
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;
const int PUZZLE_TRANSITION = 5;
const int INCORRECT_MOVE_ANIMATION = 6;
const int PROMOTION_REVERSAL = 7;

// --- Shake Animation Constants ---
const float UNPAUSE_SHAKE_DURATION = 0.4;
const float UNPAUSE_SHAKE_INTENSITY = 4.0;
const float WRONG_MOVE_SHAKE_DURATION = 0.25;
const float WRONG_MOVE_SHAKE_INTENSITY = 2.5;

const array<int> PIECE_VALUES = {0, 100, 3, 1, 9, 5, 3};
const array<int> PIECE_ORDER = {PAWN, KNIGHT, BISHOP, ROOK, QUEEN, KING};
const array<int> PROMOTION_PIECES = {QUEEN, ROOK, BISHOP, KNIGHT};
const int board_y_offset = 20;

// Embed the puzzles file
const string EMBED_puzzles = "puzzles_1100.txt";

// 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 ---
class replay_rand {
  scene@ g;
  array<uint> encoders;
  uint initial_seed = 1;
  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;
    
    if (frame_counter == 1) {
        if (!init_encoders()) {
            puts("replay_rand: FAILED to init encoders. Randomness will not be replay-consistent.");
            _is_initialised = true;
            _failed = true;
            return;
        }
    }

    if (is_replay()) {
        if (frame_counter == 20) {
            if(decode_seed()) {
                _is_initialised = true;
            } else {
                puts("replay_rand: FAILED to decode seed. Randomness will not be replay-consistent.");
                _is_initialised = true;
                _failed = true;
            }
        }
    } else { 
        if (frame_counter == 1) {
            initial_seed = seed = timestamp_now() + get_time_us();
            _is_initialised = true;
        }
        if (frame_counter == 10) {
            encode_seed();
        }
    }

    frame_counter++;
  }

  uint rand() {
    if (!_is_initialised) {
       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);
      g.add_entity(@ent.as_entity(), true); 
      encoders.insertLast(ent.id());
    }
    return true;
  }

  private void encode_seed() {
    uint val_to_encode = initial_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;
      }
      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 a single piece's animation.
class PieceAnimation {
    int piece;
    int from_r = -1, from_c = -1;
    int to_r, to_c;
    float from_x, from_y, to_x, to_y;
    float from_scale, to_scale;
    float progress, duration;
    
    bool is_castling;
    int castle_rook_from_c, castle_rook_to_c;
    int castle_rook_piece;
    
    bool is_promotion;
    int promotion_piece = 0;
    bool is_reversing = false;
    int captured_piece = EMPTY; // Added to handle captures on incorrect moves

    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.from_scale = 1.0; this.to_scale = 1.0;
        this.duration = duration;
        this.progress = 0;
        this.is_castling = false;
        this.is_promotion = false;
    }

    PieceAnimation(int piece, int to_r, int to_c, float from_x, float from_y, float to_x, float to_y, float from_scale, float to_scale, 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.from_scale = from_scale; this.to_scale = to_scale;
        this.duration = duration;
        this.progress = 0;
        this.is_castling = false;
        this.is_promotion = false;
    }
}

// A helper class for the puzzle animation logic.
class PieceLocation {
    int r, c;
    bool handled = false;
}

// Manages the animation of the entire board from one state to another.
class BoardTransitionAnimation {
    array<PieceAnimation@> piece_anims;
    float delay;
    float duration;
    float progress;
    SimplifiedBoardState@ target_sbs;
    
    bool has_scripted_move = false;
    int scripted_move_r1, scripted_move_c1;
    int scripted_move_r2, scripted_move_c2;
    int scripted_promotion = 0;

    BoardTransitionAnimation(SimplifiedBoardState@ target_sbs, float delay, float duration) {
        @this.target_sbs = target_sbs;
        this.delay = delay;
        this.duration = duration;
        this.progress = 0;
    }
}

// Helper class for storing drawn arrows
class Arrow {
    int r1, c1, r2, c2;
}

// --- Main Game Logic Class ---
class KnightChess : enemy_base {
    array<array<int>> board;
	int PUZZLE_AI_DELAY_FRAMES = 60; // Approx. 1 second delay before puzzle op
    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;
    int last_captured_piece = EMPTY; // Added to correctly revert captures

    bool white_king_in_check = false;
    bool black_king_in_check = false;

    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;
	int last_mouse_st_marking = 0;
    
    dictionary piece_sprites;
    PieceAnimation@ active_animation = null;
    BoardTransitionAnimation@ active_board_transition = null;
    AdvancedAI@ ai;
    int ply_count = 0;
    bool game_started = false;
    int puzzles_completed = 0;
    int puzzles_to_win_goal;

    // --- Board Shake Animation ---
    float shake_timer = 0.0;
    float shake_duration_initial = 0.0;
    float shake_intensity_initial = 0.0;

    // --- State for completing puzzle via promotion ---
    float puzzle_complete_delay = -1.0;
    int completed_promo_r = -1, completed_promo_c = -1;
    int completed_promo_original_piece = -1;
	
	// --- Board Markings ---
    array<int> highlighted_squares;
    array<int> circled_squares;
	array<Arrow@> drawn_arrows;
    
	// --- Input State for Markings ---
    bool right_mouse_is_down = false;
    int drag_start_r = -1, drag_start_c = -1;
    bool is_dragging = false;

    // --- PUZZLE DATA ---
    bool is_puzzle_mode = false;
    array<string>@ puzzle_solution_moves;
    bool puzzle_just_completed = false;
    bool puzzle_penalty_applied = false;
    int current_puzzle_index = -1;
    array<string> all_puzzle_ids;
    array<string> all_puzzle_fens;
    array<string> all_puzzle_solutions;

    KnightChess(replay_rand@ rrnd, string player_character, float board_scale, int puzzles_to_win) {
        @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();
        this.puzzles_to_win_goal = puzzles_to_win;
    }
    
    void start_shake(float duration, float intensity) {
        if (shake_timer > 0 && intensity < shake_intensity_initial) {
            return;
        }
        shake_timer = duration;
        shake_duration_initial = duration;
        shake_intensity_initial = intensity;
    }

    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();
        load_puzzles_from_file();
        start_new_game();
    }


    void load_puzzles_from_file() {
        if (!has_embed_value("puzzles")) {
            puts("Error: puzzles.txt not found in embeds.");
            return;
        }
        string data = get_embed_value("puzzles");
        array<string>@ lines = data.split("\n");

        for (uint i = 1; i < lines.length(); i++) { // Start from 1 to skip header
            string line = lines[i];
            if (line.length() == 0) continue;
            
            array<string>@ parts = line.split(",");
            if (parts.length() >= 3) {
                all_puzzle_ids.insertLast(parts[0]);
                all_puzzle_fens.insertLast(parts[1]);
                all_puzzle_solutions.insertLast(parts[2]);
            }
        }
    }

    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;
        puzzles_completed = 0;
		clear_all_markings();
    }
	
	void clear_all_markings() {
		highlighted_squares.resize(0);
		circled_squares.resize(0);
		drawn_arrows.resize(0);
	}

    void handle_marking_input() {
		if (game_state == PUZZLE_TRANSITION || @active_board_transition != null) {
			right_mouse_is_down = false;
			is_dragging = false; 
			return;
		}

		int player_index = self.player_index();
		if (player_index == -1) return;

		int mst = g.mouse_state(player_index);
		bool right_press = (mst & 8) != 0 && (last_mouse_st_marking & 8) == 0;
		bool right_release = (mst & 8) == 0 && (last_mouse_st_marking & 8) != 0;

		float lft = self.x() - board_size_px / 2.0;
		float top = self.y() - board_size_px / 2.0 - board_y_offset;
		float mx = g.mouse_x_world(player_index, self.layer());
		float my = g.mouse_y_world(player_index, self.layer());
		int r = int(floor((my - top) / tile_size));
		int c = int(floor((mx - lft) / tile_size));

		if (right_press) {
			if (is_on_board(r, c)) {
				right_mouse_is_down = true;
				drag_start_r = r;
				drag_start_c = c;
				is_dragging = false;
			} else {
				clear_all_markings();
			}
		}

		if (right_mouse_is_down) {
			if (r != drag_start_r || c != drag_start_c) {
				is_dragging = true;
			}
		}

		if (right_release) {
			if (right_mouse_is_down) {
				if (is_dragging) {
					if (is_on_board(r, c) && (r != drag_start_r || c != drag_start_c)) {
						Arrow@ new_arrow = Arrow();
						new_arrow.r1 = drag_start_r;
						new_arrow.c1 = drag_start_c;
						new_arrow.r2 = r;
						new_arrow.c2 = c;
						drawn_arrows.insertLast(new_arrow);
					}
				} else { // It's a click, not a drag.
					int clicked_r = drag_start_r;
					int clicked_c = drag_start_c;
					int key = clicked_r * 8 + clicked_c;
					bool removed = false;

					
					// Check for highlighted squares to remove
					int highlight_index = highlighted_squares.find(key);
					if (highlight_index != -1) {
						highlighted_squares.removeAt(highlight_index);
						removed = true;
					}

					// Check for circled squares to remove
					int circle_index = circled_squares.find(key);
					if (circle_index != -1) {
						circled_squares.removeAt(circle_index);
						removed = true;
					}

					// If no mark was at the location, add a new one.
					if (!removed) {
						int piece = board[clicked_r][clicked_c];
						if ((piece & COLOR_MASK) == BLACK) {
							circled_squares.insertLast(key);
						} else { // Empty or White
							highlighted_squares.insertLast(key);
						}
					}
				}
			}
			right_mouse_is_down = false;
			is_dragging = false;
		}
		
		last_mouse_st_marking = mst;
	}
    
	void step() {
		// --- Handle Shake ---
		if(shake_timer > 0) {
			shake_timer -= DT;
			if(shake_timer < 0) {
				shake_timer = 0;
			}
		}

		handle_marking_input();
		
		if(puzzle_complete_delay > 0) {
			puzzle_complete_delay -= DT;
			if(puzzle_complete_delay <= 0) {
				puzzle_complete_delay = -1.0;
				
				if(completed_promo_r != -1) {
					board[completed_promo_r][completed_promo_c] = completed_promo_original_piece;
				}

				puzzles_completed++;
				if (puzzles_completed >= puzzles_to_win_goal) {
					game_state = CHECKMATE_BLACK;
					g.end_level(self.x(), self.y());
				} else {
					load_next_puzzle();
				}
			}
			return; // Block other processing during this delay.
		}

		if (!game_started) {
			if (rrnd._is_initialised) {
				load_next_puzzle();
				game_started = true;
			} else {
				return;
			}
		}

		if (game_state == PUZZLE_TRANSITION) {
			if (@active_board_transition != null) {
				active_board_transition.progress += DT;
				if (active_board_transition.progress >= active_board_transition.delay + active_board_transition.duration) {
					SimplifiedBoardState@ sbs = active_board_transition.target_sbs;
					board = sbs.board;
					current_player = sbs.current_player;
					can_castle_w_k = sbs.can_castle_w_k; can_castle_w_q = sbs.can_castle_w_q;
					can_castle_b_k = sbs.can_castle_b_k; can_castle_b_q = sbs.can_castle_b_q;
					en_passant_r = sbs.en_passant_r; en_passant_c = sbs.en_passant_c;
					king_pos_w_r = sbs.king_pos_w_r; king_pos_w_c = sbs.king_pos_w_c;
					king_pos_b_r = sbs.king_pos_b_r; king_pos_b_c = sbs.king_pos_b_c;

					white_king_in_check = is_king_in_check(WHITE);
					black_king_in_check = is_king_in_check(BLACK);
					
					captured_by_white.deleteAll();
					captured_by_black.deleteAll();
					calculate_captured_pieces_from_sbs(sbs, captured_by_white, captured_by_black);
					
					bool has_move = active_board_transition.has_scripted_move;
					int r1 = active_board_transition.scripted_move_r1, c1 = active_board_transition.scripted_move_c1;
					int r2 = active_board_transition.scripted_move_r2, c2 = active_board_transition.scripted_move_c2;
					int promo = active_board_transition.scripted_promotion;

					@active_board_transition = null;
					game_state = PLAYING;
					
					if (current_player != WHITE && has_move) {
						make_move(r1, c1, r2, c2, promo);
					}
				}
			}
			return;
		}

		if (game_state == AWAITING_PROMOTION) {
			handle_promotion_input();
			return;
		}

		if (game_state == PROMOTION_REVERSAL) {
			if (@active_animation != null) {
				active_animation.progress += DT;
				if (active_animation.progress >= active_animation.duration) {
					// Restore the board to its state before the pawn move
					board[active_animation.to_r][active_animation.to_c] = active_animation.piece; // Pawn back to original square
					board[active_animation.from_r][active_animation.from_c] = active_animation.captured_piece; // Restore captured piece
					
					// Revert ply count.
					ply_count--;

					// Decrement the captured count from when the pawn first moved.
					if (active_animation.captured_piece != EMPTY) {
						int captured_type = active_animation.captured_piece & PIECE_MASK;
						// The pawn was white, so it captured a black piece. White's capture list needs updating.
						dictionary@ captured_dict = @captured_by_white;
						int count = 0;
						if (captured_dict.get(''+captured_type, count) && count > 0) {
							if (count > 1) {
								captured_dict.set(''+captured_type, count - 1);
							} else {
								captured_dict.delete(''+captured_type);
							}
						}
					}

					@active_animation = null;
					game_state = PLAYING;
				}
			}
			return;
		}

		if (game_state == INCORRECT_MOVE_ANIMATION) {
			if (@active_animation != null) {
				if (!active_animation.is_reversing) {
					active_animation.progress += DT;
					if (active_animation.progress >= active_animation.duration) {
						active_animation.progress = active_animation.duration;
						active_animation.is_reversing = true;
						start_shake(WRONG_MOVE_SHAKE_DURATION, WRONG_MOVE_SHAKE_INTENSITY);
					}
				} else {
					active_animation.progress -= DT;
					if (active_animation.progress <= 0) {
						board[active_animation.from_r][active_animation.from_c] = active_animation.piece;
						board[active_animation.to_r][active_animation.to_c] = active_animation.captured_piece;
						@active_animation = null;
						game_state = PLAYING;
					}
				}
			}
			return;
		}

		if (game_state != PLAYING) return;

		if (@active_animation != null) {
			active_animation.progress += DT;
			if (active_animation.progress >= active_animation.duration) {
				// Put the piece on the board after its animation completes
				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;
				int promo_piece_type = active_animation.promotion_piece;
				
				@active_animation = null;
				
				if (is_puzzle_mode && puzzle_just_completed) {
					puzzle_just_completed = false;
					
					// Calculate check status on final board state before delay
					white_king_in_check = is_king_in_check(WHITE);
					black_king_in_check = is_king_in_check(BLACK);
					
					puzzle_complete_delay = 1.0; 
					completed_promo_r = promo_r;
					completed_promo_c = promo_c;
					completed_promo_original_piece = PAWN | promo_color;
					return;
				}

				if (was_promotion) {
					if (promo_piece_type != EMPTY) { // Handle scripted AI promotions
						board[promo_r][promo_c] = promo_piece_type | promo_color;
						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);
						if (is_puzzle_mode && puzzle_solution_moves.length() > 0 && current_player != WHITE) {
							self.state_timer(PUZZLE_AI_DELAY_FRAMES);
						} else {
							check_for_game_end();
							if (current_player == BLACK) self.state_timer(30);
						}
					} else if (promo_color == WHITE) { // Handle all player promotions (normal and puzzle)
						game_state = AWAITING_PROMOTION;
						promotion_r = promo_r;
						promotion_c = promo_c;
					} else { // Handle normal AI promotion
						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);

					if (is_puzzle_mode && puzzle_solution_moves.length() > 0 && current_player != WHITE) {
						 self.state_timer(PUZZLE_AI_DELAY_FRAMES); 
					} else {
						 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 { // current_player is BLACK
				self.state_timer(self.state_timer() - 1);

				if (self.state_timer() <= 0) {
					if (is_puzzle_mode && puzzle_solution_moves.length() > 0) {
						int r1, c1, r2, c2, promotion;
						parse_move_string(puzzle_solution_moves[0], r1, c1, r2, c2, promotion);
						make_move(r1, c1, r2, c2, promotion);
					} else if (!is_puzzle_mode) {
						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) {
				clear_all_markings();
				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) {
						
						int clicked_piece = PROMOTION_PIECES[i];

						if (is_puzzle_mode) {
							string move_made_str = move_to_algebraic(last_move_r1, last_move_c1, promotion_r, promotion_c, clicked_piece);

							if (puzzle_solution_moves.length() > 0 && move_made_str == puzzle_solution_moves[0]) {
								// CORRECT promotion choice
								board[promotion_r][promotion_c] = clicked_piece | WHITE;
								puzzle_solution_moves.removeAt(0);
								if (puzzle_solution_moves.length() == 0) {
									puzzle_just_completed = true;
								}
							} else {
								// INCORRECT promotion choice - Setup reversal animation
								if (!puzzle_penalty_applied) {
									int current_breaks = g.combo_break_count();
									int new_breaks;
									if (current_breaks == 0) new_breaks = 1; else if (current_breaks == 1) new_breaks = 2; else if (current_breaks <= 3) new_breaks = 4; else if (current_breaks <= 5) new_breaks = 6; else new_breaks = current_breaks + 1;
									g.combo_break_count(new_breaks);
									puzzle_penalty_applied = true;
								}
								
								// Animate the pawn returning to its original square.
								float from_x = lft + promotion_c * tile_size + tile_size / 2;
								float from_y = top + promotion_r * tile_size + tile_size / 2;
								float to_x = lft + last_move_c1 * tile_size + tile_size / 2;
								float to_y = top + last_move_r1 * tile_size + tile_size / 2;

								// The piece to animate is the pawn that was on the promotion square.
								int moving_piece = board[promotion_r][promotion_c];
								
								@active_animation = PieceAnimation(moving_piece, last_move_r1, last_move_c1, from_x, from_y, to_x, to_y, 0.4);
								// The 'from' square for the board state restoration is the promotion square
								active_animation.from_r = promotion_r;
								active_animation.from_c = promotion_c;
								active_animation.captured_piece = last_captured_piece;
								
								// Remove the pawn from the board so it doesn't draw statically during the animation
								board[promotion_r][promotion_c] = EMPTY;
								
								game_state = PROMOTION_REVERSAL;
								start_shake(WRONG_MOVE_SHAKE_DURATION, WRONG_MOVE_SHAKE_INTENSITY);
								
								// Clear the last move highlight
								last_move_r1 = -1;
								
								promotion_r = -1;
								promotion_c = -1;
								
								last_mouse_st = mst;
								return;
							}
						} else {
							// Normal game promotion
							board[promotion_r][promotion_c] = clicked_piece | WHITE;
						}
						
						if (puzzle_just_completed) {
							puzzle_just_completed = false; // Prevent normal turn switch
							// Calculate final check status before delay
							white_king_in_check = is_king_in_check(WHITE);
							black_king_in_check = is_king_in_check(BLACK);
							puzzle_complete_delay = 1.0; 
							completed_promo_r = promotion_r;
							completed_promo_c = promotion_c;
							completed_promo_original_piece = PAWN | WHITE;
							game_state = PLAYING;
							promotion_r = -1;
							promotion_c = -1;
							last_mouse_st = mst;
							return;
						}

						// Common logic for successful promotion (puzzle or normal)
						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);
						
						if (!is_puzzle_mode) {
							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) {
			clear_all_markings();
			if (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) {
			clear_all_markings();
			if (is_on_board(r, c)) {
				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) { // No piece is pre-selected
					// First click: select a piece to premove
					if ((piece_at_rc & COLOR_MASK) == WHITE) {
						pre_move_r1 = r; pre_move_c1 = c;
					}
				} else { // A piece is already pre-selected
					// Second click: either set destination or deselect
					if (r == pre_move_r1 && c == pre_move_c1) {
						// Clicked the same piece again: deselect.
						pre_move_r1 = -1; pre_move_c1 = -1;
						pre_move_r2 = -1; pre_move_c2 = -1;
					} else {
						// Clicked any other square: set as destination for the premove.
						// This now correctly handles "recaptures" where the destination
						// square is currently occupied by a friendly piece.
						// The move's validity will be checked when it's our turn.
						pre_move_r2 = r;
						pre_move_c2 = c;
					}
				}
			}
		}
		last_mouse_st = mst;
	}

    void handle_ai_turn() {
        if (@active_animation != null) return;
        
        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
        );

        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 {
            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, int promotion_choice = EMPTY) {
		clear_all_markings();
		if (@active_animation != null) return;

		last_captured_piece = board[r2][c2]; // Store captured piece for potential revert
		int piece_to_move = board[r1][c1];
		int piece_type = piece_to_move & PIECE_MASK;
		bool is_potential_promotion = (piece_type == PAWN && (r2 == 0 || r2 == 7));

		if (is_puzzle_mode) {
			if (current_player == WHITE) {
				// For player moves in puzzle mode, defer promotion verification.
				// Verify non-promotion moves immediately.
				if (!is_potential_promotion) {
					string move_made_str = move_to_algebraic(r1, c1, r2, c2);
					if (puzzle_solution_moves.length() > 0 && move_made_str == puzzle_solution_moves[0]) {
						puzzle_solution_moves.removeAt(0);
						if (puzzle_solution_moves.length() == 0) {
							puzzle_just_completed = true;
							
							// Pre-calculate check status on a temporary board to show highlight during animation.
							array<array<int>> temp_board = board;
							int piece = temp_board[r1][c1];
							temp_board[r2][c2] = piece;
							temp_board[r1][c1] = EMPTY;
							
							SimplifiedBoardState@ sbs = SimplifiedBoardState(
								temp_board, BLACK,
								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
							);
							
							// Account for king movement in the temporary state
							if ((piece & PIECE_MASK) == KING) {
								sbs.king_pos_w_r = r2; sbs.king_pos_w_c = c2;
							}
							
							// Set the check flag for the opponent
							black_king_in_check = ai.is_king_in_check(sbs, BLACK);
						}
					} else {
						// Incorrect non-promotion move
						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(board[r1][c1], r2, c2, from_x, from_y, to_x, to_y, 0.4);
						active_animation.from_r = r1;
						active_animation.from_c = c1;
						active_animation.captured_piece = board[r2][c2];
						board[r1][c1] = EMPTY;
						
						if (!puzzle_penalty_applied) {
							int current_breaks = g.combo_break_count();
							int new_breaks;
							if (current_breaks == 0) new_breaks = 1; else if (current_breaks == 1) new_breaks = 2; else if (current_breaks <= 3) new_breaks = 4; else if (current_breaks <= 5) new_breaks = 6; else new_breaks = current_breaks + 1;
							g.combo_break_count(new_breaks);
							puzzle_penalty_applied = true;
						}
						game_state = INCORRECT_MOVE_ANIMATION;
						return;
					}
				}
			} else { // AI's scripted move
				puzzle_solution_moves.removeAt(0);
			}
		}
		
		ply_count++;
		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) {
				last_captured_piece = ep_captured_piece; // Store en-passant capture
				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; }
		
		if (last_captured_piece != EMPTY) {
			int captured_type = last_captured_piece & PIECE_MASK;
			dictionary@ captured_dict = (last_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 (is_potential_promotion) {
			active_animation.is_promotion = true;
			if (promotion_choice != EMPTY) { // For AI moves
				active_animation.promotion_piece = promotion_choice;
			}
		}
		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;
		}
	}
    // --- PUZZLE MODE LOGIC ---

    string move_to_algebraic(int r1, int c1, int r2, int c2, int promotion_piece = EMPTY) {
        string h_chars = "abcdefgh";
        string v_chars = "87654321";
        
        if (c1 < 0 || c1 > 7 || r1 < 0 || r1 > 7 || c2 < 0 || c2 > 7 || r2 < 0 || r2 > 7) {
            return "err";
        }

        string move_str = h_chars.substr(c1, 1) + v_chars.substr(r1, 1) + 
                          h_chars.substr(c2, 1) + v_chars.substr(r2, 1);
        
        if (promotion_piece != EMPTY) {
            int piece_type = promotion_piece & PIECE_MASK;
            if (piece_type == QUEEN) move_str += "q";
            else if (piece_type == ROOK) move_str += "r";
            else if (piece_type == BISHOP) move_str += "b";
            else if (piece_type == KNIGHT) move_str += "n";
        }
        return move_str;
    }

    void load_next_puzzle() {
        if(game_state == PUZZLE_TRANSITION || all_puzzle_fens.length() == 0) return;
        
        const float POST_PUZZLE_DELAY = 0.3; // Delay in seconds.
        float delay = (puzzles_completed > 0) ? POST_PUZZLE_DELAY : 0.0;
        
        current_puzzle_index = rrnd.rand() % all_puzzle_fens.length();
        puts("Selected puzzle: " + all_puzzle_ids[current_puzzle_index]);
        transition_to_puzzle(all_puzzle_fens[current_puzzle_index], all_puzzle_solutions[current_puzzle_index], delay);
    }
    
    void transition_to_puzzle(string fen, string moves, float delay) {
        is_puzzle_mode = true;
        @puzzle_solution_moves = moves.split(" ");
        puzzle_penalty_applied = false;
        
        SimplifiedBoardState@ sbs_to = parse_fen(fen);
        if (sbs_to is null) {
            puts("Error: Could not parse FEN string.");
            game_state = PLAYING;
            return;
        }

        @active_board_transition = BoardTransitionAnimation(sbs_to, delay, 1.5);
        
        if (sbs_to.current_player != WHITE && puzzle_solution_moves.length() > 0) {
            active_board_transition.has_scripted_move = true;
            parse_move_string(puzzle_solution_moves[0], 
                active_board_transition.scripted_move_r1, active_board_transition.scripted_move_c1,
                active_board_transition.scripted_move_r2, active_board_transition.scripted_move_c2,
                active_board_transition.scripted_promotion);
        }
        
        game_state = PUZZLE_TRANSITION;
        
        dictionary to_captured_by_white, to_captured_by_black;
        calculate_captured_pieces_from_sbs(sbs_to, to_captured_by_white, to_captured_by_black);
        
        dictionary from_lists, to_lists;
        get_piece_locations(board, from_lists, null);
        get_piece_locations(sbs_to.board, to_lists, null);

        array<string> piece_keys = from_lists.getKeys();
        array<string> to_keys = to_lists.getKeys();
        for (uint i = 0; i < to_keys.length(); i++) {
            if(piece_keys.find(to_keys[i]) < 0) piece_keys.insertLast(to_keys[i]);
        }

        for (uint i = 0; i < piece_keys.length(); i++) {
            string key = piece_keys[i];
            int piece = parseInt(key);

            array<PieceLocation@>@ from_board = cast<array<PieceLocation@>>(from_lists[key]);
            array<PieceLocation@>@ to_board = cast<array<PieceLocation@>>(to_lists[key]);
            if(from_board is null) { @from_board = array<PieceLocation@>(); }
            if(to_board is null) { @to_board = array<PieceLocation@>(); }
            
            for (uint f = 0; f < from_board.length(); f++) {
                PieceLocation@ from_loc = from_board[f];
                int best_t = -1;
                float best_dist_sq = 99999;

                for (uint t = 0; t < to_board.length(); t++) {
                    PieceLocation@ to_loc = to_board[t];
                    if (to_loc.handled) continue;

                    float dist_sq = dist_sqr(from_loc.c, from_loc.r, to_loc.c, to_loc.r);
                    if (dist_sq < best_dist_sq) {
                        best_dist_sq = dist_sq;
                        best_t = t;
                    }
                }

                if (best_t != -1) {
                    PieceLocation@ to_loc = to_board[best_t];
                    float from_x, from_y, to_x, to_y;
                    get_board_or_capture_xy(piece, from_loc.r, from_loc.c, 0, from_x, from_y);
                    get_board_or_capture_xy(piece, to_loc.r, to_loc.c, 0, to_x, to_y);
                    active_board_transition.piece_anims.insertLast(PieceAnimation(piece, 0, 0, from_x, from_y, to_x, to_y, active_board_transition.duration));
                    from_loc.handled = true;
                    to_loc.handled = true;
                }
            }
            
            int to_capture_count = 0;
            dictionary@ target_capture_dict = (piece & COLOR_MASK) == WHITE ? @to_captured_by_black : @to_captured_by_white;
            target_capture_dict.get(''+(piece & PIECE_MASK), to_capture_count);
            int new_capture_offset = 0;
            
            for (uint f = 0; f < from_board.length(); f++) {
                PieceLocation@ from_loc = from_board[f];
                if (from_loc.handled) continue;
                
                float from_x, from_y, to_x, to_y;
                get_board_or_capture_xy(piece, from_loc.r, from_loc.c, 0, from_x, from_y);
                get_board_or_capture_xy(piece, -1, -1, to_capture_count + new_capture_offset++, to_x, to_y);
                active_board_transition.piece_anims.insertLast(PieceAnimation(piece, 0, 0, from_x, from_y, to_x, to_y, 1.0, 0.5, active_board_transition.duration));
                from_loc.handled = true;
            }

            int from_capture_count = 0;
            dictionary@ source_capture_dict = (piece & COLOR_MASK) == WHITE ? @captured_by_black : @captured_by_white;
            source_capture_dict.get(''+(piece & PIECE_MASK), from_capture_count);
            int from_capture_offset = 0;

            for (uint t = 0; t < to_board.length(); t++) {
                PieceLocation@ to_loc = to_board[t];
                if (to_loc.handled) continue;
                
                float from_x, from_y, to_x, to_y;
                get_board_or_capture_xy(piece, -1, -1, from_capture_count + from_capture_offset++, from_x, from_y);
                get_board_or_capture_xy(piece, to_loc.r, to_loc.c, 0, to_x, to_y);
                active_board_transition.piece_anims.insertLast(PieceAnimation(piece, 0, 0, from_x, from_y, to_x, to_y, 0.5, 1.0, active_board_transition.duration));
                to_loc.handled = true;
            }
        }

        board.resize(0);
        captured_by_white.deleteAll();
        captured_by_black.deleteAll();
    }

    void calculate_captured_pieces_from_sbs(SimplifiedBoardState@ sbs, dictionary &out captured_by_white, dictionary &out captured_by_black) {
        captured_by_white.deleteAll();
        captured_by_black.deleteAll();
        array<int> piece_types = {PAWN, KNIGHT, BISHOP, ROOK, QUEEN};
        for(uint i = 0; i < piece_types.length(); i++) {
            int type = piece_types[i];
            int start_count = (type == PAWN ? 8 : (type == QUEEN ? 1 : 2));
            
            int w_on_board = 0, b_on_board = 0;
            for(int r=0; r<8; r++) for(int c=0; c<8; c++) {
                if ((sbs.board[r][c] & PIECE_MASK) == type) {
                    if ((sbs.board[r][c] & COLOR_MASK) == WHITE) w_on_board++;
                    else b_on_board++;
                }
            }
            
            int w_captured = start_count - w_on_board;
            if (w_captured > 0) captured_by_black.set(''+type, w_captured);
            
            int b_captured = start_count - b_on_board;
            if (b_captured > 0) captured_by_white.set(''+type, b_captured);
        }
    }

    void get_board_or_capture_xy(int piece, int r, int c, int cap_idx, float &out x, float &out y) {
        float lft = self.x() - board_size_px / 2.0;
        float top = self.y() - board_size_px / 2.0 - board_y_offset;

        if (r != -1) {
            x = lft + c * tile_size + tile_size / 2;
            y = top + r * tile_size + tile_size / 2;
        } else {
            int capturer_color = (piece & COLOR_MASK) == WHITE ? BLACK : WHITE;
            int piece_type = piece & PIECE_MASK;
            
            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++) {
                if (PIECE_ORDER[i] == piece_type) {
                    y = start_y + current_y_offset * y_dir;
                    x = captured_start_x + cap_idx * horizontal_spacing;
                    return;
                }
                current_y_offset += vertical_spacing;
            }
        }
    }

    void get_piece_locations(array<array<int>> &in current_board, dictionary &out piece_locs, dictionary@ ignored) {
        for (int r = 0; r < 8; r++) {
            for (int c = 0; c < 8; c++) {
                int p = current_board[r][c];
                if (p == EMPTY) continue;
                string key = '' + p;

                if (!piece_locs.exists(key)) @piece_locs[key] = array<PieceLocation@>();
                array<PieceLocation@>@ locs = cast<array<PieceLocation@>>(piece_locs[key]);
                PieceLocation@ new_loc = PieceLocation();
                new_loc.r = r; new_loc.c = c;
                locs.insertLast(new_loc);
            }
        }
    }

	SimplifiedBoardState@ parse_fen(string fen) {
			array<string>@ parts = fen.split(" ");
			if (parts.length() < 2) return null;

			array<array<int>> fen_board;
			fen_board.resize(8);
			for(int i=0; i<8; i++) fen_board[i].resize(8);

			int r = 0, c = 0;
			for (uint i = 0; i < parts[0].length(); i++) {
				string char_ = parts[0].substr(i, 1);
				if (char_ == "/") { r++; c = 0; continue; }
				int digit = parseInt(char_);
				if (digit > 0 && digit < 9) { c += digit; continue; }

				int piece = EMPTY, color = WHITE;
				string lchar = string::lowercase(char_);
				if (lchar == "p") piece = PAWN; else if (lchar == "n") piece = KNIGHT; else if (lchar == "b") piece = BISHOP;
				else if (lchar == "r") piece = ROOK; else if (lchar == "q") piece = QUEEN; else if (lchar == "k") piece = KING;

				if (char_ == lchar) color = BLACK; else color = WHITE;
				
				fen_board[r][c] = piece | color;
				c++;
			}
			
			int kwr = -1, kwc = -1, kbr = -1, kbc = -1;
			for(r=0; r<8; r++) for(c=0; c<8; c++) {
				if(fen_board[r][c] == (KING | WHITE)) { kwr=r; kwc=c; }
				if(fen_board[r][c] == (KING | BLACK)) { kbr=r; kbc=c; }
			}

			bool w_k = false, w_q = false, b_k = false, b_q = false;
			// The 3rd part of the FEN string contains castling rights.
			if (parts.length() >= 3) {
				string castling = parts[2];
				if (castling.findFirst("K") >= 0) w_k = true;
				if (castling.findFirst("Q") >= 0) w_q = true;
				if (castling.findFirst("k") >= 0) b_k = true;
				if (castling.findFirst("q") >= 0) b_q = true;
			}

			// The 4th part of the FEN string contains the en passant square.
			int ep_r = -1, ep_c = -1;
			if(parts.length() >= 4 && parts[3] != "-") {
				string ep_square = parts[3];
				ep_c = int(ep_square[0]) - 97;
				ep_r = 56 - int(ep_square[1]);
			}

			return SimplifiedBoardState(fen_board, parts[1] == "w" ? WHITE : BLACK,
				w_k, w_q, b_k, b_q, ep_r, ep_c, kwr, kwc, kbr, kbc);
		}
    
    void parse_move_string(string move_str, int &out r1, int &out c1, int &out r2, int &out c2, int &out promotion_piece) {
        if (move_str.length() < 4) return;
        c1 = int(move_str[0]) - 97;
        r1 = 56 - int(move_str[1]);
        c2 = int(move_str[2]) - 97;
        r2 = 56 - int(move_str[3]);
        
        promotion_piece = EMPTY;
        if (move_str.length() >= 5) {
            string promo_char = move_str.substr(4, 1);
            if (promo_char == "q") promotion_piece = QUEEN;
            else if (promo_char == "r") promotion_piece = ROOK;
            else if (promo_char == "b") promotion_piece = BISHOP;
            else if (promo_char == "n") promotion_piece = KNIGHT;
        }
    }
    
    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() {
        if (is_puzzle_mode) return;
        
        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) {
                game_state = CHECKMATE_WHITE;
                int current_breaks = g.combo_break_count();
                int new_breaks;
                if (current_breaks == 0) new_breaks = 1; else if (current_breaks == 1) new_breaks = 2; else if (current_breaks <= 3) new_breaks = 4; else if (current_breaks <= 5) new_breaks = 6; else new_breaks = current_breaks + 1;
                g.combo_break_count(new_breaks);
                g.end_level(self.x(), self.y());
            } else {
                game_state = CHECKMATE_BLACK;
                g.end_level(self.x(), self.y());
            }
        } else { 
            game_state = STALEMATE; 
			int current_breaks = g.combo_break_count();
			int new_breaks;
			if (current_breaks == 0) new_breaks = 1; else if (current_breaks == 1) new_breaks = 2; else if (current_breaks <= 3) new_breaks = 4; else if (current_breaks <= 5) new_breaks = 6; else new_breaks = current_breaks + 1;
			g.combo_break_count(new_breaks);
			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_board_arrow(int r1, int c1, int r2, int c2, uint colour) {
		float lft = self.x() - board_size_px/2.0;
		float top = self.y() - board_size_px/2.0 - board_y_offset;

		float x1 = lft + c1 * tile_size + tile_size / 2.0;
		float y1 = top + r1 * tile_size + tile_size / 2.0;
		float x2 = lft + c2 * tile_size + tile_size / 2.0;
		float y2 = top + r2 * tile_size + tile_size / 2.0;
		
		draw_arrow(g, self.layer(), 12, x1, y1, x2, y2, 6, 20, 1.0, colour);
	}

    void draw(float subframe) {
        float lft = self.x()-board_size_px/2.0, top = self.y()-board_size_px/2.0 - board_y_offset;
        
        // --- Board Shake Animation ---
        if (shake_timer > 0 && shake_duration_initial > 0) {
            float current_intensity = shake_intensity_initial * ease_out_cubic(shake_timer / shake_duration_initial);
            lft += rand_range(-current_intensity, current_intensity);
            top += rand_range(-current_intensity, current_intensity);
        }
        
        g.draw_rectangle_world(self.layer(),0,lft,top,lft + board_size_px, top + board_size_px,0,0xFF000000); // BG
        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 (@active_board_transition != null) {
            float anim_progress = active_board_transition.progress - active_board_transition.delay;
            array<PieceAnimation@>@ anims = active_board_transition.piece_anims;
            if (anim_progress > 0) {
                float t = ease_in_out_cubic(clamp01(anim_progress / active_board_transition.duration));
                for (uint i = 0; i < anims.length(); i++) {
                    PieceAnimation@ pa = anims[i];
                    float anim_x = lerp(pa.from_x, pa.to_x, t);
                    float anim_y = lerp(pa.from_y, pa.to_y, t);
                    float anim_scale = lerp(pa.from_scale, pa.to_scale, t);
                    draw_piece(pa.piece, anim_x, anim_y, anim_scale);
                }
            } else {
                 for (uint i = 0; i < anims.length(); i++) {
                    PieceAnimation@ pa = anims[i];
					draw_piece(pa.piece, pa.from_x, pa.from_y, pa.from_scale);
                }
            }
        } else {
            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);
			
			// Draw custom highlights
			for(uint i = 0; i < highlighted_squares.length(); i++) {
				int key = highlighted_squares[i];
				int r = key / 8;
				int c = key % 8;
				g.draw_rectangle_world(self.layer(), 1, lft + c * tile_size, top + r * tile_size, lft + (c + 1) * tile_size, top + (r + 1) * tile_size, 0, 0x88FF0000);
			}

            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, anim_y;
                if (game_state == INCORRECT_MOVE_ANIMATION && active_animation.is_reversing) {
                     anim_x = lerp(active_animation.to_x, active_animation.from_x, 1-t);
                     anim_y = lerp(active_animation.to_y, active_animation.from_y, 1-t);
                } else {
                    anim_x = lerp(active_animation.from_x, active_animation.to_x, t);
                    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 custom circles over pieces
			for(uint i = 0; i < circled_squares.length(); i++) {
				int key = circled_squares[i];
				int r = key / 8;
				int c = key % 8;
				float circle_x = lft + c * tile_size + tile_size / 2;
				float circle_y = top + r * tile_size + tile_size / 2;
				drawing::circle(g, self.layer(), 11, circle_x, circle_y, tile_size * 0.4, 32, 4, 0x88FF0000);
			}

			// Draw finalized arrows
			for (uint i = 0; i < drawn_arrows.length(); i++) {
				Arrow@ arrow = drawn_arrows[i];
				draw_board_arrow(arrow.r1, arrow.c1, arrow.r2, arrow.c2, 0xAA228B22); // ForestGreen, semi-transparent
			}

			// Draw the live arrow
			if (is_dragging) {
				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());
					int r_end = int(floor((my - top) / tile_size));
					int c_end = int(floor((mx - lft) / tile_size));

					if (is_on_board(r_end, c_end) && (r_end != drag_start_r || c_end != drag_start_c)) {
						draw_board_arrow(drag_start_r, drag_start_c, r_end, c_end, 0xAA228B22);
					}
				}
			}
            
            draw_captured_pieces(WHITE, lft, top);
            draw_captured_pieces(BLACK, lft, top);
        }
        
        if (game_state == AWAITING_PROMOTION) draw_promotion_options();

        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());
                g.draw_rectangle_world(self.layer(), 10, mx - 5, my - 5, mx + 5, my + 5, 0, 0x33FF0000);
            }
        }

        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; case PUZZLE_TRANSITION: status="Setting up puzzle..."; break; case INCORRECT_MOVE_ANIMATION: status="Incorrect Move!"; 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);
    }
    
	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;
    [text] int puzzles_to_win = 5;

    // --- Pause Detection ---
    scene@ g;

    KnightChess@ game_instance;
	bool level_ended = false;
    bool is_paused = false;
    bool unpause_shake_pending = false;
    int has_stepped = 2;
    int frames = 0;

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

    void on_level_start() {
        frames = 0;
        is_paused = false;
        unpause_shake_pending = false;
        has_stepped = 2;
        @game_instance = null;
    }

	void on_level_end() {
		level_ended = true;
	}

    void step_fixed() {
        // This function runs at a fixed rate, even when paused.
        if (is_paused) return; // If already paused, do nothing more here.
		
        if (frames <= 54 || level_ended) {
            // Don't do anything until the game is running and past the initial countdown.
            return;
        }

        if (has_stepped <= 0) {
            // If step() hasn't run, it means the game is paused. Set the flags.
            is_paused = true;
            unpause_shake_pending = true;
        }
        has_stepped--;
    }

    void step(int) {
        // This function only runs when the game is NOT paused.
        has_stepped = 2; // Signal that step() has run this frame.
        
        // If the `is_paused` flag is true, it means we have just un-paused.
        if (is_paused) {
            // Apply the combo break penalty.
            int current_breaks = g.combo_break_count();
            int new_breaks;
            if (current_breaks == 0) new_breaks = 1;
            else if (current_breaks == 1) new_breaks = 2;
            else if (current_breaks <= 3) new_breaks = 4;
            else if (current_breaks <= 5) new_breaks = 6;
            else new_breaks = current_breaks + 1;
            g.combo_break_count(new_breaks);

            // Reset the pause flag.
            is_paused = false;
			puts("Game unpaused, combo broken.");
        }
        
        // Trigger shake if pending from an unpause event
        if (unpause_shake_pending) {
            // Re-acquire the game instance from the player controller
            // as the cached @game_instance handle can become null.
            controllable@ p_controllable = controller_controllable(0);
            if (@p_controllable != null) {
                scriptenemy@ game_entity = p_controllable.as_scriptenemy();
                if (@game_entity != null) {
                    KnightChess@ chess_instance = cast<KnightChess@>(game_entity.get_object());
                    if (@chess_instance != null) {
                        chess_instance.start_shake(UNPAUSE_SHAKE_DURATION, UNPAUSE_SHAKE_INTENSITY);
                    }
                }
            }
            unpause_shake_pending = false;
        }

        rrnd.step();
        frames++;
    }
    
    void spawn_player(message@ msg) {
        string player_char = msg.get_string("character", "dustman");
        // Store the KnightChess instance so we can access its state.
        @game_instance = KnightChess(@rrnd, player_char, board_scale, puzzles_to_win);
        scriptenemy@ game_entity = create_scriptenemy(game_instance);
        game_entity.x(msg.get_float("x"));
        game_entity.y(msg.get_float("y"));
        msg.set_entity("player", @game_entity.as_entity());
    }
}