1use std::{
2    collections::HashMap,
3    io::{self, Write},
4    sync::Arc,
5};
6
7use byteorder::{BigEndian, WriteBytesExt};
8use indexmap::IndexMap;
9
10use super::{MAX_STRING_LENGTH, UnsizedByteArray};
11
12fn write_utf_with_len(buf: &mut impl Write, string: &str, len: usize) -> io::Result<()> {
13    if string.len() > len {
14        panic!(
15            "String too big (was {} bytes encoded, max {})",
16            string.len(),
17            len
18        );
19    }
20    string.as_bytes().to_vec().azalea_write(buf)?;
21    Ok(())
22}
23
24pub trait AzaleaWrite {
25    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()>;
26}
27
28pub trait AzaleaWriteVar {
29    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()>;
30}
31
32impl AzaleaWrite for () {
33    fn azalea_write(&self, _buf: &mut impl Write) -> io::Result<()> {
34        Ok(())
35    }
36}
37
38impl AzaleaWrite for i32 {
39    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
40        WriteBytesExt::write_i32::<BigEndian>(buf, *self)
41    }
42}
43
44impl AzaleaWriteVar for i32 {
45    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
46        let mut buffer = [0];
47        let mut value = *self;
48        if value == 0 {
49            buf.write_all(&buffer)?;
50        }
51        while value != 0 {
52            buffer[0] = (value & 0b0111_1111) as u8;
53            value = (value >> 7) & (i32::MAX >> 6);
54            if value != 0 {
55                buffer[0] |= 0b1000_0000;
56            }
57            buf.write_all(&buffer)?;
58        }
59        Ok(())
60    }
61}
62
63impl AzaleaWrite for UnsizedByteArray {
64    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
65        buf.write_all(self)
66    }
67}
68
69impl<T: AzaleaWrite> AzaleaWrite for Vec<T> {
70    default fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
71        self[..].azalea_write(buf)
72    }
73}
74impl<T: AzaleaWrite> AzaleaWrite for Box<[T]> {
75    default fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
76        self[..].azalea_write(buf)
77    }
78}
79
80impl<T: AzaleaWrite> AzaleaWrite for [T] {
81    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
82        (self.len() as u32).azalea_write_var(buf)?;
83        for item in self {
84            T::azalea_write(item, buf)?;
85        }
86        Ok(())
87    }
88}
89
90macro_rules! impl_for_map_type {
91    ($ty: ident) => {
92        impl<K: AzaleaWrite, V: AzaleaWrite> AzaleaWrite for $ty<K, V> {
93            fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
94                u32::azalea_write_var(&(self.len() as u32), buf)?;
95                for (key, value) in self {
96                    key.azalea_write(buf)?;
97                    value.azalea_write(buf)?;
98                }
99
100                Ok(())
101            }
102        }
103        impl<K: AzaleaWrite, V: AzaleaWriteVar> AzaleaWriteVar for $ty<K, V> {
104            fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
105                u32::azalea_write_var(&(self.len() as u32), buf)?;
106                for (key, value) in self {
107                    key.azalea_write(buf)?;
108                    value.azalea_write_var(buf)?;
109                }
110
111                Ok(())
112            }
113        }
114    };
115}
116
117impl_for_map_type!(HashMap);
118impl_for_map_type!(IndexMap);
119
120impl AzaleaWrite for Vec<u8> {
121    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
122        (self.len() as u32).azalea_write_var(buf)?;
123        buf.write_all(self)
124    }
125}
126
127impl AzaleaWrite for String {
128    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
129        write_utf_with_len(buf, self, MAX_STRING_LENGTH.into())
130    }
131}
132
133impl AzaleaWrite for &str {
134    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
135        write_utf_with_len(buf, self, MAX_STRING_LENGTH.into())
136    }
137}
138
139impl AzaleaWrite for u32 {
140    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
141        i32::azalea_write(&(*self as i32), buf)
142    }
143}
144
145impl AzaleaWriteVar for u32 {
146    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
147        i32::azalea_write_var(&(*self as i32), buf)
148    }
149}
150
151impl AzaleaWriteVar for i64 {
152    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
153        let mut buffer = [0];
154        let mut value = *self;
155        if value == 0 {
156            buf.write_all(&buffer).unwrap();
157        }
158        while value != 0 {
159            buffer[0] = (value & 0b0111_1111) as u8;
160            value = (value >> 7) & (i64::MAX >> 6);
161            if value != 0 {
162                buffer[0] |= 0b1000_0000;
163            }
164            buf.write_all(&buffer)?;
165        }
166        Ok(())
167    }
168}
169
170impl AzaleaWriteVar for u64 {
171    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
172        i64::azalea_write_var(&(*self as i64), buf)
173    }
174}
175
176impl AzaleaWrite for u16 {
177    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
178        i16::azalea_write(&(*self as i16), buf)
179    }
180}
181
182impl AzaleaWriteVar for u16 {
183    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
184        i32::azalea_write_var(&(*self as i32), buf)
185    }
186}
187
188impl<T: AzaleaWriteVar> AzaleaWriteVar for [T] {
189    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
190        u32::azalea_write_var(&(self.len() as u32), buf)?;
191        for i in self {
192            i.azalea_write_var(buf)?;
193        }
194        Ok(())
195    }
196}
197impl<T: AzaleaWriteVar> AzaleaWriteVar for Vec<T> {
198    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
199        self[..].azalea_write_var(buf)
200    }
201}
202impl<T: AzaleaWriteVar> AzaleaWriteVar for Box<[T]> {
203    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
204        self[..].azalea_write_var(buf)
205    }
206}
207
208impl AzaleaWrite for u8 {
209    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
210        WriteBytesExt::write_u8(buf, *self)
211    }
212}
213
214impl AzaleaWrite for i16 {
215    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
216        WriteBytesExt::write_i16::<BigEndian>(buf, *self)
217    }
218}
219
220impl AzaleaWrite for i64 {
221    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
222        WriteBytesExt::write_i64::<BigEndian>(buf, *self)
223    }
224}
225
226impl AzaleaWrite for u64 {
227    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
228        i64::azalea_write(&(*self as i64), buf)
229    }
230}
231
232impl AzaleaWrite for bool {
233    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
234        let byte = u8::from(*self);
235        byte.azalea_write(buf)
236    }
237}
238
239impl AzaleaWrite for i8 {
240    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
241        (*self as u8).azalea_write(buf)
242    }
243}
244
245impl AzaleaWrite for f32 {
246    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
247        WriteBytesExt::write_f32::<BigEndian>(buf, *self)
248    }
249}
250
251impl AzaleaWrite for f64 {
252    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
253        WriteBytesExt::write_f64::<BigEndian>(buf, *self)
254    }
255}
256
257impl<T: AzaleaWrite> AzaleaWrite for Option<T> {
258    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
259        if let Some(s) = self {
260            true.azalea_write(buf)?;
261            s.azalea_write(buf)?;
262        } else {
263            false.azalea_write(buf)?;
264        };
265        Ok(())
266    }
267}
268
269impl<T: AzaleaWriteVar> AzaleaWriteVar for Option<T> {
270    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
271        if let Some(s) = self {
272            true.azalea_write(buf)?;
273            s.azalea_write_var(buf)?;
274        } else {
275            false.azalea_write(buf)?;
276        };
277        Ok(())
278    }
279}
280
281impl<T: AzaleaWrite, const N: usize> AzaleaWrite for [T; N] {
283    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
284        for i in self {
285            i.azalea_write(buf)?;
286        }
287        Ok(())
288    }
289}
290
291impl AzaleaWrite for simdnbt::owned::NbtTag {
292    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
293        let mut data = Vec::new();
294        self.write(&mut data);
295        buf.write_all(&data)
296    }
297}
298
299impl AzaleaWrite for simdnbt::owned::NbtCompound {
300    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
301        let mut data = Vec::new();
302        simdnbt::owned::NbtTag::Compound(self.clone()).write(&mut data);
303        buf.write_all(&data)
304    }
305}
306
307impl AzaleaWrite for simdnbt::owned::Nbt {
308    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
309        let mut data = Vec::new();
310        self.write_unnamed(&mut data);
311        buf.write_all(&data)
312    }
313}
314
315impl<T> AzaleaWrite for Box<T>
316where
317    T: AzaleaWrite,
318{
319    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
320        T::azalea_write(&**self, buf)
321    }
322}
323
324impl<A: AzaleaWrite, B: AzaleaWrite> AzaleaWrite for (A, B) {
325    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
326        self.0.azalea_write(buf)?;
327        self.1.azalea_write(buf)
328    }
329}
330
331impl<T: AzaleaWrite> AzaleaWrite for Arc<T> {
332    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
333        T::azalea_write(&**self, buf)
334    }
335}