scripts/collect_bot_data.py

#!/usr/bin/env python3
"""Collect compact macro-policy demonstrations against OpenRA's built-in AI.

Supports two Python bots:
  - scripted: PeriodicAttackBot (simple build-order + periodic grid-search attack)
  - normal:   NormalAIBot (reimplements OpenRA's ModularBot@NormalAI in Python)

Usage:
    # Start the OpenRA-RL server first:
    docker run -p 8000:8000 openra-rl

    # Collect a macro dataset with the default scripted bot:
    python scripts/collect_bot_data.py --episodes 10

    # Collect with the normal AI bot (mimics OpenRA's built-in normal AI):
    python scripts/collect_bot_data.py --episodes 10 --bot normal

    # Quick test (2 minutes per episode):
    python scripts/collect_bot_data.py --episodes 2 --max-minutes 2 --bot normal

Output:
    data/macro/macro_dataset.jsonl.gz  - compact macro-policy dataset
    data/episodes/collection_summary.json
    data/episodes/episode_001.orarep
    ...

Notes:
    - The server runs at ~23 steps/second, so 15 min is about 20,000 steps.
    - The ScriptedBot builds a full base by ~1400 steps, then marches to the enemy.
      You need 10+ minutes to see actual combat on a 128x128 map.
    - Episodes stop early if the game ends (win/lose) before the time limit.
    - The collector writes compact macro rows directly; it does not save raw
      per-step trajectory JSON.
"""

import argparse
import asyncio
import json
import shutil
import sys
import time
from collections import Counter
from pathlib import Path
from typing import Any
from urllib import parse as urlparse
from urllib import request as urlrequest

# openra_env is installed via: pip install openra-rl
# scripted_bot.py is vendored in this scripts/ directory
sys.path.insert(0, str(Path(__file__).resolve().parent))

from openra_env.client import OpenRAEnv
from openra_env.mcp_ws_client import OpenRAMCPClient
from openra_env.models import OpenRAAction, OpenRAObservation, CommandModel, ActionType
from scripted_bot import ScriptedBot
from normal_ai_bot import NormalAIBot


class ToolEnabledOpenRAEnv(OpenRAEnv):
    """OpenRA env client that can issue MCP tool calls on the same WS session."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._rpc_id = 0

    async def call_tool(self, name: str, **kwargs: Any) -> Any:
        """Call an MCP tool over the existing environment websocket session."""
        self._rpc_id += 1
        rpc_request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "params": {"name": name, "arguments": kwargs},
            "id": self._rpc_id,
        }
        response = await self._send_and_receive({"type": "mcp", "data": rpc_request})
        rpc_response = response.get("data", {})

        if "error" in rpc_response:
            error = rpc_response["error"]
            raise RuntimeError(f"Tool '{name}' failed: {error.get('message', error)}")

        return OpenRAMCPClient._unwrap_mcp_result(rpc_response.get("result"))


def get_nested(d: dict[str, Any], key: str, default: Any = None) -> Any:
    value = d.get(key, default)
    return value if value is not None else default


def count_types(items: list[dict[str, Any]], key: str = "type", top_k: int = 12) -> dict[str, int]:
    counts = Counter(str(item.get(key, "")) for item in items if item.get(key))
    ordered = sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))[:top_k]
    return {name: count for name, count in ordered}


def render_counts(counts: dict[str, int]) -> str:
    if not counts:
        return "none"
    return ", ".join(f"{count}x {name}" for name, count in counts.items())


def infer_phase(obs: dict[str, Any]) -> str:
    if obs.get("done"):
        return "terminal"

    buildings = {b.get("type") for b in get_nested(obs, "buildings", []) if b.get("type")}
    tick = int(obs.get("tick", 0) or 0)
    visible_enemy_units = len(get_nested(obs, "visible_enemies", []))
    visible_enemy_buildings = len(get_nested(obs, "visible_enemy_buildings", []))
    military = get_nested(obs, "military", {})
    combat_activity = sum(
        int(get_nested(military, field, 0) or 0)
        for field in ("units_killed", "buildings_killed", "units_lost", "buildings_lost")
    )

    if visible_enemy_units or visible_enemy_buildings:
        return "combat"
    if not buildings:
        return "eliminated"
    if tick < 3000 or "fact" not in buildings:
        return "opening"
    if not ({"tent", "barr"} & buildings) or "proc" not in buildings:
        return "opening"
    if tick < 12000 and combat_activity == 0:
        return "build_up"
    if tick >= 24000:
        return "late_game"
    return "mid_game"


def summarize_observation(obs: dict[str, Any], top_k: int) -> dict[str, Any]:
    eco = get_nested(obs, "economy", {})
    units = get_nested(obs, "units", [])
    buildings = get_nested(obs, "buildings", [])
    enemies = get_nested(obs, "visible_enemies", [])
    enemy_buildings = get_nested(obs, "visible_enemy_buildings", [])
    production = get_nested(obs, "production", [])
    available = get_nested(obs, "available_production", [])
    military = get_nested(obs, "military", {})
    map_info = get_nested(obs, "map_info", {})

    idle_units = sum(1 for unit in units if unit.get("is_idle"))
    power_balance = int(get_nested(eco, "power_provided", 0) or 0) - int(get_nested(eco, "power_drained", 0) or 0)

    return {
        "tick": int(obs.get("tick", 0) or 0),
        "phase": infer_phase(obs),
        "done": bool(obs.get("done", False)),
        "result": str(obs.get("result", "") or ""),
        "map": {
            "name": str(get_nested(map_info, "map_name", "") or ""),
            "width": int(get_nested(map_info, "width", 0) or 0),
            "height": int(get_nested(map_info, "height", 0) or 0),
        },
        "economy": {
            "cash": int(get_nested(eco, "cash", 0) or 0),
            "ore": int(get_nested(eco, "ore", 0) or 0),
            "credits": int(get_nested(eco, "cash", 0) or 0) + int(get_nested(eco, "ore", 0) or 0),
            "power_balance": power_balance,
            "harvesters": int(get_nested(eco, "harvester_count", 0) or 0),
        },
        "own": {
            "units": len(units),
            "idle_units": idle_units,
            "unit_types": count_types(units, top_k=top_k),
            "buildings": len(buildings),
            "building_types": count_types(buildings, top_k=top_k),
        },
        "enemy": {
            "visible_units": len(enemies),
            "unit_types": count_types(enemies, top_k=top_k),
            "visible_buildings": len(enemy_buildings),
            "building_types": count_types(enemy_buildings, top_k=top_k),
        },
        "production": [
            {
                "item": str(p.get("item", "") or ""),
                "progress": round(float(p.get("progress", 0.0) or 0.0), 3),
                "remaining_ticks": int(p.get("remaining_ticks", 0) or 0),
            }
            for p in production[:8]
            if p.get("item")
        ],
        "available_production": [str(item) for item in available[:20]],
        "military": {
            key: int(get_nested(military, key, 0) or 0)
            for key in (
                "units_killed",
                "buildings_killed",
                "units_lost",
                "buildings_lost",
                "kills_cost",
                "deaths_cost",
            )
        },
        "explored_percent": float(obs.get("explored_percent", 0.0) or 0.0),
    }


def summarize_command(cmd: dict[str, Any]) -> dict[str, Any]:
    action = str(cmd.get("action", "no_op") or "no_op").lower()
    macro: dict[str, Any] = {"intent": action, "count": 1}

    item_type = cmd.get("item_type")
    if item_type:
        macro["item_type"] = str(item_type)

    target_actor_id = int(cmd.get("target_actor_id", 0) or 0)
    if target_actor_id:
        macro["target_actor_id"] = target_actor_id

    target_x = cmd.get("target_x")
    target_y = cmd.get("target_y")
    if target_x is not None or target_y is not None:
        tx = int(target_x or 0)
        ty = int(target_y or 0)
        if tx or ty:
            macro["target"] = {"x": tx, "y": ty}

    queued = bool(cmd.get("queued", False))
    if queued:
        macro["queued"] = True

    for extra_key in ("stance", "stance_type"):
        extra_value = cmd.get(extra_key)
        if extra_value not in (None, "", 0):
            macro[extra_key] = extra_value

    return macro


def macro_signature(macro: dict[str, Any]) -> str:
    payload = {k: v for k, v in macro.items() if k != "count"}
    return json.dumps(payload, sort_keys=True, separators=(",", ":"))


def merge_macros(macros: list[dict[str, Any]]) -> list[dict[str, Any]]:
    merged: list[dict[str, Any]] = []
    index_by_sig: dict[str, int] = {}

    for macro in macros:
        sig = macro_signature(macro)
        existing_idx = index_by_sig.get(sig)
        if existing_idx is None:
            index_by_sig[sig] = len(merged)
            merged.append(dict(macro))
        else:
            merged[existing_idx]["count"] += int(macro.get("count", 1) or 1)

    build_idx = next((i for i, item in enumerate(merged) if item.get("intent") == "build"), None)
    place_idx = next((i for i, item in enumerate(merged) if item.get("intent") == "place_building"), None)
    if build_idx is not None and place_idx is not None:
        build_macro = merged[build_idx]
        place_macro = merged[place_idx]
        construct_macro: dict[str, Any] = {
            "intent": "construct",
            "count": min(int(build_macro.get("count", 1)), int(place_macro.get("count", 1))),
        }
        if "item_type" in build_macro:
            construct_macro["item_type"] = build_macro["item_type"]
        if "target" in place_macro:
            construct_macro["target"] = place_macro["target"]

        rebuilt: list[dict[str, Any]] = [construct_macro]
        for idx, macro in enumerate(merged):
            if idx not in (build_idx, place_idx):
                rebuilt.append(macro)
        merged = rebuilt

    return merged


def extract_macro_actions(action: dict[str, Any]) -> list[dict[str, Any]]:
    commands = get_nested(action, "commands", [])
    if not commands:
        return [{"intent": "no_op", "count": 1}]
    merged = merge_macros([summarize_command(cmd) for cmd in commands])
    return merged or [{"intent": "no_op", "count": 1}]


def primary_intent(macros: list[dict[str, Any]]) -> str:
    intents = [str(m.get("intent", "no_op")) for m in macros]
    if not intents:
        return "no_op"
    if len(set(intents)) == 1:
        return intents[0]
    if any(intent in {"attack", "attack_move"} for intent in intents):
        return "combat_mixed"
    if any(intent in {"construct", "build", "place_building"} for intent in intents):
        return "base_mixed"
    if any(intent == "train" for intent in intents):
        return "production_mixed"
    return "mixed"


def observation_signature(obs: dict[str, Any]) -> tuple[Any, ...]:
    military = get_nested(obs, "military", {})
    return (
        len(get_nested(obs, "units", [])),
        len(get_nested(obs, "buildings", [])),
        len(get_nested(obs, "visible_enemies", [])),
        len(get_nested(obs, "visible_enemy_buildings", [])),
        len(get_nested(obs, "production", [])),
        int(get_nested(military, "units_killed", 0) or 0),
        int(get_nested(military, "buildings_killed", 0) or 0),
        int(get_nested(military, "units_lost", 0) or 0),
        int(get_nested(military, "buildings_lost", 0) or 0),
    )


def is_wait_only(macros: list[dict[str, Any]]) -> bool:
    return len(macros) == 1 and macros[0].get("intent") == "no_op"


def render_prompt(summary: dict[str, Any]) -> str:
    economy = summary["economy"]
    own = summary["own"]
    enemy = summary["enemy"]
    military = summary["military"]

    lines = [
        "You are planning the next macro action for an OpenRA Red Alert bot.",
        "Return a JSON array of compact macro actions.",
        "",
        f"[tick] {summary['tick']}",
        f"[phase] {summary['phase']}",
        (
            "[economy] "
            f"cash={economy['cash']} ore={economy['ore']} total={economy['credits']} "
            f"power={economy['power_balance']:+d} harvesters={economy['harvesters']}"
        ),
        (
            "[own] "
            f"units={own['units']} idle={own['idle_units']} "
            f"buildings={own['buildings']}"
        ),
        f"[own_units] {render_counts(own['unit_types'])}",
        f"[own_buildings] {render_counts(own['building_types'])}",
        (
            "[enemy] "
            f"visible_units={enemy['visible_units']} "
            f"visible_buildings={enemy['visible_buildings']}"
        ),
        f"[enemy_units] {render_counts(enemy['unit_types'])}",
        f"[enemy_buildings] {render_counts(enemy['building_types'])}",
        (
            "[combat] "
            f"killed={military['units_killed']}u/{military['buildings_killed']}b "
            f"lost={military['units_lost']}u/{military['buildings_lost']}b "
            f"kills_cost={military['kills_cost']} deaths_cost={military['deaths_cost']}"
        ),
        (
            "[production] "
            + (
                ", ".join(
                    f"{item['item']}@{item['progress']:.0%}(~{item['remaining_ticks']}t)"
                    for item in summary["production"]
                )
                if summary["production"]
                else "idle"
            )
        ),
        (
            "[available] "
            + (", ".join(summary["available_production"]) if summary["available_production"] else "none")
        ),
        f"[explored_percent] {summary['explored_percent']:.1f}",
    ]
    return "\n".join(lines)


def should_keep_step(
    step_idx: int,
    step_data: dict[str, Any],
    obs: dict[str, Any],
    macros: list[dict[str, Any]],
    prev_sig: tuple[Any, ...] | None,
    sample_every: int,
    keep_state_changes: bool,
) -> list[str]:
    reasons: list[str] = []

    if step_idx == 0:
        reasons.append("episode_start")
    if step_data.get("done"):
        reasons.append("terminal")
    if sample_every > 0 and step_idx % sample_every == 0:
        reasons.append("periodic")
    if not is_wait_only(macros):
        reasons.append("non_noop")
    if get_nested(obs, "visible_enemies", []) or get_nested(obs, "visible_enemy_buildings", []):
        reasons.append("enemy_visible")
    if abs(float(step_data.get("reward", 0.0) or 0.0)) > 1e-9:
        reasons.append("nonzero_reward")

    if keep_state_changes:
        current_sig = observation_signature(obs)
        if prev_sig is not None and current_sig != prev_sig:
            reasons.append("state_change")

    return reasons


def build_row(
    episode_name: str,
    episode_result: str,
    step_data: dict[str, Any],
    state_summary: dict[str, Any],
    macros: list[dict[str, Any]],
    reasons: list[str],
) -> dict[str, Any]:
    return {
        "id": f"{episode_name}:{int(step_data.get('step', 0) or 0)}",
        "episode": episode_name,
        "step": int(step_data.get("step", 0) or 0),
        "tick": state_summary["tick"],
        "phase": state_summary["phase"],
        "episode_result": episode_result,
        "reward": float(step_data.get("reward", 0.0) or 0.0),
        "done": bool(step_data.get("done", False)),
        "selection_reason": reasons,
        "primary_intent": primary_intent(macros),
        "state": state_summary,
        "macro_actions": macros,
        "prompt": render_prompt(state_summary),
        "completion": json.dumps(macros, separators=(",", ":")),
    }

class PeriodicAttackBot(ScriptedBot):
    """ScriptedBot with grid-search targeting that works on any map layout.

    Instead of assuming rotational symmetry, divides the map into a grid
    and systematically searches cells far from our base for the enemy.
    Cycles to the next grid cell whenever the army fails to find enemies.
    """

    REATTACK_INTERVAL = 600   # re-issue attack order every ~600 ticks (~18s)
    CYCLE_AFTER_REATTACKS = 2  # switch target after this many re-attacks with no enemy contact

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._last_attack_tick: int = 0
        self._candidate_targets: list[tuple[int, int]] = []
        self._target_index: int = 0
        self._no_contact_reattacks: int = 0
        self._cached_map_size: tuple[int, int] | None = None
        self._enemy_base_pos: tuple[int, int] | None = None

    def _handle_rally_points(self, obs):
        """Override: set rally on both Allied ('tent') and Soviet ('barr') barracks."""
        from openra_env.models import CommandModel, ActionType
        commands = []
        cy = self._find_building(obs, "fact")
        if not cy:
            return commands
        for b in obs.buildings:
            if b.type in ("tent", "barr", "weap") and b.actor_id not in self._rally_set:
                rally_x = cy.cell_x if cy.cell_x > 0 else cy.pos_x // 1024
                rally_y = cy.cell_y if cy.cell_y > 0 else cy.pos_y // 1024
                commands.append(CommandModel(
                    action=ActionType.SET_RALLY_POINT,
                    actor_id=b.actor_id,
                    target_x=rally_x,
                    target_y=rally_y,
                ))
                self._rally_set.add(b.actor_id)
        return commands

    def _get_map_size(self, obs: OpenRAObservation) -> tuple[int, int]:
        """Return map dimensions, preferring the in-game reported size.

        The reset observation may report padded dimensions (e.g. 128x128)
        while gameplay observations report the actual playable area (e.g.
        112x54). Always update the cache when a smaller (more accurate)
        value is observed.
        """
        w, h = obs.map_info.width, obs.map_info.height
        if w > 0 and h > 0:
            if self._cached_map_size is None:
                self._cached_map_size = (w, h)
            else:
                cw, ch = self._cached_map_size
                if w < cw or h < ch:
                    self._cached_map_size = (w, h)
                    self._candidate_targets = []
        if self._cached_map_size is not None:
            return self._cached_map_size
        return (128, 128)

    def _compute_candidate_spawns(self, obs: OpenRAObservation) -> list[tuple[int, int]]:
        """Generate search targets by dividing the map into a grid.

        Produces a 4x4 grid of cell centers, excludes the cell containing
        our base, and sorts by distance (farthest first) so we check the
        most likely enemy positions before nearby ones.
        """
        cy_bldg = self._find_building(obs, "fact")
        w, h = self._get_map_size(obs)

        if not cy_bldg:
            return [(w // 2, h // 2)]

        bx = cy_bldg.cell_x if cy_bldg.cell_x > 0 else cy_bldg.pos_x // 1024
        by = cy_bldg.cell_y if cy_bldg.cell_y > 0 else cy_bldg.pos_y // 1024

        grid_n = 3
        cell_w, cell_h = w // grid_n, h // grid_n
        grid_centers = []
        for gx in range(grid_n):
            for gy in range(grid_n):
                cx = cell_w * gx + cell_w // 2
                cy = cell_h * gy + cell_h // 2
                cx = max(0, min(w - 1, cx))
                cy = max(0, min(h - 1, cy))
                grid_centers.append((cx, cy))

        min_dist_sq = (min(w, h) // grid_n) ** 2
        candidates = [
            p for p in grid_centers
            if (p[0] - bx) ** 2 + (p[1] - by) ** 2 > min_dist_sq
        ]

        if not candidates:
            candidates = [(w // 2, h // 2)]

        candidates.sort(
            key=lambda p: (p[0] - bx) ** 2 + (p[1] - by) ** 2,
            reverse=True,
        )
        return candidates

    def _find_attack_target(self, obs: OpenRAObservation):
        """Priority: visible enemy buildings > enemy units > remembered base > grid search."""
        if obs.visible_enemy_buildings:
            prod_buildings = [
                b for b in obs.visible_enemy_buildings
                if b.type in ("fact", "tent", "weap", "hpad", "afld")
            ]
            target = prod_buildings[0] if prod_buildings else obs.visible_enemy_buildings[0]
            self._enemy_base_pos = (target.cell_x, target.cell_y)
            return target.cell_x, target.cell_y

        if obs.visible_enemies:
            enemy = obs.visible_enemies[0]
            if self._enemy_base_pos is None:
                self._enemy_base_pos = (enemy.cell_x, enemy.cell_y)
            return enemy.cell_x, enemy.cell_y

        if self._enemy_base_pos is not None:
            return self._enemy_base_pos

        if not self._candidate_targets:
            self._candidate_targets = self._compute_candidate_spawns(obs)
            self._target_index = 0

        return self._candidate_targets[self._target_index % len(self._candidate_targets)]

    def _handle_combat(self, obs: OpenRAObservation):
        from openra_env.models import CommandModel, ActionType
        commands = []
        if self.phase != "attack":
            return commands

        commands.extend(self._handle_unload(obs))

        fighters = [
            u for u in obs.units
            if u.type in self.COMBAT_UNIT_TYPES
            and u.actor_id not in self._guards_assigned
        ]

        if len(fighters) < 2:
            return commands

        # Reset counter when we actually see enemies (we're at the right place)
        if obs.visible_enemies or obs.visible_enemy_buildings:
            self._no_contact_reattacks = 0

        ticks_since_attack = obs.tick - self._last_attack_tick
        if ticks_since_attack < self.REATTACK_INTERVAL:
            return commands

        # Each time we re-attack without enemy contact, increment counter.
        # After CYCLE_AFTER_REATTACKS misses, switch to next candidate.
        if not obs.visible_enemies and not obs.visible_enemy_buildings:
            self._no_contact_reattacks += 1
            if (self._no_contact_reattacks >= self.CYCLE_AFTER_REATTACKS
                    and self._candidate_targets):
                self._target_index = (self._target_index + 1) % len(self._candidate_targets)
                self._no_contact_reattacks = 0
                self._log(
                    f"[cycle] No enemy contact after {self.CYCLE_AFTER_REATTACKS} "
                    f"re-attacks, switching to target #{self._target_index}: "
                    f"{self._candidate_targets[self._target_index]}"
                )

        target_x, target_y = self._find_attack_target(obs)
        self._last_attack_tick = obs.tick

        for unit in fighters:
            commands.append(CommandModel(
                action=ActionType.ATTACK_MOVE,
                actor_id=unit.actor_id,
                target_x=target_x,
                target_y=target_y,
            ))

        self._log(
            f"[periodic] Attack-move {len(fighters)} units → "
            f"({target_x}, {target_y}) at tick {obs.tick}"
        )
        return commands


def serialize_obs(obs: OpenRAObservation) -> dict:
    """Convert observation to a JSON-serializable dict.

    Drops the spatial_map field (large binary) to keep files manageable.
    """
    d = obs.model_dump()
    # Remove large spatial tensor — not needed for text-based imitation learning
    d.pop("spatial_map", None)
    d.pop("metadata", None)
    return d


def serialize_action(action: OpenRAAction) -> dict:
    """Convert action to a JSON-serializable dict."""
    return action.model_dump()


def available_credits(obs: OpenRAObservation) -> int:
    """OpenRA spendable money = liquid cash + stored ore/resources."""
    return obs.economy.cash + obs.economy.ore


def build_artifact_url(env_url: str, route: str, query: dict[str, Any] | None = None) -> str:
    """Build an HTTP URL for the repo's artifact helper endpoints."""
    base = env_url.rstrip("/")
    route = route if route.startswith("/") else f"/{route}"
    if not query:
        return f"{base}{route}"
    encoded = urlparse.urlencode({k: v for k, v in query.items() if v is not None})
    return f"{base}{route}?{encoded}" if encoded else f"{base}{route}"


def _read_json_response(url: str, timeout: float = 10.0) -> dict[str, Any]:
    with urlrequest.urlopen(url, timeout=timeout) as response:
        body = response.read().decode("utf-8")
    return json.loads(body) if body else {}


def _post_json_response(url: str, payload: dict[str, Any] | None = None, timeout: float = 30.0) -> dict[str, Any]:
    data = json.dumps(payload or {}).encode("utf-8")
    request = urlrequest.Request(
        url,
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urlrequest.urlopen(request, timeout=timeout) as response:
        body = response.read().decode("utf-8")
    return json.loads(body) if body else {}


def resolve_server_urls(env_url: str) -> dict[str, str]:
    """Resolve wrapper/root URLs into the actual game and artifact endpoints.

    For the Hugging Face wrapper app:
      - OpenRA env lives under `/openra`
      - replay/log artifact helpers live on the wrapper root
    """
    base = env_url.rstrip("/")

    if base.endswith("/openra"):
        artifact_url = base[:-len("/openra")] or base
        return {
            "game_url": base,
            "artifact_url": artifact_url,
        }

    status_url = build_artifact_url(base, "/openra-status")
    try:
        status = _read_json_response(status_url)
    except Exception:
        return {
            "game_url": base,
            "artifact_url": base,
        }

    mount_path = str(status.get("mount_path") or "/openra").rstrip("/") or "/openra"
    if not mount_path.startswith("/"):
        mount_path = f"/{mount_path}"

    if not status.get("mounted"):
        _post_json_response(build_artifact_url(base, "/mount-openra"))

    return {
        "game_url": f"{base}{mount_path}",
        "artifact_url": base,
    }


def download_remote_replay(
    env_url: str,
    replay_path: str,
    output_dir: Path,
    episode_id: int,
) -> dict[str, Any]:
    """Download a replay from a remote OpenRA server artifact endpoint."""
    source = Path(replay_path)
    suffix = source.suffix or ".orarep"
    destination = output_dir / f"episode_{episode_id:03d}{suffix}"
    download_url = build_artifact_url(
        env_url,
        "/artifacts/replay",
        {"path": replay_path, "delete_after_download": "false"},
    )

    output_dir.mkdir(parents=True, exist_ok=True)
    with urlrequest.urlopen(download_url, timeout=120) as response, open(destination, "wb") as out_file:
        shutil.copyfileobj(response, out_file)

    return {
        "local_copy": str(destination),
        "download_url": download_url,
    }


def cleanup_remote_artifacts(
    env_url: str,
    replay_paths: list[str] | None = None,
    delete_logs: bool = True,
) -> dict[str, Any]:
    """Delete remote replays/log files from the OpenRA server after download."""
    payload = {
        "replay_paths": replay_paths or [],
        "delete_logs": delete_logs,
    }
    request = urlrequest.Request(
        build_artifact_url(env_url, "/artifacts/cleanup"),
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urlrequest.urlopen(request, timeout=30) as response:
        body = response.read().decode("utf-8")
    return json.loads(body) if body else {}


def copy_replay_artifact(replay_info: dict, output_dir: Path, episode_id: int, env_url: str) -> dict:
    """Copy or download the replay, then clean remote artifacts when applicable."""
    if not replay_info:
        try:
            cleanup_result = cleanup_remote_artifacts(env_url=env_url, replay_paths=[], delete_logs=True)
            return {"remote_cleanup": cleanup_result} if cleanup_result else {}
        except Exception as exc:
            return {"cleanup_error": f"{type(exc).__name__}: {exc}"}

    enriched = dict(replay_info)
    replay_path = str(replay_info.get("path", "") or "")
    source = Path(replay_path) if replay_path else None
    replay_downloaded = False

    if replay_path and source is not None and source.is_file():
        destination = output_dir / f"episode_{episode_id:03d}{source.suffix}"
        if source.resolve() != destination.resolve():
            shutil.copy2(source, destination)
        enriched["local_copy"] = str(destination)
        return enriched

    if replay_path:
        try:
            enriched.update(download_remote_replay(env_url, replay_path, output_dir, episode_id))
            replay_downloaded = True
        except Exception as exc:
            enriched["download_error"] = f"{type(exc).__name__}: {exc}"

    try:
        cleanup_result = cleanup_remote_artifacts(
            env_url=env_url,
            replay_paths=[replay_path] if replay_downloaded and replay_path else [],
            delete_logs=True,
        )
        if cleanup_result:
            enriched["remote_cleanup"] = cleanup_result
    except Exception as exc:
        enriched["cleanup_error"] = f"{type(exc).__name__}: {exc}"

    return enriched


def open_dataset_writer(path: Path, append: bool):
    """Open the compact macro dataset writer."""
    path.parent.mkdir(parents=True, exist_ok=True)
    mode = "at" if append else "wt"
    if path.suffix == ".gz":
        import gzip

        return gzip.open(path, mode, encoding="utf-8")
    return open(path, mode, encoding="utf-8")


def infer_outcome(
    final_obs: OpenRAObservation,
    eliminated_since_step: int | None,
    elapsed_s: float,
    max_minutes: float,
) -> str:
    """Classify why an episode stopped when the env did not emit a final result."""
    if final_obs.done and final_obs.result:
        return final_obs.result
    if eliminated_since_step is not None or (not final_obs.units and not final_obs.buildings):
        return "eliminated"
    if elapsed_s >= max_minutes * 60.0:
        return f"time_limit({max_minutes:.0f}min)"
    return "step_limit"


async def wait_for_replay_artifact(
    env: ToolEnabledOpenRAEnv,
    episode_id: int,
    verbose: bool,
    timeout_s: float,
    poll_interval_s: float = 1.0,
) -> dict[str, Any]:
    """Poll for the replay path until OpenRA has flushed it or we time out."""
    deadline = time.time() + max(timeout_s, 0.0)
    attempt = 0
    last_info: dict[str, Any] = {}
    announced_wait = False

    while True:
        attempt += 1
        try:
            replay_info = await asyncio.wait_for(env.call_tool("get_replay_path"), timeout=10.0)
        except Exception as exc:
            replay_info = {"error": f"{type(exc).__name__}: {exc}"}

        last_info = replay_info or {}
        if str(last_info.get("path", "") or ""):
            if verbose and attempt > 1:
                print(f"  Episode {episode_id}: Replay became available after {attempt} check(s).")
            return last_info

        now = time.time()
        if now >= deadline:
            return last_info

        if verbose and not announced_wait:
            print(f"  Episode {episode_id}: Waiting for replay file to be written...")
            announced_wait = True

        await asyncio.sleep(poll_interval_s)


async def collect_episode(
    game_url: str,
    episode_id: int,
    max_steps: int = 20000,
    max_minutes: float = 15.0,
    map_name: str = "singles.oramap",
    verbose: bool = False,
    bot_type: str = "scripted",
    sample_every: int = 12,
    keep_state_changes: bool = False,
    top_k_types: int = 12,
) -> dict:
    """Play one full game and record compact macro-policy rows.

    Stops when any of these conditions is met (in order):
      1. Game ends (win / lose)
      2. Wall-clock time exceeds max_minutes
      3. Step count reaches max_steps (safety cap)

    Returns:
        Dict with compact macro rows plus replay and summary data.
    """
    if bot_type == "normal":
        bot = NormalAIBot(verbose=verbose)
    else:
        bot = PeriodicAttackBot(verbose=False)
    macro_candidates = []
    replay_info = {}
    error = ""
    max_seconds = max_minutes * 60.0
    episode_start = time.time()
    step = 0
    eliminated_since_step = None
    result = None
    obs = None
    prev_sig = None
    async with ToolEnabledOpenRAEnv(base_url=game_url, message_timeout_s=300.0) as env:
        if verbose:
            print(f"  Episode {episode_id}: Resetting environment (map={map_name})...")

        try:
            result = await env.reset(map_name=map_name)
            obs = result.observation
        except Exception as exc:
            error = f"{type(exc).__name__}: {exc}"
            if verbose:
                print(f"  Episode {episode_id}: RESET ERROR | {error}")
            try:
                replay_info = await env.call_tool("get_replay_path")
            except Exception:
                replay_info = {}
            return {
                "macro_rows": [],
                "replay": replay_info,
                "error": error,
                "final_observation": {},
                "step_count": 0,
                "result": "",
            }

        # Let the bot see the reset observation so _get_map_size can initialize
        # (the correct playable-area dimensions will be refined on later steps)

        if verbose:
            print(
                f"  Episode {episode_id}: Game started on "
                f"{obs.map_info.map_name} ({obs.map_info.width}x{obs.map_info.height})"
            )

        step = 0
        eliminated_since_step = None
        try:
            while not result.done and step < max_steps:
                # Stop if wall-clock time limit exceeded
                elapsed_so_far = time.time() - episode_start
                if elapsed_so_far >= max_seconds:
                    if verbose:
                        print(
                            f"  Episode {episode_id}: Time limit reached "
                            f"({elapsed_so_far/60:.1f} min), stopping."
                        )
                    break

                # Stop early if we've lost everything (0 units + 0 buildings)
                obs_now = result.observation
                if not obs_now.units and not obs_now.buildings and step > 100:
                    if eliminated_since_step is None:
                        eliminated_since_step = step
                    elif step - eliminated_since_step >= 200:
                        if verbose:
                            print(
                                f"  Episode {episode_id}: Eliminated "
                                f"(0 units/buildings for 200 steps), stopping."
                            )
                        break
                else:
                    eliminated_since_step = None

                try:
                    # Expert decides
                    obs_before = result.observation
                    action = bot.decide(obs_before)

                    # Execute and get reward for THIS action
                    result = await env.step(action)
                    step += 1

                    step_data = {
                        "step": step - 1,
                        "observation": serialize_obs(obs_before),
                        "action": serialize_action(action),
                        "reward": result.reward or 0.0,
                        "done": result.done,
                    }
                    macros = extract_macro_actions(step_data["action"])
                    reasons = should_keep_step(
                        step_idx=step_data["step"],
                        step_data=step_data,
                        obs=step_data["observation"],
                        macros=macros,
                        prev_sig=prev_sig,
                        sample_every=sample_every,
                        keep_state_changes=keep_state_changes,
                    )
                    prev_sig = observation_signature(step_data["observation"])
                    if reasons:
                        macro_candidates.append({
                            "step_data": step_data,
                            "state_summary": summarize_observation(step_data["observation"], top_k=top_k_types),
                            "macros": macros,
                            "reasons": reasons,
                        })

                except Exception as exc:
                    error = f"{type(exc).__name__}: {exc}"
                    if verbose:
                        print(f"  Episode {episode_id}: STEP ERROR | {error}")
                    break

                if verbose and step % 200 == 0:
                    eco = result.observation.economy
                    n_units = len(result.observation.units)
                    n_buildings = len(result.observation.buildings)
                    elapsed_min = (time.time() - episode_start) / 60.0
                    credits = available_credits(result.observation)
                    attack_stats = bot.get_attack_stats(result.observation) if hasattr(bot, "get_attack_stats") else None
                    squad_stats = bot.get_squad_stats() if hasattr(bot, "get_squad_stats") else None
                    attack_info = ""
                    if attack_stats is not None:
                        attack_info = (
                            f" | Targeted:{attack_stats['unique_unit_targets']}u/{attack_stats['unique_building_targets']}b"
                            f" | Kills:{attack_stats['units_killed']}u/{attack_stats['buildings_killed']}b"
                        )
                    print(
                        f"  Episode {episode_id}: Step {step:4d} | "
                        f"Tick {result.observation.tick:5d} | "
                        f"Cash:${eco.cash:5d} Ore:{eco.ore:5d} Tot:${credits:5d} | "
                        f"Units:{n_units} Bldgs:{n_buildings} | "
                        f"{bot.phase} | {elapsed_min:.1f}min"
                        f"{attack_info}"
                    )
                    if squad_stats is not None:
                        states = squad_stats.get("states", {})
                        print(
                            "    "
                            f"Squads | idle:{squad_stats['idle_ground']} "
                            f"atk:{squad_stats['attack_squad']} "
                            f"rush:{squad_stats['rush_squad']} "
                            f"prot:{squad_stats['protection_squad']} "
                            f"thr:{squad_stats['assault_threshold']} "
                            f"state:{states}"
                        )

        except KeyboardInterrupt:
            error = "KeyboardInterrupt"
        finally:
            # If we stopped due to time/step limits, surrender to force a game-over so a replay is written.
            if result is not None and not result.done:
                try:
                    await asyncio.wait_for(env.call_tool("surrender"), timeout=30.0)
                except Exception:
                    pass

        replay_timeout_s = 25.0 if result is not None and not result.done else 10.0
        replay_info = await wait_for_replay_artifact(
            env=env,
            episode_id=episode_id,
            verbose=verbose,
            timeout_s=replay_timeout_s,
        )

        # Record the final observation as a terminal no-op example.
        final_obs = result.observation if result is not None else obs
        if final_obs is not None:
            final_step_data = {
                "step": step,
                "observation": serialize_obs(final_obs),
                "action": {"commands": [{"action": "no_op"}]},
                "reward": 0.0,
                "done": True,
            }
            final_macros = extract_macro_actions(final_step_data["action"])
            final_reasons = should_keep_step(
                step_idx=final_step_data["step"],
                step_data=final_step_data,
                obs=final_step_data["observation"],
                macros=final_macros,
                prev_sig=prev_sig,
                sample_every=sample_every,
                keep_state_changes=keep_state_changes,
            )
            if final_reasons:
                macro_candidates.append({
                    "step_data": final_step_data,
                    "state_summary": summarize_observation(final_step_data["observation"], top_k=top_k_types),
                    "macros": final_macros,
                    "reasons": final_reasons,
                })

        outcome = ""
        macro_rows = []
        final_obs_dict = {}
        if final_obs is not None:
            elapsed_total_s = time.time() - episode_start
            outcome = infer_outcome(final_obs, eliminated_since_step, elapsed_total_s, max_minutes)
            final_obs_dict = serialize_obs(final_obs)
            episode_name = f"episode_{episode_id:03d}"
            macro_rows = [
                build_row(
                    episode_name=episode_name,
                    episode_result=outcome,
                    step_data=item["step_data"],
                    state_summary=item["state_summary"],
                    macros=item["macros"],
                    reasons=item["reasons"],
                )
                for item in macro_candidates
            ]

        if verbose and final_obs is not None:
            mil = final_obs.military
            elapsed_total = elapsed_total_s / 60.0
            attack_stats = bot.get_attack_stats(final_obs) if hasattr(bot, "get_attack_stats") else None
            attack_info = ""
            if attack_stats is not None:
                attack_info = (
                    f" | Targeted: {attack_stats['unique_unit_targets']}u/{attack_stats['unique_building_targets']}b"
                    f" | Orders: {attack_stats['attack_commands']} ATTACK, {attack_stats['attack_move_commands']} AMOVE"
                )
            print(
                f"  Episode {episode_id}: DONE — {outcome.upper()} | "
                f"{step} steps | Tick {final_obs.tick} | "
                f"Real time: {elapsed_total:.1f}min | "
                f"Kills: {mil.units_killed}u/{mil.buildings_killed}b | "
                f"Lost: {mil.units_lost}u/{mil.buildings_lost}b"
                f"{attack_info}"
            )

    return {
        "macro_rows": macro_rows,
        "replay": replay_info,
        "error": error,
        "final_observation": final_obs_dict,
        "step_count": step,
        "result": outcome,
    }


async def collect_all(
    game_url: str,
    artifact_url: str,
    num_episodes: int,
    max_steps: int,
    max_minutes: float,
    map_name: str,
    output_dir: Path,
    dataset_path: Path,
    verbose: bool,
    bot_type: str = "scripted",
    sample_every: int = 12,
    keep_state_changes: bool = False,
    top_k_types: int = 12,
    append_dataset: bool = False,
):
    """Collect multiple episodes sequentially."""
    output_dir.mkdir(parents=True, exist_ok=True)

    summaries = []
    total_rows_written = 0

    with open_dataset_writer(dataset_path, append=append_dataset) as dataset_writer:
        for i in range(1, num_episodes + 1):
            print(f"\n{'='*60}")
            print(f"Collecting episode {i}/{num_episodes}")
            print(f"{'='*60}")

            t0 = time.time()
            try:
                episode_data = await collect_episode(
                    game_url=game_url,
                    episode_id=i,
                    max_steps=max_steps,
                    max_minutes=max_minutes,
                    map_name=map_name,
                    verbose=verbose,
                    bot_type=bot_type,
                    sample_every=sample_every,
                    keep_state_changes=keep_state_changes,
                    top_k_types=top_k_types,
                )
            except Exception as e:
                print(f"  Episode {i} FAILED: {e}")
                continue

            elapsed = time.time() - t0
            replay_info = copy_replay_artifact(
                episode_data.get("replay", {}),
                output_dir,
                i,
                env_url=artifact_url,
            )
            episode_error = episode_data.get("error", "")
            macro_rows = episode_data.get("macro_rows", [])
            final = episode_data.get("final_observation", {})

            for row in macro_rows:
                dataset_writer.write(json.dumps(row, separators=(",", ":")) + "\n")
            dataset_writer.flush()
            total_rows_written += len(macro_rows)

            mil = final.get("military", {})
            final_units = len(final.get("units", []))
            final_buildings = len(final.get("buildings", []))
            summary_result = episode_data.get("result") or final.get("result")
            if not summary_result:
                if final_units == 0 and final_buildings == 0:
                    summary_result = "eliminated"
                elif elapsed >= max_minutes * 60.0:
                    summary_result = f"time_limit({max_minutes:.0f}min)"
                else:
                    summary_result = "step_limit"
            summary = {
                "episode": i,
                "steps": episode_data.get("step_count", 0),
                "macro_rows": len(macro_rows),
                "ticks": final.get("tick", 0),
                "result": summary_result,
                "kills_cost": mil.get("kills_cost", 0),
                "deaths_cost": mil.get("deaths_cost", 0),
                "explored_percent": final.get("explored_percent", 0),
                "final_buildings": final_buildings,
                "final_units": final_units,
                "elapsed_s": round(elapsed, 1),
                "dataset_path": str(dataset_path),
                "replay_path": replay_info.get("path", ""),
                "replay_local_copy": replay_info.get("local_copy", ""),
                "replay_lookup_error": replay_info.get("error", ""),
                "replay_download_error": replay_info.get("download_error", ""),
                "replay_cleanup_error": replay_info.get("cleanup_error", ""),
                "error": episode_error,
            }
            summaries.append(summary)

            print(f"  Episode {i} finished ({episode_data.get('step_count', 0)} steps, {elapsed:.0f}s)")
            if replay_info.get("local_copy"):
                print(f"  Replay saved to {Path(replay_info['local_copy']).name}")
            elif replay_info.get("path"):
                print(f"  Replay available at {replay_info['path']}")
            elif replay_info.get("error"):
                print(f"  Replay unavailable: {replay_info['error']}")
            if replay_info.get("download_error"):
                print(f"  Replay download failed: {replay_info['download_error']}")
            cleanup_info = replay_info.get("remote_cleanup", {})
            if cleanup_info:
                print(
                    "  Remote cleanup:"
                    f" {len(cleanup_info.get('deleted_replays', []))} replay(s),"
                    f" {len(cleanup_info.get('deleted_logs', []))} log file(s)"
                )
            if replay_info.get("cleanup_error"):
                print(f"  Remote cleanup failed: {replay_info['cleanup_error']}")
            if episode_error:
                print(f"  Episode {i} completed with error: {episode_error}")

    # Save collection summary
    summary_file = output_dir / "collection_summary.json"
    with open(summary_file, "w") as f:
        json.dump(summaries, f, indent=2)

    print(f"\n{'='*60}")
    print(f"Collection complete: {len(summaries)}/{num_episodes} episodes")
    print(f"Summary: {summary_file}")
    print(f"{'='*60}")

    # Print results table
    if summaries:
        print(f"\n{'Episode':>8} {'Result':>10} {'Steps':>6} {'Kills$':>8} {'Deaths$':>8} {'Bldgs':>6} {'Units':>6}")
        print("-" * 62)
        for s in summaries:
            result_str = s['result'] if s['result'] else 'timeout'
            print(
                f"{s['episode']:>8} {result_str:>10} {s['steps']:>6} "
                f"{s['kills_cost']:>8} {s['deaths_cost']:>8} "
                f"{s['final_buildings']:>6} {s['final_units']:>6}"
            )


def main():
    parser = argparse.ArgumentParser(
        description="Collect compact macro-policy demonstrations for behavior cloning"
    )
    parser.add_argument(
        "--url",
        default="http://localhost:8000",
        help="OpenRA-RL server URL (default: http://localhost:8000)",
    )
    parser.add_argument(
        "--episodes",
        type=int,
        default=10,
        help="Number of episodes to collect (default: 10)",
    )
    parser.add_argument(
        "--map",
        default="singles.oramap",
        help=(
            "Map filename to play on (default: singles.oramap). "
            "Must be one of the .oramap files inside the server container. "
            "List available maps with: "
            "docker exec openra-rl-server find /opt/openra/mods/ra/maps -name '*.oramap'"
        ),
    )
    parser.add_argument(
        "--max-minutes",
        type=float,
        default=15.0,
        help="Max real-time minutes per episode (default: 15). Episodes stop early if game ends.",
    )
    parser.add_argument(
        "--max-steps",
        type=int,
        default=200000,
        help="Hard step cap per episode (default: 200000, effectively unlimited). Use --max-minutes instead.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/episodes"),
        help="Output directory for replays and collection summary (default: data/episodes)",
    )
    parser.add_argument(
        "--dataset-path",
        type=Path,
        default=Path("data/macro/macro_dataset.jsonl.gz"),
        help="Compact macro dataset path (default: data/macro/macro_dataset.jsonl.gz)",
    )
    parser.add_argument(
        "--bot",
        choices=["scripted", "normal"],
        default="scripted",
        help=(
            "Which Python bot to use (default: scripted). "
            "'normal' uses NormalAIBot that mimics OpenRA's built-in normal AI."
        ),
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print per-step progress",
    )
    parser.add_argument(
        "--sample-every",
        type=int,
        default=12,
        help="Keep every Nth step as a periodic macro snapshot (default: 12, 0 disables periodic sampling)",
    )
    parser.add_argument(
        "--keep-state-changes",
        action="store_true",
        help="Also keep steps where coarse unit/building/combat counts changed",
    )
    parser.add_argument(
        "--top-k-types",
        type=int,
        default=12,
        help="Maximum number of unit/building types to keep in summaries (default: 12)",
    )
    parser.add_argument(
        "--append-dataset",
        action="store_true",
        help="Append new rows to an existing macro dataset instead of overwriting it",
    )
    args = parser.parse_args()
    server_urls = resolve_server_urls(args.url)

    if server_urls["game_url"] != args.url.rstrip("/"):
        print(f"Using OpenRA env at {server_urls['game_url']}")
        print(f"Using artifact helper endpoints at {server_urls['artifact_url']}")

    try:
        asyncio.run(
            collect_all(
                game_url=server_urls["game_url"],
                artifact_url=server_urls["artifact_url"],
                num_episodes=args.episodes,
                max_steps=args.max_steps,
                max_minutes=args.max_minutes,
                map_name=args.map,
                output_dir=args.output_dir,
                dataset_path=args.dataset_path,
                verbose=args.verbose,
                bot_type=args.bot,
                sample_every=args.sample_every,
                keep_state_changes=args.keep_state_changes,
                top_k_types=args.top_k_types,
                append_dataset=args.append_dataset,
            )
        )
    except KeyboardInterrupt:
        print("\nInterrupted by user")
        sys.exit(0)
    except ConnectionRefusedError:
        print(f"\nCould not connect to {args.url}")
        print("Is the OpenRA-RL server running?")
        print("  docker run -p 8000:8000 openra-rl")
        sys.exit(1)


if __name__ == "__main__":
    main()