From 7b61d41f867090a120ca2493d4a499e3d08b018f Mon Sep 17 00:00:00 2001
From: mat <github@matdoes.dev>
Date: Sat, 7 May 2022 17:29:51 -0500
Subject: [PATCH] fixes

---
 data-code-generator/main.py          |  2 ++
 data-code-generator/packetcodegen.py | 29 +++++++++++++++++++++++++---
 data-code-generator/utils.py         |  2 +-
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/data-code-generator/main.py b/data-code-generator/main.py
index 50550d76..b79571ae 100644
--- a/data-code-generator/main.py
+++ b/data-code-generator/main.py
@@ -39,3 +39,5 @@ with open('burger.json', 'r') as f:
 
 burger_packets_data = burger_data[0]['packets']['packet']
 packetcodegen.generate(burger_packets_data, mappings)
+
+os.system('cd .. && cargo fmt')
diff --git a/data-code-generator/packetcodegen.py b/data-code-generator/packetcodegen.py
index 6d67f7c1..ccbb3845 100644
--- a/data-code-generator/packetcodegen.py
+++ b/data-code-generator/packetcodegen.py
@@ -1,5 +1,6 @@
 from utils import to_snake_case, to_camel_case
 from mappings import Mappings
+import os
 
 
 def burger_type_to_rust_type(burger_type):
@@ -33,18 +34,24 @@ def burger_type_to_rust_type(burger_type):
 
     elif burger_type == 'chatcomponent':
         field_type_rs = 'Component'
+        uses.add('azalea_chat::component::Component')
     elif burger_type == 'identifier':
         field_type_rs = 'ResourceLocation'
+        uses.add('azalea_core::resource_location::ResourceLocation')
     elif burger_type == 'uuid':
         field_type_rs = 'Uuid'
+        uses.add('uuid::Uuid')
     elif burger_type == 'position':
         field_type_rs = 'BlockPos'
+        uses.add('azalea_core::BlockPos')
     elif burger_type == 'nbtcompound':
         field_type_rs = 'azalea_nbt::Tag'
     elif burger_type == 'itemstack':
         field_type_rs = 'Slot'
+        uses.add('azalea_core::Slot')
     elif burger_type == 'metadata':
         field_type_rs = 'EntityMetadata'
+        uses.add('crate::mc_buf::EntityMetadata')
     elif burger_type == 'enum':
         # enums are too complicated, leave those to the user
         field_type_rs = 'todo!()'
@@ -59,18 +66,26 @@ def burger_type_to_rust_type(burger_type):
 
 
 def write_packet_file(state, packet_name_snake_case, code):
-    with open(f'../azalea-protocol/src/packets/{state}/{packet_name_snake_case}.rs', 'w') as f:
+    path = os.path.join(
+        '..', f'azalea-protocol/src/packets/{state}/{packet_name_snake_case}.rs')
+    with open(path, 'w') as f:
         f.write(code)
 
 
 def generate(burger_packets, mappings: Mappings):
+    packet_ids = [75]
     for packet in burger_packets.values():
+        if packet['id'] not in packet_ids:
+            continue
+
         direction = packet['direction'].lower()  # serverbound or clientbound
-        state = packet['state'].lower()
+        state = {'PLAY': 'game'}.get(packet['state'], packet['state'].lower())
 
         generated_packet_code = []
+        uses = set()
         generated_packet_code.append(
             f'#[derive(Clone, Debug, {to_camel_case(state)}Packet)]')
+        uses.add(f'packet_macros::{to_camel_case(state)}Packet')
 
         obfuscated_class_name = packet['class'].split('.')[0]
         class_name = mappings.get_class(obfuscated_class_name).split('.')[-1]
@@ -90,17 +105,25 @@ def generate(burger_packets, mappings: Mappings):
                     continue
 
                 field_type = instruction['type']
-                field_type_rs, is_var, uses = burger_type_to_rust_type(
+                field_type_rs, is_var, instruction_uses = burger_type_to_rust_type(
                     field_type)
                 if is_var:
                     generated_packet_code.append('#[var]')
                 generated_packet_code.append(
                     f'pub {to_snake_case(field_name)}: {field_type_rs},')
+                uses.update(instruction_uses)
             else:
                 generated_packet_code.append(f'// TODO: {instruction}')
                 continue
 
         generated_packet_code.append('}')
+
+        if uses:
+            # empty line before the `use` statements
+            generated_packet_code.insert(0, '')
+        for use in uses:
+            generated_packet_code.insert(0, f'use {use};')
+
         print(generated_packet_code)
         write_packet_file(state, to_snake_case(class_name),
                           '\n'.join(generated_packet_code))
diff --git a/data-code-generator/utils.py b/data-code-generator/utils.py
index 1052bb87..5336d574 100644
--- a/data-code-generator/utils.py
+++ b/data-code-generator/utils.py
@@ -7,7 +7,7 @@ import io
 
 def to_snake_case(name):
     s = re.sub('([A-Z])', r'_\1', name)
-    return s.lower()
+    return s.lower().strip('_')
 
 
 def to_camel_case(name):