defmodule Elswisser.Square.Sees do
  defstruct up: [],
            down: [],
            left: [],
            right: [],
            up_right: [],
            up_left: [],
            down_left: [],
            down_right: [],
            knight: [],
            all: MapSet.new()
end

defmodule Elchesser.Square do
  alias Elchesser.Square.Sees

  defstruct file: nil,
            rank: nil,
            loc: {},
            piece: nil,
            sees: %Sees{}
end

defmodule Elchesser.Move do
  defstruct from: {},
            to: {},
            piece: nil,
            capture: nil,
            promotion: nil,
            castle: false,
            checking: nil
end

defmodule Elchesser.Piece do
  alias Elchesser.{Square, Game, Move}

  @doc """
  Generates psuedo-legal moves for a given piece type.
  """
  @callback moves(Square.t(), Game.t()) :: [Move.t()]
  @callback attacks(Square.t(), Game.t()) :: [Move.t()]
end

  @spec move_range(%Square{}, %Game{}, Square.Sees.t()) :: [Move.t()]
  def move_range(%Square{} = square, %Game{} = game, direction) do
    get_in(square.sees, [Access.key!(direction)])
    |> Enum.reduce_while([], fn {file, rank}, acc ->
      s = Game.get_square(game, {file, rank})

      case friendly?(square.piece, s.piece) do
        true -> {:halt, acc}
        false -> {:halt, [Move.from(square, {s.file, s.rank}, capture: true) | acc]}
        nil -> {:cont, [Move.from(square, {s.file, s.rank}) | acc]}
      end
    end)
    |> Enum.reverse()
  end

defmodule Elchesser.Piece.Bishop do
  alias Elchesser.{Square, Game, Piece}

  @behaviour Piece

  @impl true
  def moves(%Square{} = square, %Game{} = game) do
    for d <- [:up_right, :up_left, :down_left, :down_right], reduce: [] do
      acc -> [Square.move_range(square, game, d) | acc]
    end
    |> List.flatten()
  end

  @impl true
  def attacks(square, game), do: moves(square, game)
end

defmodule Elchesser.Piece.Knight do
  alias Elchesser.{Square, Game, Piece, Move}

  @behaviour Piece

  @impl true
  def moves(%Square{} = square, %Game{} = game) do
    Enum.reduce(square.sees.knight, [], fn {file, rank}, acc ->
      s = Map.get(game.board, {file, rank})

      case Piece.friendly?(square.piece, s.piece) do
        true -> acc
        false -> [Move.from(square, {s.file, s.rank}, capture: true) | acc]
        nil -> [Move.from(square, {s.file, s.rank}) | acc]
      end
    end)
  end

  @impl true
  def attacks(square, game), do: moves(square, game)
end

defmodule Elchesser.Piece.Pawn do
  alias Elchesser.{Square, Game, Move, Piece}

  @behaviour Piece
  def moves(_, _), do: []

  @impl true
  def attacks(%Square{} = square, _), do: attacks(square) |> Enum.map(&Move.from(square, &1))

  defp attacks(%Square{piece: :P, rank: rank, file: file}) do
    [{file + 1, rank + 1}, {file - 1, rank + 1}]
    |> Enum.filter(&Square.valid?/1)
  end

  defp attacks(%Square{piece: :p, file: file, rank: rank}) do
    [{file + 1, rank - 1}, {file - 1, rank - 1}]
    |> Enum.filter(&Square.valid?/1)
  end
end

defmodule Elchesser.Piece.Pawn do
  # ... 
  @impl true
  def moves(%Square{piece: :P, rank: 2, file: file} = square, %Game{} = game),
    do: moves(square, game, [{file, 3}, {file, 4}], false)

  def moves(%Square{piece: :P, rank: rank, file: file} = square, %Game{} = game),
    do: moves(square, game, [{file, rank + 1}], rank + 1 == 8)

  def moves(%Square{piece: :p, rank: 7, file: file} = square, %Game{} = game),
    do: moves(square, game, [{file, 6}, {file, 5}], false)

  def moves(%Square{piece: :p, rank: rank, file: file} = square, %Game{} = game),
    do: moves(square, game, [{file, rank - 1}], rank - 1 == 1)
  # ... 
end

defmodule Elchesser.Piece.Pawn do
  # ...
  defp moves(square, game, candidates, promotion) do
    m =
      candidates
      # moves are only valid if the candidate square is empty
      |> Enum.filter(fn candidate -> Game.get_square(game, candidate) |> Square.empty?() end)
      |> Enum.map(fn {file, rank} ->
        Move.from(square, {file, rank}, promotion: promote(promotion, square.piece))
      end)

    a =
      attacks(square)
      # attacks are only valid if the candidate square is occupied by an enemy
      # piece or the current en passant square of the given game state
      |> Enum.filter(fn s ->
        sq = Game.get_square(game, s)
        Piece.enemy?(square.piece, sq.piece) || en_passant?(sq, game)
      end)
      |> Enum.map(fn s ->
        sq = Game.get_square(game, s)
        ep? = en_passant?(sq, game)

        capture =
          cond do
            ep? and square.piece == :P -> :p
            ep? and square.piece == :p -> :P
            true -> sq.piece
          end

        Move.from(square, s, capture: capture, promotion: promote(promotion, square.piece))
      end)

    Enum.concat(m, a)
  end
  # ...

  defp en_passant?(%Square{} = square, %Game{} = game) do
    Square.empty?(square) && Square.eq?(square, game.en_passant)
  end
end

defmodule Elchesser.Piece.King do
  # ...
  def attacks(%Square{file: file, rank: rank}) do
    for(f <- -1..1, r <- -1..1, not (f == 0 and r == 0), do: {file + f, rank + r})
    |> Enum.filter(&Square.valid?/1)
    |> Enum.reduce([], fn {file, rank}, acc ->
      s = Map.get(game.board, {file, rank})

      case Piece.friendly?(square.piece, s.piece) do
        true -> acc
        false -> [Move.from(square, {s.file, s.rank}, capture: true) | acc]
        nil -> [Move.from(square, {s.file, s.rank}) | acc]
      end
    end)
  end
  # ...
end 

defmodule Elchesser.Board do
  def white_occupied(%Game{board: board}) do
    Map.filter(board, fn {_, square} -> Square.white?(square) end)
    |> Map.values()
  end

  def black_occupied(%Game{board: board}) do
    Map.filter(board, fn {_, square} -> Square.black?(square) end)
    |> Map.values()
  end
end

defmodule Elchesser.Board do
  def white_attacks(%Game{} = game) do
    white_occupied(game)
    |> Enum.map(&Square.attacks(&1, game))
    |> List.flatten()
    |> Enum.uniq()
    |> Enum.map(& &1.to)
  end

  def black_attacks(%Game{} = game) do
    black_occupied(game)
    |> Enum.map(&Square.attacks(&1, game))
    |> List.flatten()
    |> Enum.uniq()
    |> Enum.map(& &1.to)
  end
end

defmodule Elchesser.Piece.King do
  # ...
  # handle the case where castling is illegal by the game rules
  defp maybe_castle_kingside(%Square{piece: piece}, %Game{castling: c}, _)
       when not is_map_key(c.map, piece),
       do: []

  # check that castling is legal for the given position
  defp maybe_castle_kingside(%Square{piece: :K} = square, %Game{} = game, attacks) do
    through_squares = [{?f, 1}, {?g, 1}] |> Enum.map(&Game.get_square(game, &1))

    if can_castle?(through_squares, attacks),
      do: [Move.from(square, {?g, 1}, castle: true)],
      else: []
  end
  # ...
  
  # Going kingside has the same through and attacked squares, but when castling
  # queenside, the b-file can be attacked, so we have the two function heads.
  defp can_castle?(through_squares, attacks),
    do: can_castle?(through_squares, through_squares, attacks)

  defp can_castle?(empty_through, attack_through, attacks) do
    empty? = Enum.all?(empty_through, &Square.empty?/1)

    attacked? =
      attacks |> Enum.map(&Square.from/1) |> Enum.any?(&(&1 in attack_through))

    empty? and not attacked?
  end
end

defmodule Elchesser.Piece.King do
  # ...
  # Note this is the private method shared between the black and white kings.
  # The difference would be which set of attacks are fed into the function
  defp moves(%Square{} = square, %Game{} = game, %MapSet{} = attacks) do
    # starting with all the attacked squares
    attacks(square, game) 
    # add kingside castling if it's legal
    |> Enum.concat(maybe_castle_kingside(square, game, attacks)) 
    # add queenside castling if it's legal
    |> Enum.concat(maybe_castle_queenside(square, game, attacks))
    # filter out all moves where the given square is attacked by an enemy piece
    |> Enum.reject(&MapSet.member?(attacks, &1.to))
  end
  # ...
end 

defmodule Elchesser.Game
  # ...
  defp ensure_valid_move(%Game{} = game, %Move{} = move) do
    with from <- Game.get_square(game, move.from),
         to <- Game.get_square(game, move.to),
         :ok <- not Square.empty?(from) |> or_(:empty_square),
         :ok <- (game.active == Board.color_at(game, from)) |> or_(:invalid_from_color),
         :ok <- (game.active != Board.color_at(game, to)) |> or_(:invalid_to_color),
         :ok <- (move.to in Square.legal_locs(from, game)) |> or_(:invalid_move) do
      :ok
    end
  end
  # ...
  @spec or_(boolean(), atom()) :: :ok | {:error, atom()}
  defp or_(true, _), do: :ok
  defp or_(false, reason), do: {:error, reason}
  # ...
end

defmodule Elchesser.Board do
  # ...
  @spec move(Elchesser.Game.t(), Elchesser.Move.t()) ::
          {:ok, {Move.t(), Game.t()}} | {:error, atom()}
  def move(%Game{} = game, %Move{} = move) do
    with {:ok, {piece, game}} <- move_from(game, move.from),
         {:ok, {capture, game}} <- move_to(game, move.to, piece),
         {:ok, game} <- castle(game, move),
         {:ok, game} <- promote(game, move) do
      check = Game.Check.opponent_in_check?(game)

      move = %{
        move
        | checking: if(check == true, do: :check, else: nil),
          capture: capture,
          piece: piece
      }

      move = %{move | san: Move.san(move)}

      {:ok, {move, game}}
    end
  end
  # ...
end

def move_from(%Game{board: board} = game, loc) do
  case Map.get_and_update(board, loc, fn current ->
          {current, %Square{current | piece: nil}}
        end) do
    {%Square{piece: nil}, _} -> {:error, :empty_square}
    {%Square{piece: p}, board} -> {:ok, {p, %Game{game | board: board}}}
  end
end

def move_to(%Game{board: board} = game, loc, piece) do
  {%Square{piece: capture}, board} =
    Map.get_and_update(board, loc, fn current ->
      {current, %Square{current | piece: piece}}
    end)

  cond do
    Piece.friendly?(piece, capture) -> {:error, :invalid_to_color}
    true -> {:ok, {capture, %Game{game | board: board}}}
  end
end

def castle(game, %Move{castle: false}), do: {:ok, game}

def castle(%Game{} = game, %Move{to: to}) do
  move =
    case to do
      {?g, 1} -> Move.from({?h, 1, :R}, {?f, 1})
      {?g, 8} -> Move.from({?h, 8, :r}, {?f, 8})
      {?c, 1} -> Move.from({?a, 1, :R}, {?d, 1})
      {?c, 8} -> Move.from({?a, 8, :r}, {?d, 8})
    end

  with {:ok, {_, game}} <- move(game, move) do
    {:ok, game}
  end
end

def promote(game, %Move{promotion: nil}), do: {:ok, game}

def promote(%Game{board: board} = game, %Move{to: to, promotion: promotion}) do
  with {_, board} <-
          Map.get_and_update(board, to, fn current ->
            {current, %Square{current | piece: promotion}}
          end) do
    {:ok, %Game{game | board: board}}
  end
end

defmodule Elchesser.Game do
  # ...
  def move(%Game{} = game, %Move{} = move) do
    with :ok <- ensure_valid_move(game, move),
         {:ok, {move, game}} <- Board.move(game, move) do
      game =
        game
        |> flip_color()
        |> set_in_check()
        |> add_move(move)
        |> add_fen()
        |> add_capture(move.capture)
        |> set_castling_rights(move)
        |> set_en_passant(move)
        |> set_half_move_count(move)
        |> set_full_move_count()

      {:ok, game}
    end
  # ...
end

game = Game.new()

{:ok, game} = Game.move(game, Move.from({?e, 2, :P}, {?e, 4}))
{:ok, game} = Game.move(game, Move.from({?e, 7, :p}, {?e, 5}))
{:ok, game} = Game.move(game, Move.from({?d, 2, :P}, {?d, 4}))
{:ok, game} = Game.move(game, Move.from({?e, 5, :p}, {?d, 4}))