#!/usr/bin/env python3
import sys
import torch
import chess
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "zhoudoe23/ChessQween3-base-puzzled"

print(f"[Engine] 正在加载模型: {MODEL_ID} ...", file=sys.stderr)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
model.eval()
print("[Engine] 模型加载成功，准备下棋！", file=sys.stderr)

def format_fen_for_model(fen_str: str) -> str:
    """将标准 FEN 转化为模型训练时的格式"""
    parts = fen_str.split(" ")
    # 1. 棋盘字符间加空格
    ranks = parts[0].split("/")
    formatted_ranks = [" ".join(list(rank)) for rank in ranks]
    formatted_board = " / ".join(formatted_ranks)
    
    # 2. 颜色转换
    color = "<|WHITE|>" if parts[1] == "w" else "<|BLACK|>"
    
    # 3. 组合其他字段
    castling = parts[2]
    ep = parts[3]
    halfmove = parts[4]
    fullmove = parts[5]
    
    return f"{formatted_board} {color} {castling} {ep} 0 1 >"

def is_suicide_move(board: chess.Board, move: chess.Move) -> bool:
    """精准判断：这一步是不是把重子（车/后）送给对方白吃（Hanging Piece）"""
    moving_piece = board.piece_at(move.from_square)
    if not moving_piece:
        return False
        
    # 只针对车 (Rook) 和 后 (Queen) 这种大子做防御（送兵/送马有时是战术弃子，但送车后绝对崩盘）
    if moving_piece.piece_type not in [chess.ROOK, chess.QUEEN]:
        return False

    my_color = board.turn
    dest_square = move.to_square

    # 模拟走这步棋
    board.push(move)

    # 1. 如果走完这步直接【绝杀（Checkmate）】，那送子也无所谓，这是好棋！
    if board.is_checkmate():
        board.pop()
        return False

    opponent_color = board.turn # 走完后轮到对方

    # 2. 检查落子点：落子后，这个格子是否在对方的攻击火力下？
    is_under_attack = board.is_attacked_by(opponent_color, dest_square)

    # 3. 检查落子点：落子后，这个格子是否有我方其他棋子在保护？
    is_defended = board.is_attacked_by(my_color, dest_square)

    board.pop()

    # 【核心逻辑】：如果落子点【暴露在对方火力下】且【我方没有任何棋子守护】，这就是纯送死！
    if is_under_attack and not is_defended:
        return True # 判定为纯送死（自杀）

    return False

def get_best_move(board: chess.Board) -> chess.Move:
    raw_fen = board.fen()
    prompt = format_fen_for_model(raw_fen)

    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits[0, -1, :]

    # 将所有 Token 按模型预测的概率从高到低排序
    sorted_indices = torch.argsort(logits, descending=True)

    candidate_moves = []
    for token_id in sorted_indices:
        move_str = tokenizer.decode(token_id.item()).strip()
        try:
            move = chess.Move.from_uci(move_str)
            if move in board.legal_moves:
                candidate_moves.append(move)
        except ValueError:
            continue

    # 🌟 第一优先级：找【概率高】且【不是自杀送车送后】的走法
    for move in candidate_moves:
        if not is_suicide_move(board, move):
            return move  # 找到了兼具战术直觉与安全的完美走法！

    # 🌟 保底：如果模型脑子抽风，认为所有走法都是自杀，则被迫返回概率最高的合法走法
    return candidate_moves[0]

def main():
    board = chess.Board()
    while True:
        line = sys.stdin.readline()
        if not line:
            break
        line = line.strip()
        if not line:
            continue

        if line == "uci":
            print("id name ChessQween3-base")
            print("id author zhoudoe23")
            print("uciok")
            sys.stdout.flush()

        elif line == "isready":
            print("readyok")
            sys.stdout.flush()

        elif line == "ucinewgame":
            board = chess.Board()

        elif line.startswith("position"):
            tokens = line.split()
            if "startpos" in tokens:
                board = chess.Board()
                if "moves" in tokens:
                    move_idx = tokens.index("moves") + 1
                    for m in tokens[move_idx:]:
                        board.push_uci(m)
            elif "fen" in tokens:
                fen_idx = tokens.index("fen") + 1
                if "moves" in tokens:
                    moves_idx = tokens.index("moves")
                    fen_str = " ".join(tokens[fen_idx:moves_idx])
                    board = chess.Board(fen_str)
                    for m in tokens[moves_idx + 1:]:
                        board.push_uci(m)
                else:
                    fen_str = " ".join(tokens[fen_idx:])
                    board = chess.Board(fen_str)

        elif line.startswith("go"):
            best_move = get_best_move(board)
            print(f"bestmove {best_move.uci()}")
            sys.stdout.flush()

        elif line == "quit":
            break

if __name__ == "__main__":
    main()