from __future__ import annotations
import os
import random
import sys
import chess
from search import Searcher

ENGINE_NAME = "ChessMamba"
ENGINE_AUTHOR = "CatGirlHtfi"
DEFAULT_CKPT = os.environ.get(
    "CHESSMAMBA_CKPT", os.path.join(os.path.dirname(__file__), "ckpt", "model.pt")
)
POLICY_ONLY = os.environ.get("CHESSMAMBA_POLICY_ONLY", "") == "1"


def send(msg: str) -> None:
    sys.stdout.write(msg + "\n")
    sys.stdout.flush()


def main() -> None:
    board = chess.Board()
    searcher = Searcher(checkpoint_path=DEFAULT_CKPT, policy_only=POLICY_ONLY)
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        tokens = line.split()
        cmd = tokens[0]
        if cmd == "uci":
            send(f"id name {ENGINE_NAME}")
            send(f"id author {ENGINE_AUTHOR}")
            send("uciok")
        elif cmd == "isready":
            send("readyok")
        elif cmd == "ucinewgame":
            board = chess.Board()
            searcher.reset()
        elif cmd == "position":
            _handle_position(board, tokens)
            searcher.sync(board)
        elif cmd == "go":
            move = searcher.choose_move(board, tokens)
            if move is None:
                legal = list(board.legal_moves)
                move = random.choice(legal) if legal else None
            if move is not None:
                send(f"bestmove {move.uci()}")
            else:
                send("bestmove 0000")
        elif cmd == "quit":
            break
        elif cmd == "stop":
            pass


def _handle_position(board: chess.Board, tokens: list[str]) -> None:
    if "startpos" in tokens:
        board.reset()
        idx = tokens.index("startpos") + 1
    elif "fen" in tokens:
        fen_idx = tokens.index("fen")
        moves_idx = tokens.index("moves") if "moves" in tokens else len(tokens)
        fen = " ".join(tokens[fen_idx + 1 : moves_idx])
        board.set_fen(fen)
        idx = moves_idx
    else:
        return
    if idx < len(tokens) and tokens[idx] == "moves":
        for uci_move in tokens[idx + 1 :]:
            board.push_uci(uci_move)


if __name__ == "__main__":
    main()