1use std::{fmt, fmt::Debug};
2
3use azalea_client::{
4 Client,
5 inventory::{CloseContainerEvent, ContainerClickEvent},
6 packet::game::ReceiveGamePacketEvent,
7};
8use azalea_core::position::BlockPos;
9use azalea_entity::inventory::Inventory;
10use azalea_inventory::{
11 ItemStack, Menu,
12 operations::{ClickOperation, PickupClick, QuickMoveClick},
13};
14use azalea_physics::collision::BlockWithShape;
15use azalea_protocol::packets::game::ClientboundGamePacket;
16use bevy_app::{App, Plugin, Update};
17use bevy_ecs::{component::Component, prelude::MessageReader, system::Commands};
18use derive_more::Deref;
19use futures_lite::Future;
20
21use crate::bot::BotClientExt;
22
23pub struct ContainerPlugin;
24impl Plugin for ContainerPlugin {
25 fn build(&self, app: &mut App) {
26 app.add_systems(Update, handle_menu_opened_event);
27 }
28}
29
30pub trait ContainerClientExt {
31 fn open_container_at(
54 &self,
55 pos: BlockPos,
56 ) -> impl Future<Output = Option<ContainerHandle>> + Send;
57
58 fn open_container_at_with_timeout_ticks(
70 &self,
71 pos: BlockPos,
72 timeout_ticks: Option<usize>,
73 ) -> impl Future<Output = Option<ContainerHandle>> + Send;
74
75 fn wait_for_container_open(
82 &self,
83 timeout_ticks: Option<usize>,
84 ) -> impl Future<Output = Option<ContainerHandle>> + Send;
85
86 fn open_inventory(&self) -> Option<ContainerHandle>;
98 fn get_inventory(&self) -> ContainerHandleRef;
110 fn get_held_item(&self) -> ItemStack;
113}
114
115impl ContainerClientExt for Client {
116 async fn open_container_at(&self, pos: BlockPos) -> Option<ContainerHandle> {
117 self.open_container_at_with_timeout_ticks(pos, Some(20 * 5))
118 .await
119 }
120
121 async fn open_container_at_with_timeout_ticks(
122 &self,
123 pos: BlockPos,
124 timeout_ticks: Option<usize>,
125 ) -> Option<ContainerHandle> {
126 let mut ticks = self.get_tick_broadcaster();
127 for _ in 0..10 {
129 let block = self.world().read().get_block_state(pos).unwrap_or_default();
130 if !block.is_collision_shape_empty() {
131 break;
132 }
133 let _ = ticks.recv().await;
134 }
135
136 self.ecs
137 .lock()
138 .entity_mut(self.entity)
139 .insert(WaitingForInventoryOpen);
140 self.block_interact(pos);
141
142 self.wait_for_container_open(timeout_ticks).await
143 }
144
145 async fn wait_for_container_open(
146 &self,
147 timeout_ticks: Option<usize>,
148 ) -> Option<ContainerHandle> {
149 let mut ticks = self.get_tick_broadcaster();
150 let mut elapsed_ticks = 0;
151 while ticks.recv().await.is_ok() {
152 let ecs = self.ecs.lock();
153 if ecs.get::<WaitingForInventoryOpen>(self.entity).is_none() {
154 break;
155 }
156
157 elapsed_ticks += 1;
158 if let Some(timeout_ticks) = timeout_ticks
159 && elapsed_ticks >= timeout_ticks
160 {
161 return None;
162 }
163 }
164
165 let ecs = self.ecs.lock();
166 let inventory = ecs.get::<Inventory>(self.entity).expect("no inventory");
167 if inventory.id == 0 {
168 None
169 } else {
170 Some(ContainerHandle::new(inventory.id, self.clone()))
171 }
172 }
173
174 fn open_inventory(&self) -> Option<ContainerHandle> {
175 let ecs = self.ecs.lock();
176 let inventory = ecs.get::<Inventory>(self.entity).expect("no inventory");
177 if inventory.id == 0 {
178 Some(ContainerHandle::new(0, self.clone()))
179 } else {
180 None
181 }
182 }
183
184 fn get_inventory(&self) -> ContainerHandleRef {
185 self.query_self::<&Inventory, _>(|inv| ContainerHandleRef::new(inv.id, self.clone()))
186 }
187
188 fn get_held_item(&self) -> ItemStack {
189 self.query_self::<&Inventory, _>(|inv| inv.held_item().clone())
190 }
191}
192
193pub struct ContainerHandleRef {
198 id: i32,
199 client: Client,
200}
201impl Debug for ContainerHandleRef {
202 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
203 f.debug_struct("ContainerHandle")
204 .field("id", &self.id())
205 .finish()
206 }
207}
208impl ContainerHandleRef {
209 pub fn new(id: i32, client: Client) -> Self {
210 Self { id, client }
211 }
212
213 pub fn close(&self) {
214 self.client.ecs.lock().trigger(CloseContainerEvent {
215 entity: self.client.entity,
216 id: self.id,
217 });
218 }
219
220 pub fn id(&self) -> i32 {
226 self.id
227 }
228
229 pub fn menu(&self) -> Option<Menu> {
237 let ecs = self.client.ecs.lock();
238 let inventory = ecs
239 .get::<Inventory>(self.client.entity)
240 .expect("no inventory");
241
242 if inventory.id == self.id {
244 if self.id == 0 {
245 Some(inventory.inventory_menu.clone())
246 } else {
247 Some(inventory.container_menu.clone().unwrap())
248 }
249 } else {
250 None
251 }
252 }
253
254 pub fn contents(&self) -> Option<Vec<ItemStack>> {
259 self.menu().map(|menu| menu.contents())
260 }
261
262 pub fn slots(&self) -> Option<Vec<ItemStack>> {
266 self.menu().map(|menu| menu.slots())
267 }
268
269 pub fn left_click(&self, slot: impl Into<usize>) {
271 self.click(PickupClick::Left {
272 slot: Some(slot.into() as u16),
273 });
274 }
275 pub fn shift_click(&self, slot: impl Into<usize>) {
277 self.click(QuickMoveClick::Left {
278 slot: slot.into() as u16,
279 });
280 }
281 pub fn right_click(&self, slot: impl Into<usize>) {
283 self.click(PickupClick::Right {
284 slot: Some(slot.into() as u16),
285 });
286 }
287
288 pub fn click(&self, operation: impl Into<ClickOperation>) {
291 let operation = operation.into();
292 self.client.ecs.lock().trigger(ContainerClickEvent {
293 entity: self.client.entity,
294 window_id: self.id,
295 operation,
296 });
297 }
298}
299
300#[derive(Deref)]
304pub struct ContainerHandle(ContainerHandleRef);
305
306impl Drop for ContainerHandle {
307 fn drop(&mut self) {
308 self.0.close();
309 }
310}
311impl Debug for ContainerHandle {
312 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
313 f.debug_struct("ContainerHandle")
314 .field("id", &self.id())
315 .finish()
316 }
317}
318impl ContainerHandle {
319 fn new(id: i32, client: Client) -> Self {
320 Self(ContainerHandleRef { id, client })
321 }
322
323 pub fn close(self) {
325 }
327}
328
329#[derive(Component, Debug)]
330pub struct WaitingForInventoryOpen;
331
332pub fn handle_menu_opened_event(
333 mut commands: Commands,
334 mut events: MessageReader<ReceiveGamePacketEvent>,
335) {
336 for event in events.read() {
337 if let ClientboundGamePacket::ContainerSetContent { .. } = event.packet.as_ref() {
338 commands
339 .entity(event.entity)
340 .remove::<WaitingForInventoryOpen>();
341 }
342 }
343}