run.py

#!/usr/bin/env python3
"""Game launcher for SC2 Ultra Bot.

This script sets up the Python path to include sharpy-sc2 and its
bundled python-sc2 fork, then launches a game against the AI.

Usage:
    python run.py                 # Normal game, no IPC
    python run.py --ipc           # Enable stdin/stdout IPC for voice control
    python run.py -r protoss      # Play as Protoss instead of Zerg
    python run.py --vs zerg       # Set opponent race
    python run.py --attach        # Attach to already-running SC2 (for online play)
"""

import argparse
import asyncio
import os
import sys

# Add sharpy-sc2 and its python-sc2 fork to the path
# Must be done before importing any sc2 or sharpy modules
script_dir = os.path.dirname(os.path.abspath(__file__))
sharpy_path = os.path.join(script_dir, "sharpy-sc2")
python_sc2_path = os.path.join(sharpy_path, "python-sc2")

sys.path.insert(0, sharpy_path)
sys.path.insert(0, python_sc2_path)

import aiohttp

from sc2 import maps
from sc2.client import Client
from sc2.data import Race, Difficulty
from sc2.main import run_game
from sc2.player import Bot, Computer, Human
from sc2.protocol import ConnectionAlreadyClosedError
import sc2.main


RACE_MAP = {
    "zerg": Race.Zerg,
    "protoss": Race.Protoss,
    "terran": Race.Terran,
    "random": Race.Random,
}

DIFFICULTY_MAP = {
    "easy": Difficulty.Easy,
    "medium": Difficulty.Medium,
    "hard": Difficulty.Hard,
    "harder": Difficulty.Harder,
    "veryhard": Difficulty.VeryHard,
    "elite": Difficulty.CheatVision,
}


def get_bot_for_race(race: str, ipc_enabled: bool):
    """Import and instantiate the appropriate bot for the selected race."""
    if race == "zerg":
        from bot.zerg_bot import UltraBot
        return Race.Zerg, UltraBot(ipc_enabled=ipc_enabled)
    elif race == "protoss":
        from bot.protoss_bot import ProtossBot
        return Race.Protoss, ProtossBot(ipc_enabled=ipc_enabled)
    elif race == "terran":
        from bot.terran_bot import TerranBot
        return Race.Terran, TerranBot(ipc_enabled=ipc_enabled)
    else:
        raise ValueError(f"Unknown race: {race}")


async def attach_to_game(bot, port: int, realtime: bool):
    """Attach to an already-running SC2 instance via WebSocket.

    Used for online play where SC2 is launched separately with Battle.net.
    """
    ws_url = f"ws://127.0.0.1:{port}/sc2api"
    print(f"Connecting to SC2 at {ws_url}...")

    try:
        session = aiohttp.ClientSession()
        ws = await session.ws_connect(ws_url, timeout=120)
        client = Client(ws)

        print("Connected! Waiting for game to start...")
        print("Create or join a game in SC2, then the bot will take over.")
        print("Use in-game chat to send commands.")
        print()

        # Play the game using the internal _play_game function
        result = await sc2.main._play_game(bot, client, realtime, portconfig=None)

        print(f"Game ended: {result}")
        return result

    except ConnectionAlreadyClosedError:
        print("Connection closed unexpectedly.")
        return None
    except aiohttp.ClientError as e:
        print(f"Could not connect to SC2: {e}")
        print(f"Make sure SC2 is running with: python launch_sc2.py --port {port}")
        return None
    finally:
        await session.close()


def main():
    """Launch a game with UltraBot vs Computer."""
    parser = argparse.ArgumentParser(description="SC2 Ultra Bot launcher")
    parser.add_argument(
        "--ipc",
        action="store_true",
        help="Enable IPC mode (stdin commands, stdout state)"
    )
    parser.add_argument(
        "-r", "--race",
        choices=["zerg", "protoss", "terran"],
        default="zerg",
        help="Race to play as (default: zerg)"
    )
    parser.add_argument(
        "--vs",
        choices=["zerg", "protoss", "terran", "random"],
        default="protoss",
        help="Opponent race (default: protoss)"
    )
    parser.add_argument(
        "--difficulty",
        choices=["easy", "medium", "hard", "harder", "veryhard", "elite"],
        default="hard",
        help="AI difficulty (default: hard)"
    )
    parser.add_argument(
        "--map",
        default="PylonAIE_v4",
        help="Map name (default: PylonAIE_v4)"
    )
    parser.add_argument(
        "--realtime",
        action="store_true",
        default=True,
        help="Run in realtime mode (default: True)"
    )
    parser.add_argument(
        "--no-realtime",
        action="store_false",
        dest="realtime",
        help="Run in fastest mode"
    )
    parser.add_argument(
        "--human",
        action="store_true",
        help="Play against a human (local multiplayer)"
    )
    parser.add_argument(
        "--attach",
        action="store_true",
        help="Attach to already-running SC2 (for online play via Battle.net)"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=5162,
        help="Port for SC2 WebSocket API when using --attach (default: 5162)"
    )

    args = parser.parse_args()

    # Get the appropriate bot for the selected race
    play_race, bot_instance = get_bot_for_race(args.race, args.ipc)
    print(f"Playing as {args.race.capitalize()}")

    # Attach mode: connect to already-running SC2
    if args.attach:
        bot = Bot(play_race, bot_instance)
        asyncio.run(attach_to_game(bot, args.port, args.realtime))
        return

    # Normal mode: launch SC2 and play
    opponent_race = RACE_MAP[args.vs]
    difficulty = DIFFICULTY_MAP[args.difficulty]

    # Set up opponent
    if args.human:
        opponent = Human(opponent_race)
        print(f"Starting game vs Human ({args.vs})")
        print("The human player will control the opponent race normally.")
        print("You control the bot via chat commands.")
    else:
        opponent = Computer(opponent_race, difficulty)

    try:
        run_game(
            maps.get(args.map),
            [
                Bot(play_race, bot_instance),
                opponent
            ],
            realtime=args.realtime
        )
    except (ConnectionAlreadyClosedError, KeyboardInterrupt):
        # Normal exit - game closed or Ctrl+C
        pass
    except Exception as e:
        # Unexpected error
        print(f"Game ended with error: {e}")


if __name__ == "__main__":
    main()