updates

protossbot/src/bot.rs

@@ -1,10 +1,18 @@
-use std::sync::{Arc, Mutex};
-
 use rsbwapi::*;
+use std::sync::{Arc, Mutex};
+use crate::{state::game_state::GameState, utils::{build_manager, worker_management}};
 
-use crate::game_state::GameState;
+fn draw_unit_ids(game: &Game) {
+  for unit in game.get_all_units() {
+    if unit.exists() {
+      let pos = unit.get_position();
+      let unit_id = unit.get_id();
+      game.draw_text_map(pos, &format!("{}", unit_id));
+    }
+  }
+}
 
-impl AiModule for RustBot {
+impl AiModule for ProtosBot {
   fn on_start(&mut self, game: &Game) {
     // SAFETY: rsbwapi uses interior mutability (RefCell) for the command queue.
     // enable_flag only adds a command to the queue.
@@ -15,7 +23,6 @@ impl AiModule for RustBot {
     }
 
     println!("Game started on map: {}", game.map_file_name());
-
   }
 
   fn on_frame(&mut self, game: &Game) {
@@ -24,6 +31,15 @@ impl AiModule for RustBot {
       return;
     };
 
+    let Some(player) = game.self_() else {
+      return;
+    };
+
+    build_manager::on_frame(game, &player, &mut locked_state);
+    worker_management::assign_idle_workers_to_minerals(game, &player, &mut locked_state);
+
+    build_manager::print_debug_build_status(game, &player, &locked_state);
+    draw_unit_ids(game);
   }
 
   fn on_unit_create(&mut self, game: &Game, unit: Unit) {
@@ -33,19 +49,14 @@ impl AiModule for RustBot {
     println!("unit created: {:?}", unit.get_type());
   }
 
-  fn on_unit_morph(&mut self, game: &Game, unit: Unit) {
+  fn on_unit_morph(&mut self, _game: &Game, _unit: Unit) {}
 
-  }
-
-  fn on_unit_destroy(&mut self, _game: &Game, unit: Unit) {
-  
-  }
-
-  fn on_unit_complete(&mut self, game: &Game, unit: Unit) {
-    let Some(player) = game.self_() else {
-      return;
-    };
+  fn on_unit_destroy(&mut self, _game: &Game, _unit: Unit) {}
 
+  fn on_unit_complete(&mut self, _game: &Game, _unit: Unit) {
+    // let Some(player) = game.self_() else {
+    //   return;
+    // };
   }
 
   fn on_end(&mut self, _game: &Game, is_winner: bool) {
@@ -56,14 +67,13 @@ impl AiModule for RustBot {
     }
   }
 }
-pub struct RustBot {
+
+pub struct ProtosBot {
   game_state: Arc<Mutex<GameState>>,
 }
 
-impl RustBot {
+impl ProtosBot {
   pub fn new(game_state: Arc<Mutex<GameState>>) -> Self {
-    Self {
-      game_state,
-    }
+    Self { game_state }
   }
-}
+}

protossbot/src/game_state.rs

@@ -1,14 +0,0 @@
-
-
-pub struct GameState {
-
-}
-
-
-impl Default for GameState {
-  fn default() -> Self {
-    Self {
-
-    }
-  }
-}

protossbot/src/main.rs

@@ -1,9 +1,10 @@
 mod bot;
-pub mod game_state;
+mod state;
+mod utils;
 
-use bot::RustBot;
+use bot::ProtosBot;
 use std::sync::{Arc, Mutex};
-use game_state::GameState;
+use state::game_state::GameState;
 
 
 fn main() {
@@ -11,5 +12,5 @@ fn main() {
 
   let game_state = Arc::new(Mutex::new(GameState::default()));
 
-  rsbwapi::start(move |_game| RustBot::new(game_state.clone() ));
+  rsbwapi::start(move |_game| ProtosBot::new(game_state.clone() ));
 }

protossbot/src/state/build_stages.rs

@@ -0,0 +1,46 @@
+use rsbwapi::UnitType;
+use std::collections::HashMap;
+
+#[derive(Clone, Debug)]
+pub struct BuildStage {
+  pub name: String,
+  pub desired_counts: HashMap<UnitType, i32>,
+}
+
+impl BuildStage {
+  pub fn new(name: &str) -> Self {
+    Self {
+      name: name.to_string(),
+      desired_counts: HashMap::new(),
+    }
+  }
+
+  pub fn with_unit(mut self, unit_type: UnitType, count: i32) -> Self {
+    self.desired_counts.insert(unit_type, count);
+    self
+  }
+}
+
+pub fn get_build_stages() -> Vec<BuildStage> {
+  vec![
+    BuildStage::new("Start")
+      .with_unit(UnitType::Protoss_Probe, 10)
+      .with_unit(UnitType::Protoss_Pylon, 1),
+
+    BuildStage::new("Basic Production")
+      .with_unit(UnitType::Protoss_Probe, 12)
+      .with_unit(UnitType::Protoss_Pylon, 2)
+      .with_unit(UnitType::Protoss_Gateway, 1)
+      .with_unit(UnitType::Protoss_Forge, 1),
+
+
+    // Stage 2: Defense cannons
+    BuildStage::new("Defense Cannons")
+      .with_unit(UnitType::Protoss_Probe, 16)
+      .with_unit(UnitType::Protoss_Pylon, 3)
+      .with_unit(UnitType::Protoss_Nexus, 1)
+      .with_unit(UnitType::Protoss_Gateway, 1)
+      .with_unit(UnitType::Protoss_Forge, 1)
+      .with_unit(UnitType::Protoss_Photon_Cannon, 4),
+  ]
+}

protossbot/src/state/game_state.rs

@@ -0,0 +1,46 @@
+use std::collections::HashMap;
+use rsbwapi::{Order, Position, Unit, UnitType, UpgradeType};
+
+use crate::state::build_stages::BuildStage;
+
+pub struct GameState {
+  pub intended_commands: HashMap<usize, IntendedCommand>,
+  pub unit_build_history: Vec<BuildHistoryEntry>,
+  pub build_stages: Vec<BuildStage>,
+  pub current_stage_index: usize,
+  pub desired_game_speed: i32,
+}
+
+
+impl Default for GameState {
+  fn default() -> Self {
+    Self {
+      intended_commands: HashMap::new(),
+      unit_build_history: Vec::new(),
+      build_stages: crate::state::build_stages::get_build_stages(),
+      current_stage_index: 0,
+      desired_game_speed: 20,
+    }
+  }
+}
+
+#[derive(Clone, Debug)]
+pub struct IntendedCommand {
+  pub order: Order,
+  pub target_position: Option<Position>,
+  pub target_unit: Option<Unit>,
+}
+
+#[derive(Clone, Debug)]
+pub enum BuildStatus {
+  Assigned,
+  Started,
+}
+
+#[derive(Clone, Debug)]
+pub struct BuildHistoryEntry {
+  pub unit_type: Option<UnitType>,
+  pub upgrade_type: Option<UpgradeType>,
+  pub assigned_unit_id: Option<usize>,
+  // pub status: BuildStatus,
+}

protossbot/src/state/mod.rs

@@ -0,0 +1,2 @@
+pub mod build_stages;
+pub mod game_state;

protossbot/src/utils/build_location_utils.rs

@@ -0,0 +1,53 @@
+use rsbwapi::{Game, TilePosition, Unit, UnitType};
+
+pub fn find_build_location(
+  game: &Game,
+  builder: &Unit,
+  building_type: UnitType,
+  max_range: i32,
+) -> Option<TilePosition> {
+
+  let start_tile = builder.get_tile_position();
+
+  for distance in 0..max_range {
+    for dx in -distance..=distance {
+      for dy in -distance..=distance {
+        if dx.abs() != distance && dy.abs() != distance {
+          continue;
+        }
+
+        let tile = TilePosition {
+          x: start_tile.x + dx,
+          y: start_tile.y + dy,
+        };
+
+        if is_valid_build_location(game, building_type, tile, builder) {
+          return Some(tile);
+        }
+      }
+    }
+  }
+
+  None
+}
+
+fn is_valid_build_location(
+  game: &Game,
+  building_type: UnitType,
+  position: TilePosition,
+  builder: &Unit,
+) -> bool {
+  if !game.can_build_here(builder, position, building_type, false).unwrap_or(false) {
+    return false;
+  }
+
+  // if building_type.requires_psi() && !game.has_power(position, building_type) {
+  //   return false;
+  // }
+
+  // if building_type.requires_creep() && !game.has_creep(position) {
+  //   return false;
+  // }
+
+  true
+}

protossbot/src/utils/build_manager.rs

@@ -0,0 +1,268 @@
+use rsbwapi::{Game, Order, Player, UnitType};
+
+use crate::{
+  state::game_state::{BuildHistoryEntry, GameState, IntendedCommand},
+  utils::build_location_utils,
+};
+
+pub fn on_frame(game: &Game, player: &Player, state: &mut GameState) {
+  check_and_advance_stage(player, state);
+  try_start_next_build(game, player, state);
+}
+
+fn try_start_next_build(game: &Game, player: &Player, state: &mut GameState) {
+  let Some(unit_type) = get_next_thing_to_build(game, player, state) else {
+    return;
+  };
+
+  let Some(builder) = find_builder_for_unit(player, unit_type) else {
+    return;
+  };
+
+  let builder_id = builder.get_id();
+
+  if assign_builder_to_construct(game, &builder, unit_type, state) {
+    let entry = BuildHistoryEntry {
+      unit_type: Some(unit_type),
+      upgrade_type: None,
+      assigned_unit_id: Some(builder_id),
+    };
+
+    state.unit_build_history.push(entry);
+
+    let current_stage = &state.build_stages[state.current_stage_index];
+    println!(
+      "Started building {} with unit {} (Stage: {})",
+      unit_type.name(),
+      builder_id,
+      current_stage.name
+    );
+  }
+}
+
+fn get_next_thing_to_build(game: &Game, player: &Player, state: &GameState) -> Option<UnitType> {
+  let current_stage = state.build_stages.get(state.current_stage_index)?;
+
+  if let Some(pylon) = check_need_more_supply(game, player) {
+    return Some(pylon);
+  }
+
+  let mut candidates = Vec::new();
+
+  for (unit_type, &desired_count) in &current_stage.desired_counts {
+    let current_count = count_units_of_type(player, state, *unit_type);
+
+    if current_count >= desired_count {
+      continue;
+    }
+
+    if !can_afford_unit(player, *unit_type) {
+      continue;
+    }
+
+    if unit_type.is_building() {
+      let builder = find_builder_for_unit(player, *unit_type)?;
+      let build_location = build_location_utils::find_build_location(game, &builder, *unit_type, 20);
+      if build_location.is_none() {
+        continue;
+      }
+    }
+
+    candidates.push(*unit_type);
+  }
+
+  candidates.into_iter().max_by_key(|unit_type| {
+    unit_type.mineral_price() + unit_type.gas_price()
+  })
+}
+
+fn check_need_more_supply(game: &Game, player: &Player) -> Option<UnitType> {
+  let supply_used = player.supply_used();
+  let supply_total = player.supply_total();
+  
+  if supply_total == 0 {
+    return None;
+  }
+
+  let supply_remaining = supply_total - supply_used;
+  let threshold = ((supply_total as f32) * 0.15).ceil() as i32;
+  
+  if supply_remaining <= threshold && supply_total < 400 {
+    let pylon_type = UnitType::Protoss_Pylon;
+    
+    if can_afford_unit(player, pylon_type) {
+      if let Some(builder) = find_builder_for_unit(player, pylon_type) {
+        let build_location = build_location_utils::find_build_location(game, &builder, pylon_type, 20);
+        if build_location.is_some() {
+          return Some(pylon_type);
+        }
+      }
+    }
+  }
+
+  None
+}
+
+fn find_builder_for_unit(player: &Player, unit_type: UnitType) -> Option<rsbwapi::Unit> {
+  let builder_type = unit_type.what_builds().0;
+
+  player
+    .get_units()
+    .iter()
+    .find(|u| {
+      u.get_type() == builder_type && !u.is_constructing() && !u.is_training() && u.is_idle()
+    })
+    .cloned()
+}
+
+fn assign_builder_to_construct(game: &Game, builder: &rsbwapi::Unit, unit_type: UnitType, state: &mut GameState) -> bool {
+  let builder_id = builder.get_id();
+  
+  if unit_type.is_building() {
+    let build_location = build_location_utils::find_build_location(game, builder, unit_type, 20);
+
+    if let Some(pos) = build_location {
+      println!(
+        "Attempting to build {} at {:?} with worker {} (currently at {:?})",
+        unit_type.name(),
+        pos,
+        builder_id,
+        builder.get_position()
+      );
+      
+      match builder.build(unit_type, pos) {
+        Ok(_) => {
+          println!("Build command succeeded for {}", unit_type.name());
+          let intended_cmd = IntendedCommand {
+            order: Order::PlaceBuilding,
+            target_position: Some(pos.to_position()),
+            target_unit: None,
+          };
+          state.intended_commands.insert(builder_id, intended_cmd);
+          true
+        }
+        Err(e) => {
+          println!("Build command FAILED for {}: {:?}", unit_type.name(), e);
+          false
+        }
+      }
+    } else {
+      println!(
+        "No valid build location found for {} by builder {}",
+        unit_type.name(),
+        builder.get_id()
+      );
+      false
+    }
+  } else {
+    match builder.train(unit_type) {
+      Ok(_) => {
+        let intended_cmd = IntendedCommand {
+          order: Order::Train,
+          target_position: None,
+          target_unit: None,
+        };
+        state.intended_commands.insert(builder_id, intended_cmd);
+        true
+      }
+      Err(e) => {
+        println!("Train command FAILED for {}: {:?}", unit_type.name(), e);
+        false
+      }
+    }
+  }
+}
+
+fn count_units_of_type(player: &Player, _state: &GameState, unit_type: UnitType) -> i32 {
+  let existing = player
+    .get_units()
+    .iter()
+    .filter(|u| u.get_type() == unit_type)
+    .count() as i32;
+
+
+
+  existing
+}
+
+fn can_afford_unit(player: &Player, unit_type: UnitType) -> bool {
+  let minerals = player.minerals();
+  let gas = player.gas();
+
+  minerals >= unit_type.mineral_price() && gas >= unit_type.gas_price()
+}
+
+fn check_and_advance_stage(player: &Player, state: &mut GameState) {
+  let Some(current_stage) = state.build_stages.get(state.current_stage_index) else {
+    return;
+  };
+
+  let stage_complete = current_stage
+    .desired_counts
+    .iter()
+    .all(|(unit_type, &desired_count)| {
+      let current_count = count_units_of_type(player, state, *unit_type);
+      current_count >= desired_count
+    });
+
+  if stage_complete {
+    let next_stage_index = state.current_stage_index + 1;
+
+    if next_stage_index < state.build_stages.len() {
+      println!(
+        "Stage '{}' complete! Advancing to stage {}",
+        current_stage.name, state.build_stages[next_stage_index].name
+      );
+      state.current_stage_index = next_stage_index;
+    }
+  }
+}
+
+pub fn print_debug_build_status(game: &Game, player: &Player, state: &GameState) {
+  let mut y = 10;
+  let x = 3;
+
+  if let Some(current_stage) = state.build_stages.get(state.current_stage_index) {
+    let next_build = get_next_thing_to_build(game, player, state);
+    let next_build_str = if let Some(unit_type) = next_build {
+      format!(
+        "Next: {} ({}/{} M, {}/{} G)",
+        unit_type.name(),
+        player.minerals(),
+        unit_type.mineral_price(),
+        player.gas(),
+        unit_type.gas_price()
+      )
+    } else {
+      "Next: None".to_string()
+    };
+    game.draw_text_screen((x, y), &next_build_str);
+    y += 10;
+
+    if let Some(last_entry) = state.unit_build_history.last() {
+      let unit_name = if let Some(unit_type) = last_entry.unit_type {
+        unit_type.name()
+      } else if let Some(_upgrade) = last_entry.upgrade_type {
+        "Upgrade"
+      } else {
+        "Unknown"
+      };
+      game.draw_text_screen((x, y), &format!("Last Built: {}", unit_name));
+    } else {
+      game.draw_text_screen((x, y), "Last Built: None");
+    }
+    y += 10;
+
+    game.draw_text_screen((x, y), "Stage Progress:");
+    y += 10;
+
+    for (unit_type, &desired_count) in &current_stage.desired_counts {
+      let current_count = count_units_of_type(player, state, *unit_type);
+      game.draw_text_screen(
+        (x + 10, y),
+        &format!("{}: {}/{}", unit_type.name(), current_count, desired_count),
+      );
+      y += 10;
+    }
+  }
+}

protossbot/src/utils/mod.rs

@@ -0,0 +1,3 @@
+pub mod build_location_utils;
+pub mod build_manager;
+pub mod worker_management;

protossbot/src/utils/worker_management.rs

@@ -0,0 +1,122 @@
+use rsbwapi::{Game, Order, Player, Unit};
+
+use crate::state::game_state::{GameState, IntendedCommand};
+
+pub fn assign_idle_workers_to_minerals(game: &Game, player: &Player, state: &mut GameState) {
+  let all_units = player.get_units();
+  let workers: Vec<Unit> = all_units
+    .iter()
+    .filter(|u| u.get_type().is_worker() && u.is_completed())
+    .cloned()
+    .collect();
+
+  // First, clean up workers that finished building
+  reassign_finished_builders(game, &workers, state);
+
+  // Then assign idle workers to mining
+  for worker in workers {
+    assign_worker_to_mineral(game, &worker, state);
+  }
+}
+
+fn reassign_finished_builders(_game: &Game, workers: &[Unit], state: &mut GameState) {
+  for worker in workers {
+    let worker_id = worker.get_id();
+    
+    if let Some(cmd) = state.intended_commands.get(&worker_id) {
+      // For building commands, only remove if the worker was constructing and now is not
+      // This prevents premature removal when the worker is still moving to the build site
+      if cmd.order == Order::PlaceBuilding {
+        // Worker finished if it WAS constructing but no longer is and is idle
+        // We check the actual order to see if it's moved on from building
+        let current_order = worker.get_order();
+        if worker.is_idle() && current_order != Order::PlaceBuilding && current_order != Order::ConstructingBuilding {
+          println!("Worker {} finished building, reassigning to minerals", worker_id);
+          state.intended_commands.remove(&worker_id);
+        }
+      } else if cmd.order == Order::Train && worker.is_idle() && !worker.is_training() {
+        // For training, the original logic is fine
+        state.intended_commands.remove(&worker_id);
+      }
+    }
+  }
+}
+
+fn assign_worker_to_mineral(game: &Game, worker: &Unit, state: &mut GameState) {
+  let worker_id = worker.get_id();
+
+  // Don't reassign workers that have a non-mining intended command
+  if let Some(cmd) = state.intended_commands.get(&worker_id) {
+    if cmd.order != Order::MiningMinerals {
+      return;
+    }
+  }
+
+  // Only assign idle workers
+  if !worker.is_idle() {
+    return;
+  }
+
+  if worker.is_gathering_minerals() || worker.is_gathering_gas() {
+    return;
+  }
+
+  let Some(mineral) = find_available_mineral(game, worker, state) else {
+    return;
+  };
+
+  let intended_cmd = IntendedCommand {
+    order: Order::MiningMinerals,
+    target_position: None,
+    target_unit: Some(mineral.clone()),
+  };
+
+  state.intended_commands.insert(worker_id, intended_cmd);
+
+  println!(
+    "Worker {} current order: {:?}, assigning to mine from mineral at {:?}",
+    worker_id,
+    worker.get_order(),
+    mineral.get_position()
+  );
+
+  if worker.gather(&mineral).is_ok() {
+    println!(
+      "Assigned worker {} to mine from mineral at {:?}",
+      worker_id,
+      mineral.get_position()
+    );
+  }
+}
+
+fn find_available_mineral(game: &Game, worker: &Unit, state: &GameState) -> Option<Unit> {
+  let worker_pos = worker.get_position();
+  let minerals = game.get_static_minerals();
+  let mut mineral_list: Vec<Unit> = minerals.iter().filter(|m| m.exists()).cloned().collect();
+
+  mineral_list.sort_by_key(|m| {
+    let pos = m.get_position();
+    ((pos.x - worker_pos.x).pow(2) + (pos.y - worker_pos.y).pow(2)) as i32
+  });
+
+  for mineral in mineral_list.iter() {
+    let mineral_id: usize = mineral.get_id();
+
+    let worker_count = state
+      .intended_commands
+      .values()
+      .filter(|cmd| {
+        if let Some(target) = &cmd.target_unit {
+          target.get_id() == mineral_id
+        } else {
+          false
+        }
+      })
+      .count();
+
+    if worker_count < 2 {
+      return Some(mineral.clone());
+    }
+  }
+  mineral_list.first().cloned()
+}