fixes

1 files changed, 26 insertions, 3 deletions

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))