import ../nimfish


import std/strformat
import std/strutils
import std/times
import std/math

from std/lenientops import `/`


type
    CountData = tuple[nodes: uint64, captures: uint64, castles: uint64, checks: uint64,  promotions: uint64, enPassant: uint64, checkmates: uint64]


proc perft*(self: ChessBoard, ply: int, verbose: bool = false, divide: bool = false, bulk: bool = false): CountData =
    ## Counts (and debugs) the number of legal positions reached after
    ## the given number of ply
    
    var moves = MoveList()
    self.generateMoves(moves)
    if not bulk:
        if len(moves) == 0 and self.inCheck():
            result.checkmates = 1
        # TODO: Should we count stalemates/draws?
        if ply == 0:
            result.nodes = 1
            return
    elif ply == 1 and bulk:
        if divide:
            var postfix = ""
            for move in moves:
                case move.getPromotionType():
                    of PromoteToBishop:
                        postfix = "b"
                    of PromoteToKnight:
                        postfix = "n"
                    of PromoteToRook:
                        postfix = "r"
                    of PromoteToQueen:
                        postfix = "q"
                    else:
                        postfix = ""
                echo &"{move.startSquare.toAlgebraic()}{move.targetSquare.toAlgebraic()}{postfix}: 1"
                if verbose:
                    echo ""
        return (uint64(len(moves)), 0, 0, 0, 0, 0, 0)

    for move in moves:
        if verbose:
            let canCastle = self.canCastle(self.getSideToMove())
            echo &"Ply (from root): {self.getPlyFromRoot()}"
            echo &"Move: {move.startSquare.toAlgebraic()}{move.targetSquare.toAlgebraic()}"
            echo &"Turn: {self.getSideToMove()}"
            echo &"Piece: {self.getPiece(move.startSquare).kind}"
            echo &"Flags: {move.getFlags()}"
            echo &"In check: {(if self.inCheck(): \"yes\" else: \"no\")}"
            echo &"Can castle:\n  - King side: {(if canCastle.king: \"yes\" else: \"no\")}\n  - Queen side: {(if canCastle.queen: \"yes\" else: \"no\")}"
            echo &"Position before move: {self.toFEN()}"
            stdout.write("En Passant target: ")
            if self.getEnPassantTarget() != nullSquare():
                echo self.getEnPassantTarget().toAlgebraic()
            else:
                echo "None"
            echo "\n", self.pretty()
        self.doMove(move)
        if ply == 1:
            if move.isCapture():
                inc(result.captures)
            if move.isCastling():
                inc(result.castles)
            if move.isPromotion():
                inc(result.promotions)
            if move.isEnPassant():
                inc(result.enPassant)
        if self.inCheck():
            # Opponent king is in check
            inc(result.checks)
        if verbose:
            let canCastle = self.canCastle(self.getSideToMove())
            echo "\n"
            echo &"Opponent in check: {(if self.inCheck(): \"yes\" else: \"no\")}"
            echo &"Opponent can castle:\n  - King side: {(if canCastle.king: \"yes\" else: \"no\")}\n  - Queen side: {(if canCastle.queen: \"yes\" else: \"no\")}"
            echo &"Position after move: {self.toFEN()}"
            echo "\n", self.pretty()
            stdout.write("nextpos>> ")
            try:
                discard readLine(stdin)
            except IOError:
                discard
            except EOFError:
                discard
        let next = self.perft(ply - 1, verbose, bulk=bulk)
        self.unmakeMove()
        if divide and (not bulk or ply > 1):
            var postfix = ""
            if move.isPromotion():
                case move.getPromotionType():
                    of PromoteToBishop:
                        postfix = "b"
                    of PromoteToKnight:
                        postfix = "n"
                    of PromoteToRook:
                        postfix = "r"
                    of PromoteToQueen:
                        postfix = "q"
                    else:
                        discard
            echo &"{move.startSquare.toAlgebraic()}{move.targetSquare.toAlgebraic()}{postfix}: {next.nodes}"
            if verbose:
                echo ""
        result.nodes += next.nodes
        result.captures += next.captures
        result.checks += next.checks
        result.promotions += next.promotions
        result.castles += next.castles
        result.enPassant += next.enPassant
        result.checkmates += next.checkmates


proc handleGoCommand(board: ChessBoard, command: seq[string]) =
    if len(command) < 2:
        echo &"Error: go: invalid number of arguments"
        return
    case command[1]:
        of "perft":
            if len(command) == 2:
                echo &"Error: go: perft: invalid number of arguments"
                return
            var 
                args = command[2].splitWhitespace()
                bulk = false
                verbose = false
            if args.len() > 1:
                var ok = true
                for arg in args[1..^1]:
                    case arg:
                        of "bulk":
                            bulk = true
                        of "verbose":
                            verbose = true
                        else:
                            echo &"Error: go: perft: invalid argument '{args[1]}'"
                            ok = false
                            break
                if not ok:
                    return
            try:
                let ply = parseInt(args[0])
                if bulk:
                    let t = cpuTime()
                    let nodes = board.perft(ply, divide=true, bulk=true, verbose=verbose).nodes
                    let tot = cpuTime() - t
                    echo &"\nNodes searched (bulk-counting: on): {nodes}"
                    echo &"Time taken: {round(tot, 3)} seconds\nNodes per second: {round(nodes / tot).uint64}"
                else:
                    let t = cpuTime()
                    let data = board.perft(ply, divide=true, verbose=verbose)
                    let tot = cpuTime()
                    echo &"\nNodes searched (bulk-counting: off): {data.nodes}"
                    echo &"  - Captures: {data.captures}"
                    echo &"  - Checks: {data.checks}"
                    echo &"  - E.P: {data.enPassant}"
                    echo &"  - Checkmates: {data.checkmates}"
                    echo &"  - Castles: {data.castles}"
                    echo &"  - Promotions: {data.promotions}"
                    echo ""
                    echo &"Time taken: {round(t, 3)} seconds\nNodes per second: {round(data.nodes / tot).uint64}"
            except ValueError:
                echo "Error: go: perft: invalid depth"
        else:
            echo &"Error: go: unknown subcommand '{command[1]}'"


proc handleMoveCommand(board: ChessBoard, command: seq[string]): Move {.discardable.} =
    if len(command) != 2:
        echo &"Error: move: invalid number of arguments"
        return
    let moveString = command[1]
    if len(moveString) notin 4..5:
        echo &"Error: move: invalid move syntax"
        return
    var
        startSquare: Square
        targetSquare: Square
        flags: seq[MoveFlag]
    
    try:
        startSquare = moveString[0..1].toSquare()
    except ValueError:
        echo &"Error: move: invalid start square ({moveString[0..1]})"
        return
    try:
        targetSquare = moveString[2..3].toSquare()
    except ValueError:
        echo &"Error: move: invalid target square ({moveString[2..3]})"
        return

    # Since the user tells us just the source and target square of the move,
    # we have to figure out all the flags by ourselves (whether it's a double
    # push, a capture, a promotion, etc.)
    
    if board.getPiece(targetSquare).kind != Empty:
        flags.add(Capture)

    elif board.getPiece(targetSquare).kind == Pawn and abs(rankFromSquare(startSquare) - rankFromSquare(targetSquare)) == 2:
        flags.add(DoublePush)

    if len(moveString) == 5:
        # Promotion
        case moveString[4]:
            of 'b':
                flags.add(PromoteToBishop)
            of 'n':
                flags.add(PromoteToKnight)
            of 'q':
                flags.add(PromoteToQueen)
            of 'r':
                flags.add(PromoteToRook)
            else:
                echo &"Error: move: invalid promotion type"
                return
    
    
    var move = createMove(startSquare, targetSquare, flags)
    let piece = board.getPiece(move.startSquare)
    if piece.kind == King and move.startSquare == board.getSideToMove().getKingStartingSquare():
        if move.targetSquare == longCastleKing(piece.color):
            move.flags = move.flags or CastleLong.uint16
        elif move.targetSquare == shortCastleKing(piece.color):
            move.flags = move.flags or CastleShort.uint16
    if move.targetSquare == board.getEnPassantTarget():
        move.flags = move.flags or EnPassant.uint16
    result = board.makeMove(move)
    if result == nullMove():
        echo &"Error: move: {moveString} is illegal"


proc handlePositionCommand(board: var ChessBoard, command: seq[string]) =
    if len(command) < 2:
        echo "Error: position: invalid number of arguments"
        return
    # Makes sure we don't leave the board in an invalid state if
    # some error occurs
    var tempBoard: ChessBoard
    case command[1]:
        of "startpos":
            tempBoard = newDefaultChessboard()
            if command.len() > 2:
                let args = command[2].splitWhitespace()
                if args.len() > 0:
                    var i = 0
                    while i < args.len():
                        case args[i]:
                            of "moves":
                                var j = i + 1
                                while j < args.len():
                                    if handleMoveCommand(tempBoard, @["move", args[j]]) == nullMove():
                                        return
                                    inc(j)
                        inc(i)
            board = tempBoard
        of "fen":
            if len(command) == 2:
                echo &"Current position: {board.toFEN()}"
                return
            var 
                args = command[2].splitWhitespace()
                fenString = ""
                stop = 0
            for i, arg in args:
                if arg in ["moves", ]:
                    break
                if i > 0:
                    fenString &= " "
                fenString &= arg
                inc(stop)
            args = args[stop..^1]
            try:
                tempBoard = newChessboardFromFEN(fenString)
            except ValueError:
                echo &"error: position: {getCurrentExceptionMsg()}"
                return
            if args.len() > 0:
                var i = 0
                while i < args.len():
                    case args[i]:
                        of "moves":
                            var j = i + 1
                            while j < args.len():
                                if handleMoveCommand(tempBoard, @["move", args[j]]) == nullMove():
                                    return
                                inc(j)
                    inc(i)
            board = tempBoard
        of "print":
            echo board
        of "pretty":
            echo board.pretty()
        else:
            echo &"error: position: unknown subcommand '{command[1]}'"
            return


proc handleUCICommand(board: var ChessBoard, command: seq[string]) =
    echo "id name Nimfish 0.1"
    echo "id author Nocturn9x & Contributors (see LICENSE)"
    # TODO
    echo "uciok"


const HELP_TEXT = """Nimfish help menu:
    - go: Begin a search
          Subcommands: 
            - perft <depth> [options]: Run the performance test at the given depth (in ply) and
              print the results
              Options:
                - bulk: Enable bulk-counting (significantly faster, gives less statistics)
                - verbose: Enable move debugging (for each and every move, not recommended on large searches)
        Example: go perft 5 bulk
    - position: Get/set board position
                Subcommands:
                  - fen [string]: Set the board to the given fen string if one is provided, or print
                    the current position as a FEN string if no arguments are given
                  - startpos: Set the board to the starting position
                  - pretty: Pretty-print the current position
                  - print: Print the current position using ASCII characters only
                  Options:
                    - moves {moveList}: Perform the given moves (space-separated, all-lowercase)
                        in algebraic notation after the position is loaded. This option only applies
                        to the "startpos" and "fen" subcommands: it is ignored otherwise
                Examples:
                    - position startpos
                    - position fen "..." moves a2a3 a7a6
    - clear: Clear the screen
    - move <move>: Perform the given move in algebraic notation
    - castle: Print castlingRights rights for each side
    - check: Print if the current side to move is in check
    - unmove, u: Unmakes the last move. Can be used in succession
    - stm: Print which side is to move
    - ep: Print the current en passant target
    - pretty: Shorthand for "position pretty"
    - print: Shorthand for "position print"
    - fen: Shorthand for "position fen"
    - pos <args>: Shorthand for "position <args>"
    - get <square>: Get the piece on the given square
    - atk <square>: Print the attack bitboard of the given square for the side to move
    - skip: Swap the side to move
    - uci: enter UCI mode (WIP)
    - quit: exit
"""


proc commandLoop*: int =
    ## Nimfish's control interface
    echo "Nimfish by nocturn9x (see LICENSE)"
    var 
        board = newDefaultChessboard()
        uciMode = false
    while true:
        var
            cmd: seq[string]
            cmdStr: string
        try:
            if not uciMode:
                stdout.write(">>> ")
                stdout.flushFile()
            cmdStr = readLine(stdin).strip(leading=true, trailing=true, chars={'\t', ' '})
            if cmdStr.len() == 0:
                continue
            cmd = cmdStr.splitWhitespace(maxsplit=2)

            case cmd[0]:
                of "uci":
                    handleUCICommand(board, cmd)
                    uciMode = true
                of "clear":
                    echo "\x1Bc"
                of "help":
                    echo HELP_TEXT
                of "skip":
                    board.setSideToMove(board.getSideToMove().opposite())
                of "go":
                    handleGoCommand(board, cmd)
                of "position", "pos":
                    handlePositionCommand(board, cmd)
                of "move":
                    handleMoveCommand(board, cmd)
                of "pretty", "print", "fen":
                    handlePositionCommand(board, @["position", cmd[0]])
                of "unmove", "u":
                    board.unmakeMove()
                of "stm":
                    echo &"Side to move: {board.getSideToMove()}"
                of "atk":
                    if len(cmd) != 2:
                        echo "error: atk: invalid number of arguments"
                        continue
                    try:
                        echo board.getAttacksTo(cmd[1].toSquare(), board.getSideToMove())
                    except ValueError:
                        echo "error: atk: invalid square"
                        continue
                of "ep":
                    let target = board.getEnPassantTarget()
                    if target != nullSquare():
                        echo &"En passant target: {target.toAlgebraic()}"
                    else:
                        echo "En passant target: None"
                of "get":
                    if len(cmd) != 2:
                        echo "error: get: invalid number of arguments"
                        continue
                    try:
                        echo board.getPiece(cmd[1])
                    except ValueError:
                        echo "error: get: invalid square"
                        continue                    
                of "castle":
                    let canCastle = board.canCastle(board.getSideToMove())
                    echo &"Castling rights for {($board.getSideToMove()).toLowerAscii()}:\n  - King side: {(if canCastle.king: \"yes\" else: \"no\")}\n  - Queen side: {(if canCastle.queen: \"yes\" else: \"no\")}"
                of "check":
                    echo &"{board.getSideToMove()} king in check: {(if board.inCheck(): \"yes\" else: \"no\")}"
                else:
                    echo &"Unknown command '{cmd[0]}'. Type 'help' for more information."
        except IOError:
            echo ""
            return 0
        except EOFError:
            echo ""
            return 0