import chess
import chess.engine
import os

# Path to Stockfish executable
STOCKFISH_PATH = "/home/tun08104/senior_design/data_processing/stockfish/src/stockfish"

# Paths to the two NNUE files
NNUE_1_PATH = "/home/tun08104/senior_design/nnue-pytorch/carlsen35.nnue"
NNUE_2_PATH = "/home/tun08104/senior_design/nnue-pytorch/carlsen500.nnue"

# Names for printed output
NNUE_1_NAME = "Carlsen 35 epochs"
NNUE_2_NAME = "Carlsen 500 epochs"

# Match settings
SEARCH_DEPTH = 12
NUM_GAMES = 100

# Engine settings
HASH_MB = 64
THREADS = 1

# Safety cap so games do not run forever
MAX_PLIES = 300


# ============================================================
# ENGINE SETUP
# ============================================================
def configure_engine(engine, nnue_path):
    """
    Configure one Stockfish engine instance to use a specific NNUE file.
    """
    engine.configure({
        "EvalFile": os.path.abspath(nnue_path),
        "Hash": HASH_MB,
        "Threads": THREADS
    })


def validate_engine(engine, engine_name):
    """
    Test whether the engine is still alive after configuration by asking
    it to analyze the starting position once.

    If this fails, the NNUE file is probably invalid, incompatible,
    or the engine crashed while loading it.
    """
    board = chess.Board()

    try:
        result = engine.play(board, chess.engine.Limit(depth=1))
        if result.move is None:
            raise RuntimeError(f"{engine_name} returned no move during validation.")
    except Exception as e:
        raise RuntimeError(f"{engine_name} failed validation: {e}")


def start_engine(nnue_path, engine_name):
    """
    Start a Stockfish process, load the given NNUE file, and confirm
    that the engine can successfully make a move.
    """
    print(f"Starting {engine_name}...")
    print(f"Using NNUE file: {nnue_path}")

    if not os.path.exists(STOCKFISH_PATH):
        raise FileNotFoundError(f"Stockfish executable not found: {STOCKFISH_PATH}")

    if not os.path.exists(nnue_path):
        raise FileNotFoundError(f"NNUE file not found: {nnue_path}")

    engine = chess.engine.SimpleEngine.popen_uci(STOCKFISH_PATH)

    try:
        configure_engine(engine, nnue_path)
        validate_engine(engine, engine_name)
        print(f"{engine_name} loaded successfully.\n")
        return engine
    except Exception:
        # If configuration or validation fails, try to close the engine
        try:
            engine.quit()
        except Exception:
            pass
        raise


# ============================================================
# PLAY ONE GAME
# ============================================================
def play_one_game(engine1, engine2, game_number):
    """
    Play one game between the two engines.

    Colors alternate every game:
    - even-numbered game: engine1 = White, engine2 = Black
    - odd-numbered game:  engine2 = White, engine1 = Black

    Returns:
        winner_name -> name of the winning net, or None for draw
        result_str  -> '1-0', '0-1', or '1/2-1/2'
    """
    board = chess.Board()

    # Alternate colors for fairness
    if game_number % 2 == 0:
        white_engine = engine1
        black_engine = engine2
        white_name = NNUE_1_NAME
        black_name = NNUE_2_NAME
    else:
        white_engine = engine2
        black_engine = engine1
        white_name = NNUE_2_NAME
        black_name = NNUE_1_NAME

    ply_count = 0

    while not board.is_game_over() and ply_count < MAX_PLIES:
        # Pick the engine whose turn it is
        current_engine = white_engine if board.turn == chess.WHITE else black_engine

        # Ask that engine for the best move at fixed depth
        result = current_engine.play(board, chess.engine.Limit(depth=SEARCH_DEPTH))

        if result.move is None:
            # If for some reason no move is returned, stop the game
            break

        board.push(result.move)
        ply_count += 1

    # If the game hit the move cap without ending, count it as a draw
    if not board.is_game_over():
        return None, "1/2-1/2"

    result_str = board.result()

    if result_str == "1-0":
        return white_name, result_str
    elif result_str == "0-1":
        return black_name, result_str
    else:
        return None, result_str


# ============================================================
# SAFE ENGINE SHUTDOWN
# ============================================================
def safe_quit(engine, name):
    """
    Try to shut down an engine cleanly. If it already crashed,
    do not let that produce another traceback.
    """
    if engine is None:
        return

    try:
        engine.quit()
    except Exception as e:
        print(f"Warning: could not quit {name} cleanly: {e}")


# ============================================================
# MAIN MATCH LOOP
# ============================================================
def run_match():
    engine1 = None
    engine2 = None

    # Win counters
    net1_wins = 0
    net2_wins = 0
    draws = 0

    try:
        # Start and verify both engines before the match begins
        engine1 = start_engine(NNUE_1_PATH, NNUE_1_NAME)
        engine2 = start_engine(NNUE_2_PATH, NNUE_2_NAME)

        for game_number in range(NUM_GAMES):
            winner, result_str = play_one_game(engine1, engine2, game_number)

            if winner == NNUE_1_NAME:
                net1_wins += 1
            elif winner == NNUE_2_NAME:
                net2_wins += 1
            else:
                draws += 1

            print(f"Game {game_number + 1}/{NUM_GAMES}: {result_str}")
            print(f"{NNUE_1_NAME} wins: {net1_wins}")
            print(f"{NNUE_2_NAME} wins: {net2_wins}")
            print(f"Draws: {draws}")
            print("-" * 40)

        print("\n" + "=" * 40)
        print("FINAL RESULTS")
        print("=" * 40)
        print(f"{NNUE_1_NAME} total wins: {net1_wins}")
        print(f"{NNUE_2_NAME} total wins: {net2_wins}")
        print(f"Draws: {draws}")
        print(f"Total games: {NUM_GAMES}")

    except Exception as e:
        print("\nMatch stopped because of an error:")
        print(e)

    finally:
        safe_quit(engine1, NNUE_1_NAME)
        safe_quit(engine2, NNUE_2_NAME)


if __name__ == "__main__":
    run_match()