1use std::{fmt, fmt::Debug};
2
3use azalea_client::{
4 Client,
5 inventory::{CloseContainerEvent, ContainerClickEvent, Inventory},
6 packet::game::ReceiveGamePacketEvent,
7};
8use azalea_core::position::BlockPos;
9use azalea_inventory::{
10 ItemStack, Menu,
11 operations::{ClickOperation, PickupClick, QuickMoveClick},
12};
13use azalea_physics::collision::BlockWithShape;
14use azalea_protocol::packets::game::ClientboundGamePacket;
15use bevy_app::{App, Plugin, Update};
16use bevy_ecs::{component::Component, prelude::MessageReader, system::Commands};
17use derive_more::Deref;
18use futures_lite::Future;
19
20use crate::bot::BotClientExt;
21
22pub struct ContainerPlugin;
23impl Plugin for ContainerPlugin {
24 fn build(&self, app: &mut App) {
25 app.add_systems(Update, handle_menu_opened_event);
26 }
27}
28
29pub trait ContainerClientExt {
30 fn open_container_at(
54 &self,
55 pos: BlockPos,
56 timeout_ticks: Option<usize>,
57 ) -> impl Future<Output = Option<ContainerHandle>> + Send;
58
59 fn wait_for_container_open(
66 &self,
67 timeout_ticks: Option<usize>,
68 ) -> impl Future<Output = Option<ContainerHandle>> + Send;
69
70 fn open_inventory(&self) -> Option<ContainerHandle>;
82 fn get_inventory(&self) -> ContainerHandleRef;
94 fn get_held_item(&self) -> ItemStack;
97}
98
99impl ContainerClientExt for Client {
100 async fn open_container_at(
101 &self,
102 pos: BlockPos,
103 timeout_ticks: Option<usize>,
104 ) -> Option<ContainerHandle> {
105 let mut ticks = self.get_tick_broadcaster();
106 for _ in 0..10 {
108 let block = self.world().read().get_block_state(pos).unwrap_or_default();
109 if !block.is_collision_shape_empty() {
110 break;
111 }
112 let _ = ticks.recv().await;
113 }
114
115 self.ecs
116 .lock()
117 .entity_mut(self.entity)
118 .insert(WaitingForInventoryOpen);
119 self.block_interact(pos);
120
121 self.wait_for_container_open(timeout_ticks).await
122 }
123
124 async fn wait_for_container_open(
125 &self,
126 timeout_ticks: Option<usize>,
127 ) -> Option<ContainerHandle> {
128 let mut ticks = self.get_tick_broadcaster();
129 let mut elapsed_ticks = 0;
130 while ticks.recv().await.is_ok() {
131 let ecs = self.ecs.lock();
132 if ecs.get::<WaitingForInventoryOpen>(self.entity).is_none() {
133 break;
134 }
135
136 elapsed_ticks += 1;
137 if let Some(timeout_ticks) = timeout_ticks
138 && elapsed_ticks >= timeout_ticks
139 {
140 return None;
141 }
142 }
143
144 let ecs = self.ecs.lock();
145 let inventory = ecs.get::<Inventory>(self.entity).expect("no inventory");
146 if inventory.id == 0 {
147 None
148 } else {
149 Some(ContainerHandle::new(inventory.id, self.clone()))
150 }
151 }
152
153 fn open_inventory(&self) -> Option<ContainerHandle> {
154 let ecs = self.ecs.lock();
155 let inventory = ecs.get::<Inventory>(self.entity).expect("no inventory");
156 if inventory.id == 0 {
157 Some(ContainerHandle::new(0, self.clone()))
158 } else {
159 None
160 }
161 }
162
163 fn get_inventory(&self) -> ContainerHandleRef {
164 self.query_self::<&Inventory, _>(|inv| ContainerHandleRef::new(inv.id, self.clone()))
165 }
166
167 fn get_held_item(&self) -> ItemStack {
168 self.query_self::<&Inventory, _>(|inv| inv.held_item())
169 }
170}
171
172pub struct ContainerHandleRef {
177 id: i32,
178 client: Client,
179}
180impl Debug for ContainerHandleRef {
181 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
182 f.debug_struct("ContainerHandle")
183 .field("id", &self.id())
184 .finish()
185 }
186}
187impl ContainerHandleRef {
188 pub fn new(id: i32, client: Client) -> Self {
189 Self { id, client }
190 }
191
192 pub fn close(&self) {
193 self.client.ecs.lock().trigger(CloseContainerEvent {
194 entity: self.client.entity,
195 id: self.id,
196 });
197 }
198
199 pub fn id(&self) -> i32 {
205 self.id
206 }
207
208 pub fn menu(&self) -> Option<Menu> {
216 let ecs = self.client.ecs.lock();
217 let inventory = ecs
218 .get::<Inventory>(self.client.entity)
219 .expect("no inventory");
220
221 if inventory.id == self.id {
223 if self.id == 0 {
224 Some(inventory.inventory_menu.clone())
225 } else {
226 Some(inventory.container_menu.clone().unwrap())
227 }
228 } else {
229 None
230 }
231 }
232
233 pub fn contents(&self) -> Option<Vec<ItemStack>> {
238 self.menu().map(|menu| menu.contents())
239 }
240
241 pub fn slots(&self) -> Option<Vec<ItemStack>> {
245 self.menu().map(|menu| menu.slots())
246 }
247
248 pub fn left_click(&self, slot: impl Into<usize>) {
250 self.click(PickupClick::Left {
251 slot: Some(slot.into() as u16),
252 });
253 }
254 pub fn shift_click(&self, slot: impl Into<usize>) {
256 self.click(QuickMoveClick::Left {
257 slot: slot.into() as u16,
258 });
259 }
260 pub fn right_click(&self, slot: impl Into<usize>) {
262 self.click(PickupClick::Right {
263 slot: Some(slot.into() as u16),
264 });
265 }
266
267 pub fn click(&self, operation: impl Into<ClickOperation>) {
270 let operation = operation.into();
271 self.client.ecs.lock().trigger(ContainerClickEvent {
272 entity: self.client.entity,
273 window_id: self.id,
274 operation,
275 });
276 }
277}
278
279#[derive(Deref)]
283pub struct ContainerHandle(ContainerHandleRef);
284
285impl Drop for ContainerHandle {
286 fn drop(&mut self) {
287 self.0.close();
288 }
289}
290impl Debug for ContainerHandle {
291 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
292 f.debug_struct("ContainerHandle")
293 .field("id", &self.id())
294 .finish()
295 }
296}
297impl ContainerHandle {
298 fn new(id: i32, client: Client) -> Self {
299 Self(ContainerHandleRef { id, client })
300 }
301
302 pub fn close(self) {
304 }
306}
307
308#[derive(Component, Debug)]
309pub struct WaitingForInventoryOpen;
310
311pub fn handle_menu_opened_event(
312 mut commands: Commands,
313 mut events: MessageReader<ReceiveGamePacketEvent>,
314) {
315 for event in events.read() {
316 if let ClientboundGamePacket::ContainerSetContent { .. } = event.packet.as_ref() {
317 commands
318 .entity(event.entity)
319 .remove::<WaitingForInventoryOpen>();
320 }
321 }
322}