1 files changed, 33 insertions, 18 deletions

diff --git a/codegen/migrate.py b/codegen/migrate.py
index c0748400..6928cea1 100644
--- a/codegen/migrate.py
+++ b/codegen/migrate.py
@@ -1,5 +1,7 @@
+from codegen.lib.utils import PacketIdentifier, group_packets
 import lib.code.utils
 import lib.code.version
+import lib.code.packet
 import lib.download
 import sys
 import os
@@ -14,39 +16,52 @@ new_mappings = lib.download.get_mappings_for_version(new_version_id)
 new_burger_data = lib.download.get_burger_data_for_version(new_version_id)
 new_packet_list = list(new_burger_data[0]['packets']['packet'].values())
 
-old_packet_ids = {}
-new_packet_ids = {}
+
+old_packets: dict[PacketIdentifier, str] = {}
+new_packets: dict[PacketIdentifier, str] = {}
 
 for packet in old_packet_list:
     assert packet['class'].endswith('.class')
     packet_name = old_mappings.get_class(packet['class'][:-6])
-    old_packet_ids[packet_name] = packet['id']
+    old_packets[PacketIdentifier(
+        packet['id'], packet['direction'], packet['state'])] = packet_name
 for packet in new_packet_list:
     assert packet['class'].endswith('.class')
     packet_name = new_mappings.get_class(packet['class'][:-6])
-    new_packet_ids[packet_name] = packet['id']
+    new_packets[PacketIdentifier(
+        packet['id'], packet['direction'], packet['state'])] = packet_name
 
-# find packets that changed ids
-for packet_name in old_packet_ids:
-    if packet_name in new_packet_ids:
-        if old_packet_ids[packet_name] != new_packet_ids[packet_name]:
-            print(packet_name, 'id changed from',
-                  old_packet_ids[packet_name], 'to', new_packet_ids[packet_name])
+
+# find removed packets
+removed_packets: list[PacketIdentifier] = []
+for packet in old_packets:
+    if packet not in new_packets:
+        removed_packets.append(packet)
+for (direction, state), packets in group_packets(removed_packets).items():
+    lib.code.packet.remove_packet_ids(packets, direction, state)
 
 print()
 
-# find removed packets
-for packet_name in old_packet_ids:
-    if packet_name not in new_packet_ids:
-        print(packet_name, 'removed')
+# find packets that changed ids
+changed_packets: dict[PacketIdentifier, int] = {}
+for old_packet, old_packet_name in old_packets.items():
+    for new_packet, new_packet_name in new_packets.items():
+        if old_packet == new_packet and old_packet.packet_id != new_packet.packet_id:
+            changed_packets[old_packet] = new_packet.packet_id
+for (direction, state), packets in group_packets(list(changed_packets.keys())).items():
+    lib.code.packet.remove_packet_ids(packets, direction, state)
+
 
 print()
 
 # find added packets
-for packet_name in new_packet_ids:
-    if packet_name not in old_packet_ids:
-        print(packet_name, 'added')
-
+added_packets: list[PacketIdentifier] = []
+for packet in new_packets:
+    if packet not in old_packets:
+        added_packets.append(packet)
+for packet in added_packets:
+    lib.code.packet.generate_packet(
+        new_burger_data, new_mappings, packet.packet_id, packet.direction, packet.state)
 lib.code.utils.fmt()
 
 print('Done!')