# Copyright 2023 Mattia Giambirtone & All Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ../util/matrix

export matrix

import std/strutils
import std/strformat


type
    # Useful type aliases
    Location = tuple[row, col: int]

    Pieces = tuple[king: Location, queens: seq[Location], rooks: seq[Location],
                   bishops: seq[Location], knights: seq[Location],
                   pawns: seq[Location]]

    PieceColor* = enum
        ## A piece color enumeration
        None = 0,
        White, 
        Black

    PieceKind* = enum
        ## A chess piece enumeration
        Empty = '\0',    # No piece
        Bishop = 'b', 
        King = 'k'
        Knight = 'n',
        Pawn = 'p', 
        Queen = 'q',
        Rook = 'r', 

    Piece* = object
        ## A chess piece
        color*: PieceColor
        kind*: PieceKind

    MoveFlag* = enum
        ## An enumeration of move flags
        Default,   # Move is a regular move
        XRay,   # Move is an X-ray attack
        # Move is a pawn promotion
        PromoteToQueen,
        PromoteToRook,
        PromoteToBishop,
        PromoteToKnight       

    Move* = object
        ## A chess move
        piece*: Piece
        startSquare*: Location
        targetSquare*: Location
        flag*: MoveFlag

    ChessBoard* = ref object
        ## A chess board object
        grid: Matrix[Piece]
        # Currently active color
        turn: PieceColor
        # Number of half moves since
        # last piece capture or pawn movement.
        # Used for the 50-move rule
        halfMoveClock: int
        # Full move counter. Increments
        # every 2 ply
        fullMoveCount: int
        # Stores metadata for castling.
        castling: tuple[white, black: tuple[queen, king: bool]]
        # En passant target square (see https://en.wikipedia.org/wiki/En_passant)
        # If en passant is not possible, both the row and
        # column of the position will be set to -1
        enPassantSquare*: Move
        # Locations of all pieces
        pieces: tuple[white: Pieces, black: Pieces]
        # Potential attacking moves for black and white
        attacked: tuple[white: seq[Move], black: seq[Move]]


# Initialized only once, copied every time
var empty: seq[Piece] = @[]
for _ in countup(0, 63):
    empty.add(Piece(kind: Empty, color: None))


func emptyPiece*: Piece {.inline.} = Piece(kind: Empty, color: None)
func emptyLocation*: Location {.inline.} = (-1 , -1)
func opposite*(c: PieceColor): PieceColor {.inline.} = (if c == White: Black else: White)
proc algebraicToPosition*(s: string): Location {.inline.}
proc getCapture*(self: ChessBoard, move: Move): Location
proc makeMove*(self: ChessBoard, startSquare, targetSquare: string): Move
proc makeMove*(self: ChessBoard, move: Move): Move
func emptyMove*: Move {.inline.} = Move(startSquare: emptyLocation(), targetSquare: emptyLocation(), piece: emptyPiece())
func `+`*(a, b: Location): Location = (a.row + b.row, a.col + b.col)
func isValid*(a: Location): bool = a.row in 0..7 and a.col in 0..7
proc generateMoves(self: ChessBoard, location: Location): seq[Move]

func topLeftDiagonal(piece: Piece): Location {.inline.} = (if piece.color == White: (-1, -1) else: (1, 1))
func topRightDiagonal(piece: Piece): Location {.inline.} = (if piece.color == White: (-1, 1) else: (1, -1))
func bottomLeftDiagonal(piece: Piece): Location {.inline.} = (if piece.color == White: (-1, 1) else: (1, -1))
func bottomRightDiagonal(piece: Piece): Location {.inline.} = (if piece.color == White: (1, 1) else: (-1, -1))
func leftSide(piece: Piece): Location {.inline.} = (if piece.color == White: (0, -1) else: (0, 1))
func rightSide(piece: Piece): Location {.inline.} = (if piece.color == White: (0, 1) else: (0, -1))
func topSide(piece: Piece): Location {.inline.} = (if piece.color == White: (-1, 0) else: (1, 0))
func bottomSide(piece: Piece): Location {.inline.} = (if piece.color == White: (1, 0) else: (-1, 0))
func forward(piece: Piece): Location {.inline.} = (if piece.color == Black: (1, 0) else: (-1, 0))
func doublePush(piece: Piece): Location {.inline.} = (if piece.color == Black: (2, 0) else: (-2, 0))
proc testMoveOffsets(self: ChessBoard, move: Move): bool


proc getActiveColor*(self: ChessBoard): PieceColor =
    ## Returns the currently active color
    ## (turn of who has to move)
    return self.turn


func getStartRow(piece: Piece): int {.inline.} =
    ## Retrieves the starting row of
    ## the given piece inside our 8x8
    ## grid
    case piece.color:
        of None:
            return -1
        of White:
            case piece.kind:
                of Pawn:
                    return 6
                else:
                    return 5
        of Black:
            case piece.kind:
                of Pawn:
                    return 1
                else:
                    return 0


func getLastRow(color: PieceColor): int {.inline.} =
    ## Retrieves the location of the last
    ## row relative to the given color
    case color:
        of White:
            return 0
        of Black:
            return 7
        else:
            return -1


proc newChessboard: ChessBoard =
    ## Returns a new, empty chessboard
    new(result)
    # Turns our flat sequence into an 8x8 grid
    result.grid = newMatrixFromSeq[Piece](empty, (8, 8))
    result.attacked = (@[], @[])
    result.enPassantSquare = emptyMove()
    result.turn = White


proc newChessboardFromFEN*(state: string): ChessBoard =
    ## Initializes a chessboard with the
    ## state encoded by the given FEN string
    result = newChessboard()
    var
        # Current location in the grid
        row = 0
        column = 0
        # Current section in the FEN string
        section = 0
        # Current index into the FEN string
        index = 0
        # Temporary variable to store the piece
        piece: Piece
    # See https://en.wikipedia.org/wiki/Forsyth%E2%80%93Edwards_Notation
    while index <= state.high():
        var c = state[index]
        if c == ' ':
            # Next section
            inc(section)
            inc(index)
            continue
        case section:
            of 0:
                # Piece placement data
                case c.toLowerAscii():
                    # Piece
                    of 'r', 'n', 'b', 'q', 'k', 'p':
                        # We know for a fact these values are in our 
                        # enumeration, so all is good
                        {.push.}
                        {.warning[HoleEnumConv]:off.}
                        piece = Piece(kind: PieceKind(c.toLowerAscii()), color: if c.isUpperAscii(): White else: Black)
                        {.pop.}
                        case piece.color:
                            of Black:
                                case piece.kind:
                                    of Pawn:
                                        result.pieces.black.pawns.add((row, column))
                                    of Bishop:
                                        result.pieces.black.bishops.add((row, column))
                                    of Knight:
                                        result.pieces.black.knights.add((row, column))
                                    of Rook:
                                        result.pieces.black.rooks.add((row, column))
                                    of Queen:
                                        result.pieces.black.queens.add((row, column))
                                    of King:
                                        result.pieces.black.king = (row, column)
                                    else:
                                        discard
                            of White:
                                case piece.kind:
                                    of Pawn:
                                        result.pieces.white.pawns.add((row, column))
                                    of Bishop:
                                        result.pieces.white.bishops.add((row, column))
                                    of Knight:
                                        result.pieces.white.knights.add((row, column))
                                    of Rook:
                                        result.pieces.white.rooks.add((row, column))
                                    of Queen:
                                        result.pieces.white.queens.add((row, column))
                                    of King:
                                        result.pieces.white.king = (row, column)
                                    else:
                                        discard
                            else:
                                discard
                        result.grid[row, column] = piece
                        inc(column)
                    of '/':
                        # Next row
                        inc(row)
                        column = 0
                    of '0'..'9':
                        # Skip x columns
                        let x = int(uint8(c) - uint8('0')) - 1
                        if x > 7:
                            raise newException(ValueError, "invalid skip value (> 8) in FEN string")
                        column += x
                    else:
                        raise newException(ValueError, "invalid piece identifier in FEN string")
            of 1:
                # Active color
                case c:
                    of 'w':
                        result.turn = White
                    of 'b':
                        result.turn = Black
                    else:
                        raise newException(ValueError, "invalid active color identifier in FEN string")
            of 2:
                # Castling availability
                case c:
                    of '-':
                        # Neither side can castle anywhere: do nothing,
                        # as the castling metadata is set to this state
                        # by default
                        discard
                    of 'K':
                        result.castling.white.king = true
                    of 'Q':
                        result.castling.white.queen = true
                    of 'k':
                        result.castling.black.king = true
                    of 'q':
                        result.castling.black.queen = true
                    else:
                        raise newException(ValueError, "invalid castling availability in FEN string")
            of 3:
                # En passant target square
                case c:
                    of '-':
                        # Field is already uninitialized to the correct state
                        discard
                    else:
                        result.enPassantSquare.targetSquare = state[index..index+1].algebraicToPosition()
                        # Just for cleanliness purposes, we fill in the other positional metadata as
                        # well
                        result.enPassantSquare.piece.color = if result.turn == Black: White else: Black
                        result.enPassantSquare.piece.kind = Pawn
                        # Square metadata is 2 bytes long
                        inc(index)
            of 4:
                # Halfmove clock
                var s = ""
                while not state[index].isSpaceAscii():
                    s.add(state[index])
                    inc(index)
                # Backtrack so the space is seen by the
                # next iteration of the loop
                dec(index)
                result.halfMoveClock = parseInt(s)
            of 5:
                # Fullmove number
                var s = ""
                while index <= state.high():
                    s.add(state[index])
                    inc(index)
                result.fullMoveCount = parseInt(s)
            else:
                raise newException(ValueError, "too many fields in FEN string")
        inc(index)


proc newDefaultChessboard*: ChessBoard =
    ## Initializes a chessboard with the
    ## starting position
    return newChessboardFromFEN("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1")


proc countPieces*(self: ChessBoard, kind: PieceKind, color: PieceColor): int =
    ## Counts the number of pieces with
    ## the given color and type
    case color:
        of White:
            case kind:
                of Pawn:
                    return self.pieces.white.pawns.len()
                of Bishop:
                    return self.pieces.white.bishops.len()
                of Knight:
                    return self.pieces.white.knights.len()
                of Rook:
                    return self.pieces.white.rooks.len()
                of Queen:
                    return self.pieces.white.queens.len()
                of King:
                    # There shall be only one, forever
                    return 1
                else:
                    discard
        of Black:
            case kind:
                of Pawn:
                    return self.pieces.black.pawns.len()
                of Bishop:
                    return self.pieces.black.bishops.len()
                of Knight:
                    return self.pieces.black.knights.len()
                of Rook:
                    return self.pieces.black.rooks.len()
                of Queen:
                    return self.pieces.black.queens.len()
                of King:
                    # In perpetuity
                    return 1
                else:
                    discard        
        of None:
            raise newException(ValueError, "invalid piece type")


proc countPieces*(self: ChessBoard, piece: Piece): int =
    ## Returns the number of pieces on the board that
    ## are of the same type and color of the given piece
    return self.countPieces(piece.kind, piece.color)


func rankToColumn(rank: int): int =
    ## Converts a chess rank (1-indexed) 
    ## into a 0-indexed column value for our 
    ## board. This converter is necessary because
    ## chess positions are indexed differently with
    ## respect to our internal representation
    const indeces = [7, 6, 5, 4, 3, 2, 1, 0]
    return indeces[rank - 1]


proc algebraicToPosition*(s: string): Location {.inline.} =
    ## Converts a square location from algebraic
    ## notation to its corresponding row and column
    ## in the chess grid (0 indexed)
    if len(s) != 2:
        raise newException(ValueError, "algebraic position must be of length 2")

    var s = s.toLowerAscii()
    if s[0] notin 'a'..'h':
        raise newException(ValueError, &"algebraic position has invalid first character ('{s[0]}')")
    if s[1] notin '1'..'8':
        raise newException(ValueError, &"algebraic position has invalid second character ('{s[1]}')")

    let file = int(uint8(s[0]) - uint8('a'))
    # Convert the rank character to a number
    let rank = rankToColumn(int(uint8(s[1]) - uint8('0')))
    return (rank, file)


proc getPiece*(self: ChessBoard, square: string): Piece =
    ## Gets the piece on the given square
    ## in algebraic notation
    let loc = square.algebraicToPosition()
    return self.grid[loc.row, loc.col]


proc getCapture*(self: ChessBoard, move: Move): Location =
    ## Returns the location that would be captured if this
    ## move were played on the board, taking en passant and
    ## other things into account (the move is assumed to be
    ## already valid). An empty location is returned if no
    ## piece is captured by the given move
    result = emptyLocation()
    let target = self.grid[move.targetSquare.row, move.targetSquare.col]
    if target.color == None:
        if move.targetSquare != self.enPassantSquare.targetSquare:
            return
        else:
            return ((if move.piece.color == White: move.targetSquare.row + 1 else: move.targetSquare.row - 1), move.targetSquare.col)
    if target.color == move.piece.color.opposite():
        return move.targetSquare


proc isCapture*(self: ChessBoard, move: Move): bool {.inline.} =
    ## Returns whether the given move is a capture
    ## or not
    return self.getCapture(move) != emptyLocation()


proc generatePawnMoves(self: ChessBoard, location: Location): seq[Move] =
    ## Generates the possible moves for the pawn in the given
    ## location
    var 
        piece = self.grid[location.row, location.col]
        locations: seq[Location] = @[]
    doAssert piece.kind == Pawn, &"generatePawnMoves called on a {piece.kind}"
    # Pawns can move forward one square
    let forwardOffset = piece.forward()
    let forward = (forwardOffset + location)
    if forward.isValid() and self.grid[forward.row, forward.col].color == None:
        locations.add(forwardOffset)
    # If the pawn is on its first rank, it can push two squares
    if location.row == piece.getStartRow():
        locations.add(piece.doublePush())
    if self.enPassantSquare.piece.color == piece.color.opposite:
        if abs(self.enPassantSquare.targetSquare.col - location.col) == 1 and abs(self.enPassantSquare.targetSquare.row - location.row) == 1:
            # Only viable if the piece is on the diagonal of the target
            locations.add(self.enPassantSquare.targetSquare)
    # They can also move on either diagonal one
    # square, but only to capture
    if location.col in 1..6:
        # Top right diagonal
        locations.add(piece.topRightDiagonal())
    if location.row in 1..6:
        # Top left diagonal
        locations.add(piece.topLeftDiagonal())
    
    # Pawn is at the right side, can only capture
    # on the left one
    if location.col == 7 and location.row < 7:
        locations.add(piece.topLeftDiagonal())
    # Pawn is at the left side, can only capture
    # on the right one
    if location.col == 0 and location.row < 7:
        locations.add(piece.topRightDiagonal())
    var
        newLocation: Location
        targetPiece: Piece
    for target in locations:
        newLocation = location + target
        if not newLocation.isValid():
            continue
        targetPiece = self.grid[newLocation.row, newLocation.col]
        if targetPiece.color == piece.color:
            # Can't move over a friendly piece
            continue
        if location.col != newLocation.col and not self.isCapture(Move(piece: piece, startSquare: location, targetSquare: newLocation)):
            # Can only move diagonally when capturing
            continue
        if newLocation.row == piece.color.getLastRow():
            # Generate all promotion moves
            for promotionType in [PromoteToKnight, PromoteToBishop, PromoteToRook, PromoteToQueen]:
                result.add(Move(startSquare: location, targetSquare: newLocation, piece: piece, flag: promotionType))
            continue
        # Move is just a pawn push
        result.add(Move(startSquare: location, targetSquare: newLocation, piece: piece))
        

proc generateSlidingMoves(self: ChessBoard, location: Location): seq[Move] =
    ## Generates moves for the sliding piece in the given location
    var 
        piece = self.grid[location.row, location.col]
    doAssert piece.kind in [Bishop, Rook, Queen], &"generateSlidingMoves called on a {piece.kind}"
    var directions: seq[Location] = @[]
    # Only check in the right directions for the chosen piece
    if piece.kind in [Bishop, Queen]:
        directions.add(piece.topLeftDiagonal())
        directions.add(piece.topRightDiagonal())
        directions.add(piece.bottomLeftDiagonal())
        directions.add(piece.bottomRightDiagonal())
    if piece.kind in [Queen, Rook]:
        directions.add(piece.topSide())
        directions.add(piece.bottomSide())
        directions.add(piece.rightSide())
        directions.add(piece.leftSide())
    for direction in directions:
        # Slide in this direction as long as it's possible
        var
            square: Location = location
            otherPiece: Piece
        while true:
            square = square + direction
            # End of board reached
            if not square.isValid():
                break
            otherPiece = self.grid[square.row, square.col]
            # A friendly piece is in the way
            if otherPiece.color == piece.color:
                break
            if otherPiece.color == piece.color.opposite:
                result.add(Move(startSquare: location, targetSquare: square, piece: piece))
                break
            # Target square is empty
            result.add(Move(startSquare: location, targetSquare: square, piece: piece))


proc generateMoves(self: ChessBoard, location: Location): seq[Move] =
    ## Returns the list of possible moves for the
    ## piece in the given location
    let piece = self.grid[location.row, location.col]
    case piece.kind:
        of Queen, Bishop, Rook:
            return self.generateSlidingMoves(location)
        of Pawn:
            return self.generatePawnMoves(location)
        else:
            return @[]


proc testMoveOffsets(self: ChessBoard, move: Move): bool =
    ## Returns true if the piece in the given
    ## move is pseudo-legal: this does not take pins 
    ## nor checks into account, but other rules like 
    ## double pawn pushes and en passant are validated
    ## here. Note that this is an internal method called 
    ## by checkMove and it does not validate whether the
    ## target square is occupied or not (it is assumed the
    ## check has been performed beforehand, like checkMove
    ## does)
    case move.piece.kind:
        of Pawn:
            return move in self.generatePawnMoves(move.startSquare)
        of Bishop, Queen, Rook:
            return move in self.generateSlidingMoves(move.startSquare)
        of Knight:
            # TODO
            discard
        of King:
            # TODO
            discard
        else:
            return false


proc getAttackers*(self: ChessBoard, square: string): seq[Piece] =
    ## Returns all the attackers of the given square
    let loc = square.algebraicToPosition()
    for move in self.attacked.black:
        if move.targetSquare == loc:
            result.add(move.piece)
    for move in self.attacked.white:
        if move.targetSquare == loc:
            result.add(move.piece)

# We don't use getAttackers because this one only cares about whether
# the square is attacked or not (and can therefore exit earlier than
# getAttackers)
proc isAttacked*(self: ChessBoard, loc: Location): bool =
    ## Returns whether the given location is attacked
    ## by the current opponent
    let piece = self.grid[loc.row, loc.col]
    case piece.color:
        of White:
            for move in self.attacked.black:
                if move.targetSquare == loc:
                    return true
        of Black:
            for move in self.attacked.white:
                if move.targetSquare == loc:
                    return true
        of None:
            case self.turn:
                of White:
                    for move in self.attacked.black:
                        if move.targetSquare == loc:
                            return true
                of Black:
                    for move in self.attacked.white:
                        if move.targetSquare == loc:
                            return true
                else:
                    discard



proc isAttacked*(self: ChessBoard, square: string): bool = 
    ## Returns whether the given square is attacked
    ## by the current inactive color
    return self.isAttacked(square.algebraicToPosition())


proc updateAttackedSquares(self: ChessBoard) =
    ## Updates internal metadata about which squares
    ## are attacked. Called internally by doMove
    
    # We refresh the attack metadata at every move. This is an
    # O(1) operation, because we're only updating the length 
    # field without deallocating the memory, which will promptly
    # be reused by us again. Neat!
    self.attacked.white.setLen(0)
    self.attacked.black.setLen(0)
    # Go over each piece one by one and see which squares
    # it currently attacks

    # White pawns
    for loc in self.pieces.white.pawns:
        for move in self.generateMoves(loc):
            self.attacked.white.add(move)
    # Black pawns
    for loc in self.pieces.black.pawns:
        for move in self.generateMoves(loc):
            self.attacked.black.add(move)
    # White bishops
    for loc in self.pieces.white.bishops:
        for move in self.generateMoves(loc):
            self.attacked.white.add(move)
    # Black bishops
    for loc in self.pieces.black.bishops:
        for move in self.generateMoves(loc):
            self.attacked.black.add(move)
    # White rooks
    for loc in self.pieces.white.rooks:
        for move in self.generateMoves(loc):
            self.attacked.white.add(move)
    # Black rooks
    for loc in self.pieces.black.rooks:
        for move in self.generateMoves(loc):
            self.attacked.black.add(move)
    # White queens
    for loc in self.pieces.white.queens:
        for move in self.generateMoves(loc):
            self.attacked.white.add(move)
    # Black Queens
    for loc in self.pieces.black.queens:
        for move in self.generateMoves(loc):
            self.attacked.black.add(move)


proc removePiece(self: ChessBoard, location: Location) =
    ## Removes a piece from the board, updating necessary
    ## metadata
    var piece = self.grid[location.row, location.col]
    self.grid[location.row, location.col] = emptyPiece()
    case piece.color:
        of White:
            case piece.kind:
                of Pawn:
                    self.pieces.white.pawns.delete(self.pieces.white.pawns.find(location))
                of Bishop:
                    self.pieces.white.pawns.delete(self.pieces.white.bishops.find(location))
                of Knight:
                    self.pieces.white.pawns.delete(self.pieces.white.knights.find(location))
                of Rook:
                    self.pieces.white.rooks.delete(self.pieces.white.rooks.find(location))
                of Queen:
                    self.pieces.white.queens.delete(self.pieces.white.queens.find(location))
                of King:
                    doAssert false, "removePiece: attempted to remove the white king"
                else:
                    discard 
        of Black:
            case piece.kind:
                of Pawn:
                    self.pieces.black.pawns.delete(self.pieces.white.pawns.find(location))
                of Bishop:
                    self.pieces.black.bishops.delete(self.pieces.black.bishops.find(location))
                of Knight:
                    self.pieces.black.knights.delete(self.pieces.black.knights.find(location))
                of Rook:
                    self.pieces.black.rooks.delete(self.pieces.black.rooks.find(location))
                of Queen:
                    self.pieces.black.queens.delete(self.pieces.black.queens.find(location))
                of King:
                    doAssert false, "removePiece: attempted to remove the black king"
                else:
                    discard
        else:
            discard


proc updatePositions(self: ChessBoard, move: Move) =
    ## Internal helper to update the position of
    ## the pieces on the board after a move
    let capture = self.getCapture(move)
    if capture != emptyLocation():
        # Move has captured a piece: remove it as well. We call a helper instead
        # of doing it ourselves because there's a bunch of metadata that needs
        # to be updated to do this properly and I thought it'd fit into its neat 
        # little function
        self.removePiece(capture)
    # Update the positional metadata of the moving piece
    case move.piece.color:
        of White:
            case move.piece.kind:
                of Pawn:
                    # The way things are structured, we don't care about the order
                    # of this list, so we can add and remove entries as we please
                    self.pieces.white.pawns.delete(self.pieces.white.pawns.find(move.startSquare))
                    self.pieces.white.pawns.add(move.targetSquare)
                of Bishop:
                    self.pieces.white.bishops.delete(self.pieces.white.bishops.find(move.startSquare))
                    self.pieces.white.bishops.add(move.targetSquare)
                of Knight:
                    self.pieces.white.knights.delete(self.pieces.white.knights.find(move.startSquare))
                    self.pieces.white.knights.add(move.targetSquare)
                of Rook:
                    self.pieces.white.rooks.delete(self.pieces.white.rooks.find(move.startSquare))
                    self.pieces.white.rooks.add(move.targetSquare)
                of Queen:
                    self.pieces.white.queens.delete(self.pieces.white.queens.find(move.startSquare))
                    self.pieces.white.queens.add(move.targetSquare)
                of King:
                    self.pieces.white.king = move.targetSquare
                else:
                    discard
        of Black:
            case move.piece.kind:
                of Pawn:
                    self.pieces.black.pawns.delete(self.pieces.black.pawns.find(move.startSquare))
                    self.pieces.black.pawns.add(move.targetSquare)
                of Bishop:
                    self.pieces.black.bishops.delete(self.pieces.black.bishops.find(move.startSquare))
                    self.pieces.black.bishops.add(move.targetSquare)
                of Knight:
                    self.pieces.black.knights.delete(self.pieces.black.knights.find(move.startSquare))
                    self.pieces.black.knights.add(move.targetSquare)
                of Rook:
                    self.pieces.black.rooks.delete(self.pieces.black.rooks.find(move.startSquare))
                    self.pieces.black.rooks.add(move.targetSquare)
                of Queen:
                    self.pieces.black.queens.delete(self.pieces.black.queens.find(move.startSquare))
                    self.pieces.black.queens.add(move.targetSquare)
                of King:
                    self.pieces.black.king = move.targetSquare
                else:
                    discard
        else:
            discard    
    # Empty out the starting square
    self.grid[move.startSquare.row, move.startSquare.col] = emptyPiece()
    # Actually move the piece
    self.grid[move.targetSquare.row, move.targetSquare.col] = move.piece


proc inCheck(self: ChessBoard): bool =
    ## Returns whether the active color's
    ## king is in check
    case self.turn:
        of White:
            return self.isAttacked(self.pieces.white.king)
        of Black:
            return self.isAttacked(self.pieces.black.king)
        else:
            # Unreachable
            discard
    

proc doMove(self: ChessBoard, move: Move) =
    ## Internal function called by makeMove after
    ## performing legality checks on the given move. Can
    ## be used in performance-critical paths where
    ## a move is already known to be legal
    self.updatePositions(move)    
    self.updateAttackedSquares()
    # En passant is possible only immediately after the
    # pawn has moved
    if self.enPassantSquare != emptyMove() and self.enPassantSquare.piece.color == self.turn.opposite():
        self.enPassantSquare = emptyMove()


proc undoMove(self: ChessBoard, move: Move): Move {.discardable.} =
    ## Undoes the given move if possible
    result = Move(piece: move.piece, startSquare: move.targetSquare, targetSquare: move.startSquare)
    let castling = self.castling
    let enPassant = self.enPassantSquare
    # Swap start and target square and do the move in reverse
    self.doMove(result)
    # We need to reset the entire position when the move is undone!
    self.enPassantSquare = enPassant
    #self.turn = self.turn.opposite()
    self.castling = castling


proc checkMove(self: ChessBoard, move: Move): bool =
    ## Internal function called by makeMove to check a move for legality
    # Start square doesn't contain a piece (and it isn't the en passant square) 
    # or it is of the wrong color for which turn it is to move
    if move.piece.kind == Empty or move.piece.color != self.turn:
        return false
    var destination = self.grid[move.targetSquare.row, move.targetSquare.col]
    # Destination square is occupied by a piece of the same color as the piece
    # being moved: illegal!
    if destination.kind != Empty and destination.color == self.turn:
        return false
    if not self.testMoveOffsets(move):
        # Piece cannot arrive to destination (either
        # because it is blocked or because the moving
        # pattern is incorrect)
        return false
    echo move
    self.doMove(move)
    defer: self.undoMove(move)
    # Move would reveal an attack
    # on our king: not allowed
    if self.inCheck():
        echo "in check"
        return false
    # All checks have passed: move is legal
    result = true


proc makeMove*(self: ChessBoard, move: Move): Move {.discardable.} =
    ## Like the other makeMove(), but with a Move object
    result = move
    if not self.checkMove(move):
        return emptyMove()
    self.doMove(result)
    self.turn = self.turn.opposite()


proc makeMove*(self: ChessBoard, startSquare, targetSquare: string): Move {.discardable.} =
    ## Makes a move on the board from the chosen start square to
    ## the chosen target square, ensuring it is legal (turns are
    ## taken into account!). This function returns a Move object: if the move
    ## is legal and has been performed, the fields will be populated properly.
    ## For efficiency purposes, no exceptions are raised if the move is 
    ## illegal, but the move's piece kind will be Empty (its color will be None
    ## too) and the locations will both be set to the tuple (-1, -1)
    var 
        startLocation = startSquare.algebraicToPosition()
        targetLocation = targetSquare.algebraicToPosition()
    result = Move(startSquare: startLocation, targetSquare: targetLocation, piece: self.grid[startLocation.row, startLocation.col])
    return self.makeMove(result)



proc `$`*(self: ChessBoard): string =
    result &= "- - - - - - - -"
    for i, row in self.grid:
        result &= "\n"
        for piece in row:
            if piece.kind == Empty:
                result &= "x "
                continue
            if piece.color == White:
                result &= &"{char(piece.kind).toUpperAscii()} "
            else:
                result &= &"{char(piece.kind)} "
        result &= &"{rankToColumn(i + 1) + 1}"
    result &= "\n- - - - - - - -"
    result &= "\na b c d e f g h"


proc pretty*(self: ChessBoard): string =
    ## Returns a colorized version of the
    ## board for easier visualization
    result &= "- - - - - - - -"

    for i, row in self.grid:
        result &= "\n"
        for j, piece in row:
            if piece.kind == Empty:
                result &= "\x1b[36;1mx"
                # Avoids the color overflowing
                # onto the numbers
                if j < 7:
                    result &= " \x1b[0m"
                else:
                    result &= "\x1b[0m "
                continue
            if piece.color == White:
                result &= &"\x1b[37;1m{char(piece.kind).toUpperAscii()}\x1b[0m "
            else:
                result &= &"\x1b[30;1m{char(piece.kind)} "
        result &= &"\x1b[33;1m{rankToColumn(i + 1) + 1}\x1b[0m"

    result &= "\n- - - - - - - -"
    result &= "\n\x1b[31;1ma b c d e f g h"
    result &= "\x1b[0m"


when isMainModule:
    proc testPiece(piece: Piece, kind: PieceKind, color: PieceColor) =
        doAssert piece.kind == kind and piece.color == color, &"expected piece of kind {kind} and color {color}, got {piece.kind} / {piece.color} instead"

    proc testPieceCount(board: ChessBoard, kind: PieceKind, color: PieceColor, count: int) =
        let pieces = board.countPieces(kind, color)
        doAssert pieces == count, &"expected {count} pieces of kind {kind} and color {color}, got {pieces} instead"

    echo "Running tests"
    var b = newDefaultChessboard()
    # Ensure correct number of pieces
    testPieceCount(b, Pawn, White, 8)
    testPieceCount(b, Pawn, Black, 8)
    testPieceCount(b, Knight, White, 2)
    testPieceCount(b, Knight, Black, 2)
    testPieceCount(b, Bishop, White, 2)
    testPieceCount(b, Bishop, Black, 2)
    testPieceCount(b, Rook, White, 2)
    testPieceCount(b, Rook, Black, 2)
    testPieceCount(b, Queen, White, 1)
    testPieceCount(b, Queen, Black, 1)
    testPieceCount(b, King, White, 1)
    testPieceCount(b, King, Black, 1)
    

    # Ensure pieces are in the correct location

    # Pawns
    for loc in ["a2", "b2", "c2", "d2", "e2", "f2", "g2", "h2"]:
        testPiece(b.getPiece(loc), Pawn, White)
    for loc in ["a7", "b7", "c7", "d7", "e7", "f7", "g7", "h7"]:
        testPiece(b.getPiece(loc), Pawn, Black)
    # Rooks
    testPiece(b.getPiece("a1"), Rook, White)
    testPiece(b.getPiece("h1"), Rook, White)
    testPiece(b.getPiece("a8"), Rook, Black)
    testPiece(b.getPiece("h8"), Rook, Black)
    # Knights
    testPiece(b.getPiece("b1"), Knight, White)
    testPiece(b.getPiece("g1"), Knight, White)
    testPiece(b.getPiece("b8"), Knight, Black)
    testPiece(b.getPiece("g8"), Knight, Black)
    # Bishops
    testPiece(b.getPiece("c1"), Bishop, White)
    testPiece(b.getPiece("f1"), Bishop, White)
    testPiece(b.getPiece("c8"), Bishop, Black)
    testPiece(b.getPiece("f8"), Bishop, Black)
    # Kings
    testPiece(b.getPiece("e1"), King, White)
    testPiece(b.getPiece("e8"), King, Black)
    # Queens
    testPiece(b.getPiece("d1"), Queen, White)
    testPiece(b.getPiece("d8"), Queen, Black)
    echo "All tests were successful"