server.py

import argparse
import chess
import chess.engine
from flask import Flask, request, jsonify
from flask_cors import CORS
import os
import glob
import subprocess

# Configuration & Globals
app = Flask(__name__)
CORS(app)
engines = {}
LC0_PATH = "./lc0/lc0.exe"
WEIGHTS_DIR = "./weights"

# Helper Functions
def get_or_load_engine(model_name):
    """
    Loads the engine for a specific model (weight file).
    model_name should be the filename without extension or a key map.
    """
    global engines, LC0_PATH, WEIGHTS_DIR
    
    if model_name in engines:
        return engines[model_name]
    weight_file = None
    possible_paths = [
        os.path.join(WEIGHTS_DIR, f"{model_name}.pb.gz"),
        os.path.join(WEIGHTS_DIR, f"{model_name}"),
    ]
    
    for p in possible_paths:
        if os.path.exists(p):
            weight_file = p
            break
            
    if not weight_file:
        matches = glob.glob(os.path.join(WEIGHTS_DIR, f"*{model_name}*.pb.gz"))
        if matches:
            weight_file = matches[0]
    if not weight_file:
        print(f"Could not find weights for model: {model_name}")
        return None
    if not os.path.exists(LC0_PATH):
        print(f"LC0 binary not found at: {LC0_PATH}")
        return None
    try:
        print(f"Loading engine for {model_name} with weights: {weight_file}")
        
        # Determine creation flags for headless execution on Windows
        creation_flags = 0
        if os.name == 'nt':
            creation_flags = subprocess.CREATE_NO_WINDOW
            
        engine = chess.engine.SimpleEngine.popen_uci(
            [LC0_PATH, f"--weights={weight_file}"],
            creationflags=creation_flags
        )
        engines[model_name] = engine
        return engine
    except Exception as e:
        print(f"Failed to start engine for {model_name}: {e}")
        return None

# Routes
@app.route('/models', methods=['GET'])
def list_models():
    """
    Returns a list of available models based on files in weights directory.
    """
    try:
        files = glob.glob(os.path.join(WEIGHTS_DIR, "*.pb.gz"))
        
        models = []
        for f in files:
            basename = os.path.basename(f)
            name = basename.replace(".pb.gz", "")
            models.append(name)
        
        models.sort()
        return jsonify(models)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/move', methods=['POST'])
def get_move():
    data = request.json
    fen = data.get('fen')
    model_name = data.get('model', 'maia-1900')
    nodes = data.get('nodes', 1)
    
    if not fen:
        return jsonify({"error": "No FEN provided"}), 400
    engine = get_or_load_engine(model_name)
    if not engine:
        return jsonify({"error": f"Could not load engine for model: {model_name}. Check server logs."}), 500
    try:
        board = chess.Board(fen)
        
        multipv_moves = []
        with engine.analysis(board, chess.engine.Limit(nodes=nodes), multipv=3) as analysis:
            analysis.wait()
            if analysis.multipv:
                for item in analysis.multipv:
                    if item and "pv" in item:
                        multipv_moves.append(item["pv"][0].uci())
        
        if not multipv_moves:
             r = engine.play(board, chess.engine.Limit(nodes=nodes))
             multipv_moves = [r.move.uci()]

        return jsonify({
            "bestmove": multipv_moves[0],
            "all_moves": multipv_moves, 
            "evaluation": 0.0,
            "mate": None,
            "continuation": []
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# Main Function
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Chess Botting API Server')
    parser.add_argument('--port', type=int, default=5000)
    parser.add_argument('--engine', type=str, default='./lc0/lc0.exe', help='Path to lc0 binary')
    parser.add_argument('--weights_dir', type=str, default='./weights', help='Directory containing .pb.gz key files')
    
    args = parser.parse_args()
    LC0_PATH = args.engine
    WEIGHTS_DIR = args.weights_dir
    
    if not os.path.exists(LC0_PATH):
        print(f"WARNING: lc0 binary not found at {LC0_PATH}. Requests will fail.")
    if not os.path.exists(WEIGHTS_DIR):
        print(f"WARNING: Weights directory not found at {WEIGHTS_DIR}. Requests will fail.")
        os.makedirs(WEIGHTS_DIR, exist_ok=True)
        print(f"Created {WEIGHTS_DIR}. Please place your .pb.gz files there.")
    print(f"Starting Multi-Model Server on port {args.port}...")
    print(f"Pointing to Weights Directory: {os.path.abspath(WEIGHTS_DIR)}")
    app.run(port=args.port)