from typing import Union

class Piece:
  def __init__(self, char):
    self.type = char.lower()
    self.isWhite = char.isupper()
  
  def __str__(self):
    return self.type.upper() if self.isWhite else self.type

class Board:
  def __init__(self, boardString, size=5):
    pieces = [Piece(char) if char!='.' else None for char in boardString]
    self.rows = [pieces[i:i+size] for i in range(0, len(pieces), size)]
  
  def __str__(self):
    return ''.join([''.join([str(piece) if piece!=None else '.' for piece in row])for row in self.rows])

  def print(self):
    for row in self.rows:
      for piece in row:
        print(piece if piece!=None else '.', end='')
      print()

  def getPiece(self, x, y) -> Union[Piece, None]:
    x -= 1
    y = len(self.rows) - y
    return self.rows[y][x]
  
  def getLegalMoves(self, x, y):
    try:
      piece = self.getPiece(x, y)
      assert piece.type == 'p'

      frontDirection = 1 if piece.isWhite else -1

      def checkIfStartPosition(x,y):
        if piece.isWhite and y == 2 and self.getPiece(x,y+2) == None:
          return [(x,y+2)]
        elif (not piece.isWhite) and y == 4 and self.getPiece(x,y-2) == None:
          return [(x,y-2)]
        return []

      def checkInFrontOf(x, y, frontDirection):
        if not self.getPiece(x, y+frontDirection):
          return [(x, y+frontDirection)]
        return []
      
      def checkDiagonalOf(x, y, frontDirection):
        moves = []
        for xToCheck in [-1, 1]:
          pieceToCheck = self.getPiece(x+xToCheck, y+frontDirection)
          if pieceToCheck != None and pieceToCheck.isWhite != piece.isWhite:
            moves.append((x+xToCheck, y+frontDirection))
        return moves
      
      moves = checkInFrontOf(x,y,frontDirection)
      if moves != []:
        moves += checkIfStartPosition(x,y)
      moves += checkDiagonalOf(x,y,frontDirection)
      return moves
      
    except AssertionError:
      print('Piece rules not implemented yet')

if __name__ == "__main__":
  board = Board('rkn.r.p.....P..PP.PPB.K..')
  board.print()
  print()
  print(board.getPiece(5, 2))
  print(board.getPiece(2, 1))
  print()
  print(board.getLegalMoves(4, 2))
  print(board.getLegalMoves(2, 4))