Implement ALL packets (#16)

* add a couple more packets and improve codegen

* enums in packet codegen

* fix enums and MORE PACKETS

* make unsigned numbers the default

* codegen can make hashmaps

* UnsizedByteArray in codegen

* Vec and Option

* enum codgen works in more situations

* ServerboundInteractPacket

* Fix error with new error system

* More packets

* more packets

* more packets

* guess what was added

* yeah it's more packets

* add more packets

* packets

* start adding ClientboundBossEventPacket

* finish boss event packet

* improve codegen for linux

* start on command suggestions packet

* rename declare_commands to commands

* más paquetes

* fix generating custom payload packet

* more packets

* mehr Pakete

* improve codegen for movement packets

* rename move packets to have "packet" at the end

* fix some unused variable warns

* addere plus facis

* pli da pakoj

* plus de paquets

* più pacchetti

* make ChatFormatting a macro in azalea-chat

* change a match to matches! macro

* update SetPlayerTeam to use ChatFormatting

* ClientboundSetScorePacket & fix clippy warnings

* finish game state :tada:

* add remaining packets for other states

* fix error in ping.rs

2 files changed, 247 insertions, 44 deletions

diff --git a/codegen/lib/code/packet.py b/codegen/lib/code/packet.py
index ffa7841c..692a449e 100644
--- a/codegen/lib/code/packet.py
+++ b/codegen/lib/code/packet.py
@@ -27,27 +27,40 @@ def generate_packet(burger_packets, mappings: Mappings, target_packet_id, target
 
         generated_packet_code = []
         uses = set()
+        extra_code = []
+
+        packet_derive_name = f'{to_camel_case(direction)}{to_camel_case(state)}Packet'
+
         generated_packet_code.append(
-            f'#[derive(Clone, Debug, McBuf, {to_camel_case(state)}Packet)]')
-        uses.add(f'packet_macros::{to_camel_case(state)}Packet')
+            f'#[derive(Clone, Debug, McBuf, {packet_derive_name})]')
+        uses.add(f'packet_macros::{packet_derive_name}')
         uses.add(f'azalea_buf::McBuf')
 
         obfuscated_class_name = packet['class'].split('.')[0]
         class_name = mappings.get_class(
             obfuscated_class_name).split('.')[-1]
         if '$' in class_name:
-            class_name = class_name.replace('$', '')
+            class_name, extra_part = class_name.split('$')
+            if class_name.endswith('Packet'):
+                class_name = class_name[:-
+                                        len('Packet')] + extra_part + 'Packet'
 
         generated_packet_code.append(
             f'pub struct {to_camel_case(class_name)} {{')
 
-        for instruction in packet.get('instructions', []):
-            if instruction['operation'] == 'write':
-                burger_instruction_to_code(
-                    instruction, generated_packet_code, mappings, obfuscated_class_name, uses)
+        # call burger_instruction_to_code for each instruction
+        i = -1
+        instructions = packet.get('instructions', [])
+        while (i + 1) < len(instructions):
+            i += 1
+
+            if instructions[i]['operation'] == 'write':
+                skip = burger_instruction_to_code(
+                    instructions, i, generated_packet_code, mappings, obfuscated_class_name, uses, extra_code)
+                if skip:
+                    i += skip
             else:
-                generated_packet_code.append(f'// TODO: {instruction}')
-                continue
+                generated_packet_code.append(f'// TODO: {instructions[i]}')
 
         generated_packet_code.append('}')
 
@@ -56,6 +69,8 @@ def generate_packet(burger_packets, mappings: Mappings, target_packet_id, target
             generated_packet_code.insert(0, '')
         for use in uses:
             generated_packet_code.insert(0, f'use {use};')
+        for line in extra_code:
+            generated_packet_code.append(line)
 
         print(generated_packet_code)
         write_packet_file(state, to_snake_case(class_name),
@@ -204,21 +219,108 @@ def get_packets(direction: str, state: str):
     return packet_ids, packet_class_names
 
 
-def burger_instruction_to_code(instruction: dict, generated_packet_code: list[str], mappings: Mappings, obfuscated_class_name: str, uses: set):
-    field_type = instruction['type']
-    field_type_rs, is_var, instruction_uses = burger_type_to_rust_type(
-        field_type)
-
-    obfuscated_field_name = instruction['field']
-    if '.' in obfuscated_field_name or ' ' in obfuscated_field_name or '(' in obfuscated_field_name:
-        field_type_rs, obfuscated_field_name = burger_field_to_type(
-            obfuscated_field_name)
-        if not field_type_rs:
-            generated_packet_code.append(f'// TODO: {instruction}')
-            return
-    field_name = mappings.get_field(
-        obfuscated_class_name, obfuscated_field_name) or mappings.get_field(
-        obfuscated_class_name.split('$')[0], obfuscated_field_name)
+def burger_instruction_to_code(instructions: list[dict], index: int, generated_packet_code: list[str], mappings: Mappings, obfuscated_class_name: str, uses: set, extra_code: list[str]) -> Optional[int]:
+    '''
+    Generate a field for an instruction, returns the number of instructions to skip (if any).
+    '''
+    instruction = instructions[index]
+    next_instruction = instructions[index +
+                                    1] if index + 1 < len(instructions) else None
+    next_next_instruction = instructions[index +
+                                         2] if index + 2 < len(instructions) else None
+
+    is_var = False
+    skip = 0
+    field_type_rs = None
+    field_comment = None
+
+    # iterators
+    if instruction['operation'] == 'write' and instruction['field'].endswith('.size()') and next_instruction and next_instruction['type'] == 'Iterator' and next_next_instruction and next_next_instruction['operation'] == 'loop':
+        field_obfuscated_name = instruction['field'].split('.')[
+            0]
+        field_name = mappings.get_field(
+            obfuscated_class_name, field_obfuscated_name)
+
+        # figure out what kind of iterator it is
+        loop_instructions = next_next_instruction['instructions']
+        if len(loop_instructions) == 2:
+            entry_type_rs, is_var, uses, extra_code = burger_type_to_rust_type(
+                loop_instructions[1]['type'], None, loop_instructions[1], mappings, obfuscated_class_name)
+            field_type_rs = f'Vec<{entry_type_rs}>'
+        elif len(loop_instructions) == 3:
+            is_map = loop_instructions[0]['type'].startswith(
+                'Map.Entry<')
+            if is_map:
+                assert loop_instructions[1]['field'].endswith(
+                    '.getKey()')
+                assert loop_instructions[2]['field'].endswith(
+                    '.getValue()')
+
+                # generate the type for the key
+                key_type_rs, is_key_var, key_uses, key_extra_code = burger_type_to_rust_type(
+                    loop_instructions[1]['type'], None, loop_instructions[1], mappings, obfuscated_class_name)
+                uses.update(key_uses)
+                extra_code.extend(key_extra_code)
+
+                # generate the type for the value
+                value_type_rs, is_value_var, value_uses, value_extra_code = burger_type_to_rust_type(
+                    loop_instructions[2]['type'], None, loop_instructions[2], mappings, obfuscated_class_name)
+                uses.update(value_uses)
+                extra_code.extend(value_extra_code)
+
+                field_type_rs = f'HashMap<{key_type_rs}, {value_type_rs}>'
+                uses.add('std::collections::HashMap')
+
+                # only the key is var since the value can be made var in other ways
+                is_var = is_key_var
+
+        skip = 2  # skip the next 2 instructions
+
+    # Option<T>
+    elif instruction['operation'] == 'write' and (instruction['field'].endswith('.isPresent()') or instruction['field'].endswith(' != null')) and next_instruction and (next_instruction.get('condition', '').endswith('.isPresent()') or next_instruction.get('condition', '').endswith(' != null')):
+        field_obfuscated_name = instruction['field'].split('.')[
+            0].split(' ')[0]
+        field_name = mappings.get_field(
+            obfuscated_class_name, field_obfuscated_name)
+        condition_instructions = next_instruction['instructions']
+
+        condition_types_rs = []
+        for condition_instruction in condition_instructions:
+            condition_type_rs, is_var, this_uses, this_extra_code = burger_type_to_rust_type(
+                condition_instruction['type'], None, condition_instruction, mappings, obfuscated_class_name)
+            condition_types_rs.append(condition_type_rs)
+            uses.update(this_uses)
+            extra_code.extend(this_extra_code)
+        field_type_rs = f'Option<({", ".join(condition_types_rs)})>' if len(
+            condition_types_rs) != 1 else f'Option<{condition_types_rs[0]}>'
+        skip = 1
+    else:
+        field_type = instruction['type']
+        obfuscated_field_name = instruction['field']
+
+        if obfuscated_field_name.startswith('(float)'):
+            obfuscated_field_name = obfuscated_field_name[len('(float)'):]
+
+        field_name = mappings.get_field(
+            obfuscated_class_name, obfuscated_field_name) or mappings.get_field(
+            obfuscated_class_name.split('$')[0], obfuscated_field_name)
+
+        field_type_rs, is_var, instruction_uses, instruction_extra_code = burger_type_to_rust_type(
+            field_type, field_name, instruction, mappings, obfuscated_class_name)
+
+        if '.' in obfuscated_field_name or ' ' in obfuscated_field_name or '(' in obfuscated_field_name:
+            field_type_rs2, obfuscated_field_name, field_comment = burger_field_to_type(
+                obfuscated_field_name, mappings, obfuscated_class_name)
+            if not field_type_rs2:
+                generated_packet_code.append(f'// TODO: {instruction}')
+                return
+            # try to get the field name again with the new stuff we know
+            field_name = mappings.get_field(
+                obfuscated_class_name, obfuscated_field_name) or mappings.get_field(
+                obfuscated_class_name.split('$')[0], obfuscated_field_name)
+        uses.update(instruction_uses)
+        extra_code.extend(instruction_extra_code)
+
     if not field_name:
         generated_packet_code.append(
             f'// TODO: unknown field {instruction}')
@@ -226,17 +328,44 @@ def burger_instruction_to_code(instruction: dict, generated_packet_code: list[st
 
     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)
+    line = f'pub {to_snake_case(field_name)}: {field_type_rs or "todo!()"},'
+    if field_comment:
+        line += f' // {field_comment}'
+    generated_packet_code.append(line)
 
+    return skip
 
-def burger_field_to_type(field) -> tuple[Optional[str], str]:
+
+def burger_field_to_type(field, mappings: Mappings, obfuscated_class_name: str) -> tuple[Optional[str], str, Optional[str]]:
+    '''
+    Returns field_type_rs, obfuscated_field_name, field_comment
+    '''
     # match `(x) ? 1 : 0`
     match = re.match(r'\((.*)\) \? 1 : 0', field)
     if match:
-        return ('bool', match.group(1))
-    return None, field
+        return ('bool', match.group(1), None)
+    match = re.match(r'^\w+\.\w+\(\)$', field)
+    if match:
+        print('field', field)
+        obfuscated_first = field.split('.')[0]
+        obfuscated_second = field.split('.')[1].split('(')[0]
+        first = mappings.get_field(obfuscated_class_name, obfuscated_first)
+        first_type = mappings.get_field_type(
+            obfuscated_class_name, obfuscated_first)
+        first_obfuscated_class_name: Optional[str] = mappings.get_class_from_deobfuscated_name(
+            first_type)
+        if first_obfuscated_class_name:
+            try:
+                second = mappings.get_method(
+                    first_obfuscated_class_name, obfuscated_second, '')
+            except:
+                # if this happens then the field is probably from a super class
+                second = obfuscated_second
+        else:
+            second = obfuscated_second
+        first_type_short = first_type.split('.')[-1]
+        return (first_type_short, obfuscated_first, f'TODO: Does {first_type_short}::{second}, may not be implemented')
+    return None, field, None
 
 
 def change_packet_ids(id_map: dict[int, int], direction: str, state: str):
diff --git a/codegen/lib/code/utils.py b/codegen/lib/code/utils.py
index 0c22d7ba..e4671488 100644
--- a/codegen/lib/code/utils.py
+++ b/codegen/lib/code/utils.py
@@ -1,22 +1,31 @@
 
-from lib.utils import get_dir_location
+from lib.utils import to_camel_case, to_snake_case, get_dir_location
+from lib.mappings import Mappings
+from typing import Optional
 import os
 
 # utilities specifically for codegen
 
 
-def burger_type_to_rust_type(burger_type):
+def burger_type_to_rust_type(burger_type, field_name: Optional[str] = None, instruction=None, mappings: Optional[Mappings] = None, obfuscated_class_name: Optional[str] = None):
     is_var = False
     uses = set()
+    # extra code, like enum definitions
+    extra_code = []
+
+    should_be_signed = False
+    if field_name and any(map(lambda w: w in {'x', 'y', 'z', 'xa', 'ya', 'za'}, to_snake_case(field_name).split('_'))):
+        # coordinates are signed
+        should_be_signed = True
 
     if burger_type == 'byte':
-        field_type_rs = 'i8'
+        field_type_rs = 'i8' if should_be_signed else 'u8'
     elif burger_type == 'short':
-        field_type_rs = 'i16'
+        field_type_rs = 'i16' if should_be_signed else 'u16'
     elif burger_type == 'int':
-        field_type_rs = 'i32'
+        field_type_rs = 'i32' if should_be_signed else 'u32'
     elif burger_type == 'long':
-        field_type_rs = 'i64'
+        field_type_rs = 'i64' if should_be_signed else 'u64'
     elif burger_type == 'float':
         field_type_rs = 'f32'
     elif burger_type == 'double':
@@ -24,10 +33,10 @@ def burger_type_to_rust_type(burger_type):
 
     elif burger_type == 'varint':
         is_var = True
-        field_type_rs = 'i32'
+        field_type_rs = 'i32' if should_be_signed else 'u32'
     elif burger_type == 'varlong':
         is_var = True
-        field_type_rs = 'i64'
+        field_type_rs = 'i64' if should_be_signed else 'u64'
 
     elif burger_type == 'boolean':
         field_type_rs = 'bool'
@@ -39,7 +48,7 @@ def burger_type_to_rust_type(burger_type):
         uses.add('azalea_chat::component::Component')
     elif burger_type == 'identifier':
         field_type_rs = 'ResourceLocation'
-        uses.add('azalea_core::resource_location::ResourceLocation')
+        uses.add('azalea_core::ResourceLocation')
     elif burger_type == 'uuid':
         field_type_rs = 'Uuid'
         uses.add('uuid::Uuid')
@@ -53,17 +62,82 @@ def burger_type_to_rust_type(burger_type):
         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
+        uses.add('azalea_entity::EntityMetadata')
+    elif burger_type == 'abstract':
         field_type_rs = 'todo!()'
+    elif burger_type == 'enum':
+        if not instruction or not mappings or not obfuscated_class_name:
+            field_type_rs = 'todo!("enum")'
+        else:
+            # generate the whole enum :)
+            print(instruction)
+            enum_field = instruction['field']
+            # enums with a.b() as the field
+            if '.' in enum_field:
+                enum_first_part_name = mappings.get_field_type(
+                    obfuscated_class_name, enum_field.split('.')[0])
+                enum_first_part_obfuscated_name = mappings.get_class_from_deobfuscated_name(
+                    enum_first_part_name)
+                print('enum_first_part_obfuscated_name',
+                      enum_first_part_obfuscated_name)
+                enum_name = mappings.get_method_type(
+                    enum_first_part_obfuscated_name, enum_field.split('.')[1].split('(')[0], '')
+
+                print('hm', enum_name)
+            else:
+                enum_name = mappings.get_field_type(
+                    obfuscated_class_name, enum_field)
+            print('enum_name', enum_name)
+            enum_obfuscated_name = mappings.get_class_from_deobfuscated_name(
+                enum_name)
+            print('enum_obfuscated_name', enum_obfuscated_name)
+            enum_variants = []
+            for obfuscated_field_name in mappings.fields[enum_obfuscated_name]:
+                field_name = mappings.get_field(
+                    enum_obfuscated_name, obfuscated_field_name)
+
+                # get the type just to make sure it's actually a variant and not something else
+                field_type = mappings.get_field_type(
+                    enum_obfuscated_name, obfuscated_field_name)
+                if field_type != enum_name:
+                    continue
+
+                enum_variants.append(field_name)
+
+            field_type_rs = to_camel_case(
+                enum_name.split('.')[-1].split('$')[-1])
+            extra_code.append('')
+            extra_code.append(f'#[derive(McBuf, Clone, Copy, Debug)]')
+            extra_code.append(f'pub enum {field_type_rs} {{')
+            for index, variant in enumerate(enum_variants):
+                extra_code.append(
+                    f'    {to_camel_case(variant.lower())}={index},')
+            extra_code.append('}')
+
     elif burger_type.endswith('[]'):
-        field_type_rs, is_var, uses = burger_type_to_rust_type(
+        field_type_rs, is_var, uses, extra_code = burger_type_to_rust_type(
             burger_type[:-2])
         field_type_rs = f'Vec<{field_type_rs}>'
+
+        # sometimes burger gives us a slightly incorrect type
+        if mappings and instruction:
+            if field_type_rs == 'Vec<u8>':
+                field = instruction['field']
+                if field.endswith('.copy()'):
+                    field = field[:-7]
+                try:
+                    array_type = mappings.get_field_type(
+                        obfuscated_class_name, field)
+                except KeyError:
+                    print('Error getting array type', field)
+                    return field_type_rs, is_var, uses, extra_code
+                if array_type == 'net.minecraft.network.FriendlyByteBuf':
+                    field_type_rs = 'UnsizedByteArray'
+                    uses.add('azalea_buf::UnsizedByteArray')
+
     else:
         raise Exception(f'Unknown field type: {burger_type}')
-    return field_type_rs, is_var, uses
+    return field_type_rs, is_var, uses, extra_code
 
 
 def write_packet_file(state, packet_name_snake_case, code):