# Copyright 2024 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.

## Implementation of a UCI compatible server
import std/strutils
import std/strformat
import std/atomics


import board
import movegen
import search
import transpositions


type
    UCISession = object
        ## A UCI session
        debug: bool
        # Previously reached positions
        history: seq[Position]
        # The current position
        position: Position
        # Atomic boolean flag to interrupt the search
        stopFlag: ptr Atomic[bool]
        # Atomic search flag used to know whether a search
        # is in progress
        searchFlag: ptr Atomic[bool]
        # Size of the transposition table (in megabytes)
        hashTableSize: uint64
        # The transposition table
        transpositionTable: ptr TTable
        # Storage for our history heuristic
        historyTable: ptr HistoryTable
        # Storage for our killer move heuristic
        killerMoves: ptr KillersTable
    
    UCICommandType = enum
        ## A UCI command type enumeration
        Unknown,
        IsReady,
        NewGame,
        Quit,
        Debug,
        Position,
        SetOption,
        Go,
        Stop

    UCICommand = object
        ## A UCI command
        case kind: UCICommandType
            of Debug:
                on: bool
            of Position:
                fen: string
                moves: seq[string]
            of SetOption:
                name: string
                value: string
            of Unknown:
                reason: string
            of Go:
                wtime: int
                btime: int
                winc: int
                binc: int
                movesToGo: int
                depth: int
                moveTime: int
                nodes: uint64
                searchmoves: seq[Move]
            else:
                discard


proc parseUCIMove(position: Position, move: string): tuple[move: Move, command: UCICommand] =
    var
        startSquare: Square
        targetSquare: Square
        flags: seq[MoveFlag]
    if len(move) notin 4..5:
        return (nullMove(), UCICommand(kind: Unknown, reason: "invalid move syntax"))
    try:
        startSquare = move[0..1].toSquare()
    except ValueError:
        return (nullMove(), UCICommand(kind: Unknown, reason: &"invalid start square {move[0..1]}"))
    try:
        targetSquare = move[2..3].toSquare()
    except ValueError:
        return (nullMove(), UCICommand(kind: Unknown, reason: &"invalid target square {move[2..3]}"))

    # Since the client 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 position.getPiece(targetSquare).kind != Empty:
        flags.add(Capture)

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

    if len(move) == 5:
        # Promotion
        case move[4]:
            of 'b':
                flags.add(PromoteToBishop)
            of 'n':
                flags.add(PromoteToKnight)
            of 'q':
                flags.add(PromoteToQueen)
            of 'r':
                flags.add(PromoteToRook)
            else:
                return
    let piece = position.getPiece(startSquare)
    if piece.kind == King and startSquare == position.sideToMove.getKingStartingSquare():
        if targetSquare in [piece.kingSideCastling(), piece.queenSideCastling()]:
            flags.add(Castle)
    elif piece.kind == Pawn and targetSquare == position.enPassantSquare:
        # I hate en passant I hate en passant I hate en passant I hate en passant I hate en passant I hate en passant 
        flags.add(EnPassant)
    result.move = createMove(startSquare, targetSquare, flags)


proc handleUCIMove(session: var UCISession, board: var Chessboard, move: string): tuple[move: Move, cmd: UCICommand] {.discardable.} =
    if session.debug:
        echo &"info string making move {move}"
    let 
        r = board.position.parseUCIMove(move)
        move = r.move
        command = r.command
    if move == nullMove():
        return (move, command)
    else:
        result.move = board.makeMove(move)


proc handleUCIGoCommand(session: UCISession, command: seq[string]): UCICommand =
    result = UCICommand(kind: Go)
    result.wtime = 0
    result.btime = 0
    result.winc = 0
    result.binc = 0
    result.movesToGo = 0
    result.depth = -1
    result.moveTime = -1
    result.nodes = 0
    var 
        current = 1   # Skip the "go"
    while current < command.len():
        let flag = command[current]
        inc(current)
        case flag:
            of "infinite":
                result.wtime = int32.high()
                result.btime = int32.high()
            of "wtime":
                result.wtime = command[current].parseInt()
            of "btime":
                result.btime = command[current].parseInt()
            of "winc":
                result.winc = command[current].parseInt()
            of "binc":
                result.binc = command[current].parseInt()
            of "movestogo":
                result.movesToGo = command[current].parseInt()
            of "depth":
                result.depth = command[current].parseInt()
            of "movetime":
                result.moveTime = command[current].parseInt()
            of "nodes":
                result.nodes = command[current].parseBiggestUInt()
            of "searchmoves":
                while current < command.len():
                    inc(current)
                    if command[current] == "":
                        break
                    let move = session.position.parseUCIMove(command[current]).move
                    if move == nullMove():
                        return UCICommand(kind: Unknown, reason: &"invalid move '{command[current]}' for searchmoves")
                    result.searchmoves.add(move)
            else:
                discard


proc handleUCIPositionCommand(session: var UCISession, command: seq[string]): UCICommand =
    # Makes sure we don't leave the board in an invalid state if
    # some error occurs
    result = UCICommand(kind: Position)
    var chessboard = newChessboard()
    case command[1]:
        of "startpos":
            result.fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
            chessboard.position = startpos()
            if command.len() > 2:
                let args = command[2..^1]
                if args.len() > 0:
                    var i = 0
                    while i < args.len():
                        case args[i]:
                            of "moves":
                                var j = i + 1
                                while j < args.len():
                                    let r = handleUCIMove(session, chessboard, args[j])
                                    if r.move == nullMove():
                                        if r.cmd.reason.len() > 0:
                                            return UCICommand(kind: Unknown, reason: &"move {args[j]} is illegal or invalid ({r.cmd.reason})")
                                        else:
                                            return UCICommand(kind: Unknown, reason: &"move {args[j]} is illegal or invalid")
                                    result.moves.add(args[j])
                                    inc(j)
                        inc(i)
        of "fen":
            var 
                args = command[2..^1]
                fenString = ""
                stop = 0
            for i, arg in args:
                if arg in ["moves", ]:
                    break
                if i > 0:
                    fenString &= " "
                fenString &= arg
                inc(stop)
            result.fen = fenString
            args = args[stop..^1]
            chessboard.position = loadFEN(result.fen)
            if args.len() > 0:
                var i = 0
                while i < args.len():
                    case args[i]:
                        of "moves":
                            var j = i + 1
                            while j < args.len():
                                let r = handleUCIMove(session, chessboard, args[j])
                                if r.move == nullMove():
                                    if r.cmd.reason.len() > 0:
                                        return UCICommand(kind: Unknown, reason: &"move {args[j]} is illegal or invalid ({r.cmd.reason})")
                                    else:
                                        return UCICommand(kind: Unknown, reason: &"move {args[j]} is illegal or invalid")
                                result.moves.add(args[j])
                                inc(j)
                    inc(i)
        else:
            return UCICommand(kind: Unknown, reason: &"unknown subcomponent '{command[1]}'")
    session.position = chessboard.position
    session.history = chessboard.positions


proc parseUCICommand(session: var UCISession, command: string): UCICommand =
    ## Attempts to parse the given UCI command
    var cmd = command.replace("\t", "").splitWhitespace()
    result = UCICommand(kind: Unknown)
    var current = 0
    while current < cmd.len():
        case cmd[current]:
            of "isready":
                return UCICommand(kind: IsReady)
            of "stop":
                return UCICommand(kind: Stop)
            of "ucinewgame":
                return UCICommand(kind: NewGame)
            of "quit":
                return UCICommand(kind: Quit)
            of "debug":
                if current == cmd.high():
                    return
                case cmd[current + 1]:
                    of "on":
                        return UCICommand(kind: Debug, on: true)
                    of "off":
                        return UCICommand(kind: Debug, on: false)
                    else:
                        return
            of "position":
                return session.handleUCIPositionCommand(cmd)
            of "go":
                return session.handleUCIGoCommand(cmd)
            of "setoption":
                result = UCICommand(kind: SetOption)
                inc(current)
                while current < cmd.len():
                    case cmd[current]:
                        of "name":
                            inc(current)
                            result.name = cmd[current]
                        of "value":
                            inc(current)
                            result.value = cmd[current]
                        else:
                            discard
                    inc(current)
                     
            else:
                # Unknown UCI commands should be ignored. Attempt
                # to make sense of the input regardless
                inc(current)


proc bestMove(args: tuple[session: UCISession, command: UCICommand]) {.thread.} =
    ## Finds the best move in the current position
    setControlCHook(proc () {.noconv.} = quit(0))

    # Yes yes nim sure this isn't gcsafe. Now stfu and spawn a thread
    {.cast(gcsafe).}:
        var session = args.session
        var board = newChessboard()
        board.position = session.position
        board.positions = session.history
        let command = args.command
        var searcher = newSearchManager(board, session.transpositionTable, session.stopFlag,
                                        session.searchFlag, session.historyTable, session.killerMoves)
        var 
            timeRemaining = (if session.position.sideToMove == White: command.wtime else: command.btime)
            increment = (if session.position.sideToMove == White: command.winc else: command.binc)
        if command.moveTime != -1:
            timeRemaining = 0
            increment = command.moveTime
        elif timeRemaining == 0:
            timeRemaining = int32.high()
        var move = searcher.findBestMove(timeRemaining, increment, command.depth, command.nodes, command.searchmoves)
        echo &"bestmove {move.toAlgebraic()}"


proc startUCISession* =
    ## Begins listening for UCI commands
    echo "id name Nimfish 0.1"
    echo "id author Nocturn9x & Contributors (see LICENSE)"
    echo "option name Hash type spin default 64 min 1 max 33554432"
    echo "uciok"
    var
        cmd: UCICommand
        cmdStr: string
        session = UCISession(hashTableSize: 64, position: startpos())
    # God forbid we try to use atomic ARC like it was intended. Raw pointers
    # it is then... sigh
    session.transpositionTable = cast[ptr TTable](alloc0(sizeof(TTable)))
    session.stopFlag = cast[ptr Atomic[bool]](alloc0(sizeof(Atomic[bool])))
    session.searchFlag = cast[ptr Atomic[bool]](alloc0(sizeof(Atomic[bool])))
    session.transpositionTable[] = newTranspositionTable(session.hashTableSize * 1024 * 1024)
    session.historyTable = cast[ptr HistoryTable](alloc0(sizeof(HistoryTable)))
    session.killerMoves = cast[ptr KillersTable](alloc0(sizeof(KillersTable)))
    # Initialize killer move array
    for i in 0..<MAX_DEPTH:
        for j in 0..<NUM_KILLERS:
            session.killerMoves[i][j] = nullMove()
    # Fun fact, nim doesn't collect the memory of thread vars. Another stupid fucking design pitfall
    # of nim's AWESOME threading model. Someone is getting a pipebomb in their mailbox about this, mark
    # my fucking words. (for legal purposes THAT IS A JOKE). See https://github.com/nim-lang/Nim/issues/23165
    # The solution? Just reuse the same thread object so the leak is isolated to a single thread. 
    # Also the nim allocator has internal races, so we gotta lose performance by using -d:useMalloc instead.
    # At least mimalloc exists.
    # THANKS ARAQ 
    var searchThread: Thread[tuple[session: UCISession, command: UCICommand]]
    while true:
        try:
            cmdStr = readLine(stdin).strip(leading=true, trailing=true, chars={'\t', ' '})
            if cmdStr.len() == 0:
                if session.debug:
                    echo "info string received empty input, ignoring it"
                continue
            cmd = session.parseUCICommand(cmdStr)
            if cmd.kind == Unknown:
                if session.debug:
                    echo &"info string received unknown or invalid command '{cmdStr}' -> {cmd.reason}"
                continue
            if session.debug:
                echo &"info string received command '{cmdStr}' -> {cmd}"
            case cmd.kind:
                of Quit:
                    quit(0)
                of IsReady:
                    echo "readyok"
                of Debug:
                    session.debug = cmd.on
                of NewGame:
                    session.position = startpos()
                    session.history = @[]
                of Go:
                    # when not defined(noScaleHistory):
                    # Scale our history coefficients
                    for color in PieceColor.White..PieceColor.Black:
                        for source in Square(0)..Square(63):
                            for target in Square(0)..Square(63):
                                session.historyTable[][color][source][target] = session.historyTable[][color][source][target] div 2
                    if searchThread.running:
                        joinThread(searchThread)
                    createThread(searchThread, bestMove, (session, cmd))
                    if session.debug:
                        echo "info string search started"
                of Stop:
                    if not session.searchFlag[].load():
                        continue
                    session.stopFlag[].store(true)
                    joinThread(searchThread)
                    if session.debug:
                        echo "info string search stopped"
                of SetOption:
                    if session.searchFlag[].load():
                        # Cannot set options during search
                        continue
                    case cmd.name:
                        of "Hash":
                            session.hashTableSize = cmd.value.parseBiggestUInt()
                            if session.debug:
                                echo &"info string set TT hash table size to {session.hashTableSize} MiB"
                                session.transpositionTable[] = newTranspositionTable(session.hashTableSize * 1024 * 1024)
                        else:
                            discard
                of Position:
                    discard
                else:
                    discard
        except IOError:
            if session.debug:
                echo "info string I/O error while reading from stdin, exiting"
            echo ""
            quit(0)
        except EOFError:
            if session.debug:
                echo "info string EOF received while reading from stdin, exiting"
            echo ""
            quit(0)