src/util/pgn_util.py

import io
import logging
import chess.pgn
import tempfile
import os
import time
import glob
import gradio as gr


def extract_game(png_input):
    if isinstance(png_input, io.IOBase):
        file = png_input
    elif isinstance(png_input, str):
        file = io.StringIO(png_input)
    elif isinstance(png_input, gr.utils.NamedString):
        file = png_input.name
    else:
        raise ValueError("Input must be a string or a file-like object.")

    game = read_pgn(file)
    return game


def pgn_string_to_game(pgn_text: str) -> chess.pgn.Game:
    pgn_io = io.StringIO(pgn_text)
    game = chess.pgn.read_game(pgn_io)
    return game


# Load game
def read_pgn(pgn_filepath) -> str:
    with open(pgn_filepath) as f:
        game = chess.pgn.read_game(f)
    return game


def export_pgn(game, tmp_dir="tmp_pgn"):
    os.makedirs(tmp_dir, exist_ok=True)

    with tempfile.NamedTemporaryFile(
        delete=False, suffix=".pgn", dir=tmp_dir, mode="w", encoding="utf-8"
    ) as tmp:
        exporter = chess.pgn.FileExporter(tmp)
        game.accept(exporter)
        return tmp.name


def cleanup_tmp_pgn(tmp_dir="tmp_pgn", expire_seconds=300):
    if not os.path.exists(tmp_dir):
        return

    now = time.time()
    pgn_files = glob.glob(os.path.join(tmp_dir, "*.pgn"))

    for file_path in pgn_files:
        try:
            if now - os.path.getmtime(file_path) > expire_seconds:
                os.remove(file_path)
                logging.info(f"Deleted: {file_path}")
        except Exception as e:
            logging.warning(f"Error deleting {file_path}: {e}")


def format_pv(pv_moves, board):
    temp_board = board.copy()
    san_moves = []
    for move in pv_moves:
        san_moves.append(temp_board.san(move))
        temp_board.push(move)
    return " ".join(san_moves)


def add_variation(node, variation):
    board = node.board()
    current = node

    for move in variation:
        # Ensure move is legal in current position
        if move in board.legal_moves:
            current = current.add_variation(move)
            board.push(move)
        else:
            break  # Stop if PV deviates from legal moves