Lesson 24: Real-time Strategy

For to-do 1, think "check first, then spend": return False before touching self.gold. For to-do 3, a diagonal step (dx, dy) passes two side tiles: (cur[0] + dx, cur[1]) and (cur[0], cur[1] + dy). For to-do 6, the bonus comes from the unit's faction, not the unit, so it changes for everyone at once.

If your instructor hands you the lab file, you will see a few extra lines marked lab runtime near the top and and frame_budget() in the loop. They let the instructor's checker run the program automatically; when you run it yourself they do nothing.

"""Tiny RTS: Advanced Lesson 24 practice exercise (solution).

Two factions, each with its own gold, supply and research, on a generated
tile map. Workers gather gold, bases train workers, barracks train soldiers,
farms raise supply, and an AI opponent plays by the same rules with its own
purse. Units path with octile A* (no corner cutting) and wait before retrying
an unreachable goal. Every timer counts seconds.

Mouse: left-drag to select your units, right-click to move, gather (gold) or
attack (red things). Use the buttons on the right to train, build and
research. Esc cancels placing a building. Close the window to quit.
"""
import heapq
import itertools
import math
import random
from collections import defaultdict, deque
from dataclasses import dataclass
from enum import Enum, auto

import pygame


TILE = 30
COLS, ROWS = 26, 20
MAP_W, MAP_H = COLS * TILE, ROWS * TILE
PANEL_W = 180
WIDTH, HEIGHT = MAP_W + PANEL_W, MAP_H
SQRT2 = math.sqrt(2)
GATHER_TIME = 2.0        # seconds of mining per trip
GATHER_AMOUNT = 10       # gold per trip
RETRY_DELAY = 1.0        # seconds before retrying a path that failed
REPATH_DELAY = 0.5       # seconds between re-plans while chasing a moving target
AI_THINK = 1.5           # seconds between AI decisions
BASE_SITES = [(3, 9), (21, 8)]                       # top-left tiles of the two 2x2 bases
GOLD_SITES = [(3, 6), (4, 6), (21, 5), (22, 5), (12, 10), (13, 10)]

TERRAIN_COLORS = {"grass": (70, 120, 60), "water": (50, 90, 160), "rock": (110, 105, 100),
                  "gold": (220, 185, 60)}
TEXT = (235, 235, 235)
DIM = (130, 130, 130)


@dataclass(frozen=True)
class UnitSpec:
    cost: int
    build_time: float    # seconds
    hp: int
    attack: int
    defense: int
    range: float         # px
    speed: float         # px/s
    cooldown: float      # seconds between attacks


@dataclass(frozen=True)
class BuildingSpec:
    cost: int
    build_time: float    # seconds
    size: int            # tiles per side
    hp: int
    defense: int
    supply: int = 0
    trains: tuple = ()


@dataclass(frozen=True)
class Research:
    cost: int
    time: float
    attack_bonus: int
    applies_to: tuple


UNIT_SPECS = {
    "worker": UnitSpec(cost=50, build_time=3.0, hp=40, attack=4, defense=0, range=22, speed=80, cooldown=1.0),
    "soldier": UnitSpec(cost=80, build_time=4.0, hp=90, attack=12, defense=2, range=26, speed=70, cooldown=0.8),
}
BUILDING_SPECS = {
    "base": BuildingSpec(cost=0, build_time=0.0, size=2, hp=600, defense=5, supply=5, trains=("worker",)),
    "barracks": BuildingSpec(cost=150, build_time=8.0, size=2, hp=400, defense=3, trains=("soldier",)),
    "farm": BuildingSpec(cost=60, build_time=5.0, size=1, hp=150, defense=1, supply=5),
}
RESEARCH = {"weapons": Research(cost=100, time=6.0, attack_bonus=3, applies_to=("soldier",))}


class Task(Enum):
    IDLE = auto()
    MOVE = auto()
    GATHER = auto()
    RETURN = auto()
    ATTACK = auto()


def tile_center(tile):
    return pygame.Vector2(tile[0] * TILE + TILE / 2, tile[1] * TILE + TILE / 2)


def octile(a, b):
    """Exact cost on an open 8-way grid (straight 1, diagonal sqrt 2): an admissible heuristic."""
    dx, dy = abs(a[0] - b[0]), abs(a[1] - b[1])
    return max(dx, dy) + (SQRT2 - 1) * min(dx, dy)


NEIGHBORS = [(1, 0, 1.0), (-1, 0, 1.0), (0, 1, 1.0), (0, -1, 1.0),
             (1, 1, SQRT2), (1, -1, SQRT2), (-1, 1, SQRT2), (-1, -1, SQRT2)]


def astar(passable, start, goals):
    """Octile A* from start to the nearest of several goal tiles.

    Diagonal steps may not cut a blocked corner. Returns the list of tiles to
    walk (start excluded), [] if start is already a goal, or None if unreachable.
    """
    goals = {g for g in goals if passable(g)}
    if not goals:
        return None
    if start in goals:
        return []

    def h(t):
        return min(octile(t, g) for g in goals)

    tie = itertools.count()                       # tie-breaker: tiles are never compared
    open_heap = [(h(start), next(tie), start)]
    g_cost = {start: 0.0}
    came_from = {}
    closed = set()
    while open_heap:
        _, _, cur = heapq.heappop(open_heap)
        if cur in closed:
            continue
        if cur in goals:
            path = [cur]
            while path[-1] in came_from and came_from[path[-1]] != start:
                path.append(came_from[path[-1]])
            return path[::-1]
        closed.add(cur)
        for dx, dy, step in NEIGHBORS:
            nxt = (cur[0] + dx, cur[1] + dy)
            if not passable(nxt):
                continue
            if dx and dy and not (passable((cur[0] + dx, cur[1])) and passable((cur[0], cur[1] + dy))):
                continue                          # no squeezing between two blocked tiles
            new_g = g_cost[cur] + step
            if new_g < g_cost.get(nxt, math.inf):
                g_cost[nxt] = new_g
                came_from[nxt] = cur
                heapq.heappush(open_heap, (new_g + h(nxt), next(tie), nxt))
    return None


class EventBus:
    """Publish/subscribe, as in the Event Systems lesson (minimal version)."""
    def __init__(self):
        self.handlers = defaultdict(list)

    def subscribe(self, name, handler):
        self.handlers[name].append(handler)

    def publish(self, name, **data):
        for handler in list(self.handlers[name]):
            handler(**data)


class Faction:
    def __init__(self, name, color, gold=200, is_ai=False):
        self.name = name
        self.color = color
        self.gold = gold
        self.is_ai = is_ai
        self.research = set()             # finished research names
        self.researching = None           # (name, seconds left) or None

    def pay(self, cost):
        """Spend this faction's own gold. Returns False (and spends nothing) if it can't afford it."""
        if self.gold < cost:
            return False
        self.gold -= cost
        return True

    def attack_bonus(self, kind):
        return sum(RESEARCH[r].attack_bonus for r in self.research if kind in RESEARCH[r].applies_to)


class Unit:
    def __init__(self, kind, faction, pos):
        self.kind = kind
        self.spec = UNIT_SPECS[kind]
        self.faction = faction
        self.pos = pygame.Vector2(pos)
        self.hp = self.spec.hp
        self.task = Task.IDLE
        self.target = None                # a gold tile, a Unit or a Building
        self.path = []
        self.cooldown = 0.0
        self.work = 0.0
        self.carrying = 0
        self.retry = 0.0                  # >0: a path just failed; wait before searching again
        self.repath = 0.0

    @property
    def attack(self):
        """Derived at use time, so research reaches units trained before AND after it."""
        return self.spec.attack + self.faction.attack_bonus(self.kind)

    @property
    def defense(self):
        return self.spec.defense

    @property
    def tile(self):
        return int(self.pos.x // TILE), int(self.pos.y // TILE)


class Building:
    def __init__(self, kind, faction, tile, built=False):
        self.kind = kind
        self.spec = BUILDING_SPECS[kind]
        self.faction = faction
        self.tile = tile
        self.hp = self.spec.hp
        self.progress = self.spec.build_time if built else 0.0   # seconds of construction done
        self.queue = []                   # unit kinds waiting to be trained
        self.train_time = 0.0

    @property
    def complete(self):
        return self.progress >= self.spec.build_time

    @property
    def defense(self):
        return self.spec.defense

    def tiles(self):
        x, y = self.tile
        return [(x + i, y + j) for j in range(self.spec.size) for i in range(self.spec.size)]

    def rect(self):
        return pygame.FRect(self.tile[0] * TILE, self.tile[1] * TILE,
                            self.spec.size * TILE, self.spec.size * TILE)


def damage(attacker, target):
    """Defense soaks damage, but every hit does at least 1."""
    return max(1, attacker.attack - target.defense)


def distance_to(unit, target):
    if isinstance(target, Building):
        r = target.rect()
        closest = pygame.Vector2(max(r.left, min(unit.pos.x, r.right)), max(r.top, min(unit.pos.y, r.bottom)))
        return unit.pos.distance_to(closest)
    return unit.pos.distance_to(target.pos)


def generate_map(seed):
    """Random lakes and rocks, then clear both base areas so no base ever starts on water."""
    rng = random.Random(seed)
    terrain = [["grass"] * COLS for _ in range(ROWS)]
    for _ in range(7):
        kind = rng.choice(["water", "water", "rock"])
        cx, cy, r = rng.randrange(COLS), rng.randrange(ROWS), rng.randint(1, 3)
        for y in range(cy - r, cy + r + 1):
            for x in range(cx - r, cx + r + 1):
                if 0 <= x < COLS and 0 <= y < ROWS and (x - cx) ** 2 + (y - cy) ** 2 <= r * r + 1:
                    terrain[y][x] = kind
    for bx, by in BASE_SITES:
        for y in range(by - 3, by + 5):
            for x in range(bx - 3, bx + 5):
                if 0 <= x < COLS and 0 <= y < ROWS:
                    terrain[y][x] = "grass"
    (ax, ay), (bx, by) = BASE_SITES
    for i in range(61):                   # carve a corridor so the bases can always reach each other
        x = round(ax + 2 + (bx - 1 - ax - 2) * i / 60)
        y = round(ay + (by - ay) * i / 60)
        terrain[y][x] = "grass"
        terrain[y + 1][x] = "grass"
    for gx, gy in GOLD_SITES:
        terrain[gy][gx] = "gold"
    return terrain


class Game:
    def __init__(self, seed=7):
        self.terrain = generate_map(seed)
        self.occupied = {}                # tile -> Building
        self.units = []
        self.buildings = []
        self.bus = EventBus()
        self.messages = deque(maxlen=4)
        self.winner = None
        self.player = Faction("Blue", (80, 150, 255))
        self.enemy = Faction("Red", (230, 80, 80), is_ai=True)
        self.factions = [self.player, self.enemy]
        self.ai = AIPlayer(self.enemy)
        self.bus.subscribe("notice", self.on_notice)
        for faction, site in zip(self.factions, BASE_SITES):
            base = self.add_building("base", faction, site, built=True)
            for _ in range(2):
                self.spawn("worker", base)

    # --- map queries -------------------------------------------------------
    def in_bounds(self, t):
        return 0 <= t[0] < COLS and 0 <= t[1] < ROWS

    def passable(self, t):
        return self.in_bounds(t) and self.terrain[t[1]][t[0]] == "grass" and t not in self.occupied

    def neighbors_of(self, tiles):
        tiles = set(tiles)
        around = {(x + dx, y + dy) for x, y in tiles for dx in (-1, 0, 1) for dy in (-1, 0, 1)}
        return [t for t in around - tiles if self.passable(t)]

    def can_place(self, kind, tile):
        size = BUILDING_SPECS[kind].size
        return all(self.passable((tile[0] + i, tile[1] + j)) for i in range(size) for j in range(size))

    # --- economy -------------------------------------------------------------
    def supply_cap(self, faction):
        return sum(b.spec.supply for b in self.buildings if b.faction is faction and b.complete)

    def supply_used(self, faction):
        queued = sum(len(b.queue) for b in self.buildings if b.faction is faction)
        return queued + sum(1 for u in self.units if u.faction is faction)

    def base_of(self, faction):
        return next((b for b in self.buildings if b.faction is faction and b.kind == "base"), None)

    def notice(self, faction, text):
        self.bus.publish("notice", faction=faction, text=text)

    def on_notice(self, faction, text):
        if faction is self.player:
            self.messages.append(text)

    def add_building(self, kind, faction, tile, built=False):
        b = Building(kind, faction, tile, built)
        self.buildings.append(b)
        for t in b.tiles():
            self.occupied[t] = b
        return b

    def try_build(self, faction, kind, tile):
        if tile is None or not self.can_place(kind, tile):
            self.notice(faction, "Can't build there")
            return None
        if not faction.pay(BUILDING_SPECS[kind].cost):
            self.notice(faction, "Not enough gold")
            return None
        return self.add_building(kind, faction, tile)

    def train(self, building, kind):
        if building is None or not building.complete or kind not in building.spec.trains:
            return False
        f = building.faction
        if self.supply_used(f) >= self.supply_cap(f):
            self.notice(f, "Need more farms")
            return False
        if not f.pay(UNIT_SPECS[kind].cost):
            self.notice(f, "Not enough gold")
            return False
        building.queue.append(kind)
        return True

    def start_research(self, faction, name):
        barracks = [b for b in self.buildings if b.faction is faction and b.kind == "barracks" and b.complete]
        if not barracks or name in faction.research or faction.researching:
            return False
        if not faction.pay(RESEARCH[name].cost):
            self.notice(faction, "Not enough gold")
            return False
        faction.researching = (name, RESEARCH[name].time)
        return True

    def spawn(self, kind, building):
        spots = self.neighbors_of(building.tiles())
        if not spots:
            return None
        spot = min(spots, key=lambda t: (t[1], t[0]))
        unit = Unit(kind, building.faction, tile_center(spot))
        self.units.append(unit)
        self.notice(building.faction, f"{kind.title()} ready")
        return unit

    # --- orders --------------------------------------------------------------
    def order_move(self, unit, tile):
        unit.task, unit.target, unit.path, unit.retry = Task.MOVE, tile, [], 0.0

    def order_gather(self, unit, tile=None):
        if unit.kind != "worker":
            return
        if tile is None:
            golds = [(x, y) for y in range(ROWS) for x in range(COLS) if self.terrain[y][x] == "gold"]
            tile = min(golds, key=lambda g: unit.pos.distance_to(tile_center(g)))
        unit.task, unit.target, unit.path, unit.retry = Task.GATHER, tile, [], 0.0

    def order_attack(self, unit, target):
        unit.task, unit.target, unit.path, unit.retry = Task.ATTACK, target, [], 0.0

    def route(self, unit, goals):
        """Plan a path, but after a failure wait RETRY_DELAY seconds instead of searching every frame."""
        if unit.retry > 0:
            return
        path = astar(self.passable, unit.tile, goals)
        if path is None:
            unit.retry = RETRY_DELAY
        else:
            unit.path = path
            unit.repath = REPATH_DELAY

    def follow(self, unit, dt):
        """Walk along unit.path. Returns True when there is nowhere left to walk."""
        step = unit.spec.speed * dt
        while unit.path and step > 0:
            goal = tile_center(unit.path[0])
            gap = unit.pos.distance_to(goal)
            if gap <= step:
                unit.pos.update(goal)
                unit.path.pop(0)
                step -= gap
            else:
                unit.pos.move_towards_ip(goal, step)
                step = 0
        return not unit.path

    # --- simulation ------------------------------------------------------------
    def update_unit(self, u, dt):
        u.cooldown = max(0.0, u.cooldown - dt)
        u.retry = max(0.0, u.retry - dt)
        u.repath = max(0.0, u.repath - dt)
        if u.task is Task.MOVE:
            if u.path:
                self.follow(u, dt)
            elif u.tile == u.target:
                u.task = Task.IDLE
            else:
                self.route(u, [u.target] if self.passable(u.target) else self.neighbors_of([u.target]))
                if u.retry > 0:
                    u.task = Task.IDLE            # unreachable: give up the move order
        elif u.task is Task.GATHER:
            if u.carrying:
                u.task, u.path = Task.RETURN, []
            elif u.path:
                self.follow(u, dt)
            elif u.pos.distance_to(tile_center(u.target)) <= TILE * 1.5:
                u.work += dt
                if u.work >= GATHER_TIME:
                    u.work, u.carrying = 0.0, GATHER_AMOUNT
            else:
                self.route(u, self.neighbors_of([u.target]))
        elif u.task is Task.RETURN:
            base = self.base_of(u.faction)
            if base is None:
                u.task = Task.IDLE
            elif distance_to(u, base) <= TILE * 0.9:
                u.faction.gold += u.carrying      # the worker's OWN faction gets the gold
                u.carrying = 0
                u.task, u.path = Task.GATHER, []
            elif u.path:
                self.follow(u, dt)
            else:
                self.route(u, self.neighbors_of(base.tiles()))
        elif u.task is Task.ATTACK:
            t = u.target
            if t.hp <= 0:
                u.task, u.target, u.path = Task.IDLE, None, []
            elif distance_to(u, t) <= u.spec.range + (0 if isinstance(t, Building) else 8):
                u.path = []
                if u.cooldown == 0:
                    t.hp -= damage(u, t)
                    u.cooldown = u.spec.cooldown
            elif u.path and u.repath > 0:
                self.follow(u, dt)
            else:
                tiles = t.tiles() if isinstance(t, Building) else [t.tile]
                goals = self.neighbors_of(tiles) + ([t.tile] if isinstance(t, Unit) else [])
                self.route(u, goals)

    def update(self, dt):
        if self.winner:
            return
        for b in list(self.buildings):
            if not b.complete:
                b.progress = min(b.spec.build_time, b.progress + dt)
                if b.complete:
                    self.notice(b.faction, f"{b.kind.title()} complete")
            elif b.queue:
                b.train_time += dt
                if b.train_time >= UNIT_SPECS[b.queue[0]].build_time:
                    b.train_time = 0.0
                    self.spawn(b.queue.pop(0), b)
        for f in self.factions:
            if f.researching:
                name, left = f.researching
                left -= dt
                if left <= 0:
                    f.research.add(name)
                    f.researching = None
                    self.notice(f, f"Research done: {name}")
                else:
                    f.researching = (name, left)
        for u in self.units:
            self.update_unit(u, dt)
        for u in [u for u in self.units if u.hp <= 0]:
            self.units.remove(u)
        for b in [b for b in self.buildings if b.hp <= 0]:
            self.buildings.remove(b)
            for t in b.tiles():
                self.occupied.pop(t, None)
            if b.kind == "base":
                self.winner = self.enemy if b.faction is self.player else self.player
        self.ai.update(self, dt)

    def entity_at(self, pos):
        for u in self.units:
            if u.pos.distance_to(pos) <= 10:
                return u
        t = (int(pos[0] // TILE), int(pos[1] // TILE))
        return self.occupied.get(t)

    def find_site(self, faction, kind):
        base = self.base_of(faction)
        if base is None:
            return None
        bx, by = base.tile
        for r in range(2, 8):
            for dy in range(-r, r + 1):
                for dx in range(-r, r + 1):
                    t = (bx + dx, by + dy)
                    if max(abs(dx), abs(dy)) == r and self.can_place(kind, t):
                        return t
        return None


class AIPlayer:
    """Decides every AI_THINK seconds, and only ever spends its own faction's gold."""
    def __init__(self, faction):
        self.faction = faction
        self.timer = AI_THINK

    def update(self, game, dt):
        self.timer -= dt
        if self.timer > 0:
            return
        self.timer = AI_THINK
        f = self.faction
        base = game.base_of(f)
        if base is None:
            return
        mine = [u for u in game.units if u.faction is f]
        workers = [u for u in mine if u.kind == "worker"]
        soldiers = [u for u in mine if u.kind == "soldier"]
        barracks = next((b for b in game.buildings if b.faction is f and b.kind == "barracks"), None)
        for w in workers:
            if w.task is Task.IDLE:
                game.order_gather(w)
        farm_building = any(b.faction is f and b.kind == "farm" and not b.complete for b in game.buildings)
        if game.supply_used(f) + 1 >= game.supply_cap(f) and not farm_building:
            if f.gold >= BUILDING_SPECS["farm"].cost:
                game.try_build(f, "farm", game.find_site(f, "farm"))
        elif len(workers) < 5 and not base.queue:
            game.train(base, "worker")
        elif barracks is None:
            if f.gold >= BUILDING_SPECS["barracks"].cost:
                game.try_build(f, "barracks", game.find_site(f, "barracks"))
        elif barracks.complete and not barracks.queue:
            game.train(barracks, "soldier")
        if len(soldiers) >= 5:
            enemies = [e for e in game.units + game.buildings if e.faction is not f]
            for s in soldiers:
                if s.task is not Task.ATTACK and enemies:
                    target = min(enemies, key=lambda e: s.pos.distance_to(
                        e.pos if isinstance(e, Unit) else e.rect().center))
                    game.order_attack(s, target)


BUTTONS = [("train_worker", "Train Worker 50g"), ("train_soldier", "Train Soldier 80g"),
           ("build_farm", "Build Farm 60g"), ("build_barracks", "Barracks 150g"),
           ("research", "Weapons +3 100g")]


def button_rects():
    return {key: pygame.Rect(MAP_W + 10, 150 + i * 44, PANEL_W - 20, 36) for i, (key, _) in enumerate(BUTTONS)}


def press_button(game, key, ui):
    p = game.player
    if key == "train_worker":
        game.train(game.base_of(p), "worker")
    elif key == "train_soldier":
        barracks = next((b for b in game.buildings if b.faction is p and b.kind == "barracks" and b.complete), None)
        if barracks is None:
            game.notice(p, "Build a barracks first")
        else:
            game.train(barracks, "soldier")
    elif key in ("build_farm", "build_barracks"):
        ui["placing"] = key.split("_")[1]
    elif key == "research":
        game.start_research(p, "weapons")


def draw(screen, game, fonts, ui, terrain_surface):
    small, big = fonts
    screen.blit(terrain_surface, (0, 0))
    for b in game.buildings:
        r = b.rect()
        color = b.faction.color if b.complete else tuple(c // 2 for c in b.faction.color)
        pygame.draw.rect(screen, color, r.inflate(-4, -4), border_radius=4)
        label = small.render(b.kind[0].upper(), True, TEXT)
        screen.blit(label, label.get_rect(center=r.center))
        fraction = b.hp / b.spec.hp if b.complete else b.progress / b.spec.build_time
        pygame.draw.rect(screen, (30, 30, 30), (r.x + 2, r.y - 5, r.w - 4, 4))
        pygame.draw.rect(screen, (90, 230, 90) if b.complete else (240, 200, 80),
                         (r.x + 2, r.y - 5, (r.w - 4) * max(0.0, fraction), 4))
    for u in game.units:
        radius = 7 if u.kind == "worker" else 10
        pygame.draw.circle(screen, u.faction.color, u.pos, radius)
        if u.carrying:
            pygame.draw.circle(screen, TERRAIN_COLORS["gold"], u.pos, 3)
        if u in ui["selected"]:
            pygame.draw.circle(screen, (255, 255, 255), u.pos, radius + 3, 1)
        pygame.draw.rect(screen, (30, 30, 30), (u.pos.x - 10, u.pos.y - radius - 6, 20, 3))
        pygame.draw.rect(screen, (90, 230, 90), (u.pos.x - 10, u.pos.y - radius - 6, 20 * u.hp / u.spec.hp, 3))
    if ui["drag"] is not None:
        x0, y0 = ui["drag"]
        x1, y1 = ui["mouse"]
        pygame.draw.rect(screen, (255, 255, 255), (min(x0, x1), min(y0, y1), abs(x1 - x0), abs(y1 - y0)), 1)
    if ui["placing"]:
        mx, my = ui["mouse"]
        tile = (int(mx // TILE), int(my // TILE))
        size = BUILDING_SPECS[ui["placing"]].size
        ok = game.can_place(ui["placing"], tile)
        pygame.draw.rect(screen, (90, 230, 90) if ok else (230, 80, 80),
                         (tile[0] * TILE, tile[1] * TILE, size * TILE, size * TILE), 2)

    pygame.draw.rect(screen, (28, 30, 38), (MAP_W, 0, PANEL_W, HEIGHT))
    p = game.player
    lines = [f"Gold: {p.gold}", f"Supply: {game.supply_used(p)}/{game.supply_cap(p)}",
             "Weapons: done" if "weapons" in p.research else
             (f"Weapons: {p.researching[1]:.1f}s" if p.researching else "Weapons: none"),
             f"Enemy gold: {game.enemy.gold}"]
    for i, text in enumerate(lines):
        screen.blit(big.render(text, True, TEXT), (MAP_W + 10, 10 + i * 26))
    for key, label in BUTTONS:
        rect = button_rects()[key]
        pygame.draw.rect(screen, (60, 64, 80), rect, border_radius=6)
        screen.blit(small.render(label, True, TEXT), (rect.x + 8, rect.y + 10))
    for i, text in enumerate(game.messages):
        screen.blit(small.render(text, True, DIM), (MAP_W + 10, 390 + i * 20))
    if game.winner:
        msg = big.render(f"{game.winner.name} wins!", True, (255, 255, 255))
        screen.blit(msg, msg.get_rect(center=(MAP_W / 2, MAP_H / 2)))


def render_terrain(game):
    surf = pygame.Surface((MAP_W, MAP_H)).convert()
    for y in range(ROWS):
        for x in range(COLS):
            pygame.draw.rect(surf, TERRAIN_COLORS[game.terrain[y][x]], (x * TILE, y * TILE, TILE, TILE))
            pygame.draw.rect(surf, (0, 0, 0), (x * TILE, y * TILE, TILE, TILE), 1)
    return surf


def handle_click(game, ui, event):
    pos = pygame.Vector2(event.pos)
    if event.button == 1 and pos.x >= MAP_W:
        for key, rect in button_rects().items():
            if rect.collidepoint(event.pos):
                press_button(game, key, ui)
    elif event.button == 1 and ui["placing"]:
        tile = (int(pos.x // TILE), int(pos.y // TILE))
        if game.try_build(game.player, ui["placing"], tile):
            ui["placing"] = None
    elif event.button == 1:
        ui["drag"] = tuple(event.pos)
    elif event.button == 3 and ui["placing"]:
        ui["placing"] = None
    elif event.button == 3 and pos.x < MAP_W:
        target = game.entity_at(pos)
        tile = (int(pos.x // TILE), int(pos.y // TILE))
        for u in ui["selected"]:
            if target is not None and target.faction is not game.player:
                game.order_attack(u, target)
            elif game.terrain[tile[1]][tile[0]] == "gold" and u.kind == "worker":
                game.order_gather(u, tile)
            else:
                game.order_move(u, tile)


def finish_drag(game, ui, event):
    x0, y0 = ui["drag"]
    x1, y1 = event.pos
    box = pygame.Rect(min(x0, x1), min(y0, y1), abs(x1 - x0) + 1, abs(y1 - y0) + 1).inflate(12, 12)
    ui["selected"] = [u for u in game.units if u.faction is game.player and box.collidepoint(u.pos)]
    ui["drag"] = None


def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Tiny RTS")
    clock = pygame.time.Clock()
    fonts = (pygame.font.Font(None, 20), pygame.font.Font(None, 26))   # created once
    game = Game(seed=7)
    terrain_surface = render_terrain(game)
    ui = {"selected": [], "drag": None, "placing": None, "mouse": (0, 0)}

    running = True
    while running:
        dt = min(clock.tick(60) / 1000, 0.05)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEMOTION:
                ui["mouse"] = event.pos
            elif event.type == pygame.MOUSEBUTTONDOWN:
                ui["mouse"] = event.pos
                handle_click(game, ui, event)
            elif event.type == pygame.MOUSEBUTTONUP and event.button == 1 and ui["drag"] is not None:
                finish_drag(game, ui, event)
            elif event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
                ui["placing"] = None
        ui["selected"] = [u for u in ui["selected"] if u.hp > 0]
        game.update(dt)
        draw(screen, game, fonts, ui, terrain_surface)
        pygame.display.flip()

    pygame.quit()
    workers = sum(1 for u in game.units if u.faction is game.player and u.kind == "worker")
    print(f"Blue gold: {game.player.gold}  workers: {workers}")
    print("RTS closed cleanly.")


if __name__ == "__main__":
    main()