require "./eval"
require "./moves"

module Fork
  enum GameState
    Menu
    PlayerTurnStart
    PickingPiece
    ChoosingMove
    PlayerMoveReady
    PromotionDialog
    EnemyTurn
    WaitingForEnemy
    GameEnded
  end

  alias AnimTexture = Array(Raylib::Texture2D)

  struct Textures
    getter pieces : Array(AnimTexture)
    getter square : AnimTexture
    getter win : AnimTexture
    getter defeat : AnimTexture
    getter draw : AnimTexture
    getter new : AnimTexture
    getter game : AnimTexture

    def initialize
      pawn = load_animation_textures("res/pawn", 4)
      knight = load_animation_textures("res/knight", 4)
      bishop = load_animation_textures("res/bishop", 4)
      rook = load_animation_textures("res/rook", 4)
      queen = load_animation_textures("res/queen", 4)
      king = load_animation_textures("res/king", 4)

      @pieces = [pawn, knight, bishop, rook, queen, king]
      @square = load_animation_textures("res/square", 4)
      @win = load_animation_textures("res/win", 4)
      @defeat = load_animation_textures("res/defeat", 4)
      @draw = load_animation_textures("res/draw", 4)
      @new = load_animation_textures("res/new", 4)
      @game = load_animation_textures("res/game", 4)
    end

    def load_animation_textures(path, num)
      arr = [] of Raylib::Texture2D
      1.upto(num) do |i|
        tex = Raylib.load_texture("#{path}#{i}.png")
        arr << tex
      end
      arr
    end
  end

  ANIM_OFFSET = Array(Int32).new(64) { Random.rand(4) }

  class UI
    MARGIN          =  2
    SQUARE_OVERLAP  =  2
    WIDTH           = 32
    PROMO_TYPES     = [QUEEN, KNIGHT, ROOK, BISHOP]
    LAST_MOVE_COLOR = Raylib::Color.new r: 255, g: 203, b: 0, a: 127
    HOVERED_COLOR   = Raylib::Color.new r: 102, g: 191, b: 255, a: 127

    @state : GameState = GameState::Menu
    @starting_board : Board
    @board : Board = EMPTY_BOARD
    @player_side : Int32 = WHITE
    @piece_picked : Fork::Piece | Nil = nil
    @last_move : Fork::Move | Nil = nil
    @player_move : Fork::Move | Nil = nil
    @pawn_promotion_move : Fork::Move | Nil = nil
    @legal_moves : Array(Fork::Move) | Nil = nil
    @anim_tick : Int32 = 0
    @winner : Int32 | Nil = nil
    @textures : Textures

    def initialize(@starting_board : Board = STARTING_BOARD)
      Raylib.set_trace_log_level(Raylib::TraceLogLevel::Warning)
      Raylib.init_window(MARGIN * 2 + WIDTH * 8, MARGIN * 2 + WIDTH * 8, "fork")
      Raylib.set_target_fps(30)

      @textures = Textures.new
    end

    def draw_anim_texture(textures, square, color)
      x, y = get_coords(square)
      frame = (@anim_tick + ANIM_OFFSET[square]) % textures.size
      Raylib.draw_texture(textures[frame], x, y, color)
    end

    def play
      while continue?
        @anim_tick = (Raylib.get_time / 0.25).to_i
        square = check_for_input
        Raylib.begin_drawing
        show_board

        case @state
        when GameState::Menu
          menu(square)
        when GameState::PlayerTurnStart
          player_turn_start
        when GameState::PickingPiece
          picking_piece(square) if square
        when GameState::ChoosingMove
          choosing_move(square) if square
        when GameState::PromotionDialog
          show_promo_dialog
          promotion_dialog(square) if square
        when GameState::PlayerMoveReady
          player_move_ready
        when GameState::EnemyTurn
          enemy_turn
        when GameState::GameEnded
          game_ended(square)
        else
        end

        Raylib.end_drawing
      end
    end

    def to_menu
      @state = GameState::Menu
      @board = EMPTY_BOARD
      reset_game
    end

    def menu(square)
      draw_anim_texture(@textures.new, D5, Raylib::WHITE)
      draw_anim_texture(@textures.game, E5, Raylib::WHITE)

      draw_anim_texture(@textures.pieces[PAWN], C3, Raylib::WHITE)
      draw_anim_texture(@textures.pieces[PAWN], F3, Raylib::BLACK)

      if square
        if square == C3
          starting_game(WHITE)
        elsif square == F3
          starting_game(BLACK)
        end
      end
    end

    def starting_game(side)
      @board = @starting_board
      @player_side = side
      if @board.side == @player_side
        @state = GameState::PlayerTurnStart
      else
        @state = GameState::EnemyTurn
      end
    end

    def player_turn_start
      moves = Moves.generate(@board).select do |m|
        copy = @board.dup
        legal = copy.make_move(m)
        legal
      end

      if moves.empty?
        to_game_ended
      else
        @legal_moves = moves
        to_picking_piece
      end
    end

    def to_picking_piece
      @piece_picked = nil
      @state = GameState::PickingPiece
    end

    def picking_piece(square)
      piece = @board.piece_at(square)
      if piece.nil? || piece.side != @player_side
        to_picking_piece
      else
        to_choosing_move(piece)
      end
    end

    def to_choosing_move(piece)
      @piece_picked = piece
      @state = GameState::ChoosingMove
    end

    def choosing_move(square)
      return to_picking_piece if square.nil?
      row, col = square.divmod(8)
      piece = @board.piece_at(square)
      if piece.try &.side == @player_side
        return to_picking_piece
      end

      @piece_picked.try do |p|
        if p.type == PAWN && (row == 0 || row == 7)
          return to_promotion_dialog(Fork::Move.new(p.index, square, p.type, 0))
        else
          castle = castle_flag(p, square)
          move = Fork::Move.new(p.index, square, p.type, 0, castle: castle)
          return to_player_move_ready(move)
        end
      end
    end

    def to_promotion_dialog(promo_move)
      @pawn_promotion_move = promo_move
      @state = GameState::PromotionDialog
    end

    def promotion_dialog(square)
      return to_picking_piece if square.nil?
      if move = @pawn_promotion_move
        row, col = square.divmod(8)
        promo_row, promo_col = move.to.divmod(8)
        pick = @player_side == WHITE ? 7 - row : row
        if col == promo_col && pick < 4
          type = PROMO_TYPES[pick]
          move.set_promotion_piece(type)
          return to_player_move_ready(move)
        else
          return to_picking_piece
        end
      end
    end

    def to_player_move_ready(move)
      @player_move = move
      @state = GameState::PlayerMoveReady
    end

    def player_move_ready
      move = @player_move
      return to_picking_piece if move.nil?
      legal_move = @legal_moves.try &.find do |l|
        l.from == move.from && l.to == move.to && l.promote_to == move.promote_to
      end

      if legal_move.nil?
        to_picking_piece
      else
        copy = @board.dup
        legal = copy.make_move(legal_move)
        if legal
          @board = copy
          @last_move = legal_move
          @legal_moves = nil
          to_enemy_turn
        end
      end
    end

    def to_enemy_turn
      @state = GameState::EnemyTurn
    end

    def enemy_turn
      to_waiting_for_enemy
      Thread.new do
        moves = Search.new(@board, 6).run
        moves.each do |m|
          puts "#{Fork.square_to_text(m.from)}-#{Fork.square_to_text(m.to)}" unless m.nil?
        end
        puts "\n"
        if m = moves.first
          @board.make_move(m)
          @last_move = m
          to_player_turn_start
        else
          to_game_ended
        end
      end
    end

    def to_waiting_for_enemy
      @state = GameState::WaitingForEnemy
    end

    def to_player_turn_start
      @state = GameState::PlayerTurnStart
    end

    def to_game_ended
      @state = GameState::GameEnded
    end

    def game_ended(square)
      if @board.checked?(@board.side)
        @winner = Fork.enemy(@board.side)
      else
        @winner = DRAW
      end
      if square
        to_menu
      end
    end

    def castle_flag(piece, move_to)
      if piece.type == KING
        if @player_side == WHITE
          if (@board.castle & WHITE_KING_CASTLE) != 0 && move_to == G1
            return WHITE_KING_CASTLE
          elsif (@board.castle & WHITE_QUEEN_CASTLE) != 0 && move_to == C1
            return WHITE_QUEEN_CASTLE
          end
        else
          if (@board.castle & BLACK_KING_CASTLE) != 0 && move_to == G8
            return BLACK_KING_CASTLE
          elsif (@board.castle & BLACK_QUEEN_CASTLE) != 0 && move_to == C8
            return BLACK_QUEEN_CASTLE
          end
        end
      end
      0
    end

    def continue?
      !Raylib.close_window?
    end

    def get_coords(square)
      row, col = square.divmod(8)
      row = (row - 7).abs if @player_side == BLACK
      {col * WIDTH + MARGIN, (7 - row).abs * WIDTH + MARGIN}
    end

    def show_board
      Raylib.clear_background(Raylib::DARKGREEN)
      available_moves = nil
      if moves_from = @piece_picked.try &.index
        available_moves = @legal_moves.try &.select { |m| m.from == moves_from }.map(&.to)
      end

      A1.upto(H8) do |index|
        row, col = index.divmod(8)
        x, y = get_coords(index)
        x -= SQUARE_OVERLAP
        y -= SQUARE_OVERLAP

        texture = @textures.square[(row + col) % 4]
        board_color = (row + col) % 2 == 0 ? Raylib::DARKGREEN : Raylib::BEIGE
        Raylib.draw_texture(texture, x, y, board_color)

        @last_move.try do |m|
          if index == m.from || index == m.to
            Raylib.draw_texture(texture, x, y, LAST_MOVE_COLOR)
          end
        end

        if @piece_picked.try &.index == index
          Raylib.draw_texture(texture, x, y, HOVERED_COLOR)
        end

        if available_moves.try &.includes?(index)
          Raylib.draw_texture(texture, x, y, HOVERED_COLOR)
        end
      end

      A1.upto(H8) do |index|
        piece = @board.piece_at(index)
        unless piece.nil?
          color = piece.side == WHITE ? Raylib::WHITE : Raylib::BLACK
          if @state == GameState::GameEnded && piece.type == KING
            if @winner == DRAW
              draw_anim_texture(@textures.draw, index, color)
            else
              result_texture = piece.side == @winner ? @textures.win : @textures.defeat
              draw_anim_texture(result_texture, index, color)
            end
          else
            draw_anim_texture(@textures.pieces[piece.type], index, color)
          end
        end
      end
    end

    def show_promo_dialog
      if move = @pawn_promotion_move
        PROMO_TYPES.each_with_index do |type, i|
          square = move.to - (8 * i)
          x, y = get_coords(square)
          x -= SQUARE_OVERLAP
          y -= SQUARE_OVERLAP
          Raylib.draw_texture(@textures.square[i], x, y, Raylib::DARKPURPLE)
          color = @player_side == WHITE ? Raylib::WHITE : Raylib::BLACK
          draw_anim_texture(@textures.pieces[type], square, color)
        end
      end
    end

    def check_for_input
      return square_clicked if Raylib.mouse_button_pressed?(Raylib::MouseButton::Left)
    end

    def square_clicked
      pos = Raylib.get_mouse_position
      col = (pos.x - MARGIN) // WIDTH
      row = 7 - ((pos.y - MARGIN) // WIDTH)
      row = (7 - row).abs if @player_side == BLACK
      if col < 0 || col > 7 || row < 0 || row > 7
        nil
      else
        (row * 8 + col).to_i
      end
    end

    def reset_game
      @piece_picked = nil
      @last_move = nil
      @player_move = nil
      @pawn_promotion_move = nil
      @legal_moves = nil
      @winner = nil
      @player_side = WHITE
    end
  end
end