"""Example: multi-agent self-play with BlokusMultiAgentEnv."""
import numpy as np

from blokus_gym import BlokusMultiAgentEnv, GreedyBot, RandomBot


def main():
    env = BlokusMultiAgentEnv(num_players=2, board_size=7)

    # Create bots for each player
    bots = {
        "player_0": RandomBot(player_idx=0, seed=42),
        "player_1": GreedyBot(player_idx=1),
    }

    obs, info = env.reset(seed=42)
    print("=== Multi-Agent Game ===")
    print(f"Agents: {env.agents}")

    steps = 0
    while True:
        for agent in env.agents:
            obs, info = env.last() if hasattr(env, "last") else (obs.get(agent, {}), {})
            mask = env.get_action_mask(agent)
            valid_actions = np.where(mask)[0]

            if len(valid_actions) > 0:
                bot = bots[agent]
                action = bot.select_action(env.game, mask)
                if action is not None:
                    obs, rewards, terminations, truncations, infos = env.step(action)
                    steps += 1

                    if terminations[agent]:
                        print(f"Game over after {steps} steps")
                        scores = env.game.get_scores()
                        for i, score in enumerate(scores):
                            print(f"  Player {i}: {score}")
                        env.close()
                        return

    env.close()


if __name__ == "__main__":
    main()