From a48a83370754667454d03b0f0ce9c8252df54581 Mon Sep 17 00:00:00 2001
From: Andreas <andreasruden91@gmail.com>
Date: Tue, 24 Oct 2023 20:56:55 +0200
Subject: [PATCH] Add socketed items parsing

---
 d2warehouse/item.py                  | 12 +++++++++---
 d2warehouse/parser.py                | 26 ++++++++++++++++++--------
 d2warehouse/tests/test_parse_item.py | 17 ++++++++++-------
 3 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/d2warehouse/item.py b/d2warehouse/item.py
index fe56233..76b8595 100644
--- a/d2warehouse/item.py
+++ b/d2warehouse/item.py
@@ -17,6 +17,7 @@
 import json
 import os
 import re
+from typing import Optional
 from bitarray import bitarray
 from dataclasses import dataclass
 from enum import Enum
@@ -108,7 +109,7 @@ class Item:
     defense: int | None = None
     durability: int | None = None
     max_durability: int | None = None
-    sockets: int | None = None
+    sockets: list[Optional["Item"]] | None = None
     quantity: int | None = None
     stats: list[Stat] | None = None
 
@@ -157,8 +158,13 @@ class Item:
                 " " * indent,
                 f"Durability: {self.durability} out of {self.max_durability}",
             )
-        if self.sockets:
-            print(" " * indent, f"Num Sockets: {self.sockets}")
+        if self.is_socketed:
+            print(" " * indent, f"{len(self.sockets)} sockets:")
+            for socket in self.sockets:
+                if socket:
+                    socket.print(indent + 4)
+                else:
+                    print(" " * (indent + 4), "Empty")
         if self.quantity:
             print(" " * indent, f"Quantity: {self.quantity}")
         if self.stats:
diff --git a/d2warehouse/parser.py b/d2warehouse/parser.py
index 2440f5f..b5a77f9 100644
--- a/d2warehouse/parser.py
+++ b/d2warehouse/parser.py
@@ -128,8 +128,7 @@ def parse_item(data: bytes) -> tuple[bytes, Item]:
     code_end += 53
     # TODO: verify that this socket thing is really 1 bit for simple items...?
     sockets_end = code_end + 1 if is_simple else code_end + 3
-    sockets_count = ba2int(bits[code_end:sockets_end])
-    print("sockets", sockets_count)
+    filled_sockets = ba2int(bits[code_end:sockets_end])
     if is_ear:
         raise UnsupportedItemError("Ear items are not supported")
 
@@ -180,13 +179,28 @@ def parse_item(data: bytes) -> tuple[bytes, Item]:
     item, itembase_end = parse_basetype_data(bits[personalized_end:], item)
     itembase_end += personalized_end
 
-    item, enchantments_end = parse_enchantments(bits[itembase_end:], item)
+    if item.is_socketed:
+        sockets_count = ba2int(bits[itembase_end : itembase_end + 4])
+        sockets_end = itembase_end + 4
+    else:
+        sockets_end = itembase_end
+
+    item, enchantments_end = parse_enchantments(bits[sockets_end:], item)
     enchantments_end += itembase_end
 
     extended_byte_size = int((enchantments_end + 7) / 8)
 
     item.raw_data = data[:extended_byte_size]
-    return data[extended_byte_size:], item
+    remaining_data = data[extended_byte_size:]
+
+    # Parse out sockets if any exist on the item
+    if item.is_socketed:
+        item.sockets = [None] * sockets_count
+        for i in range(0, filled_sockets):
+            remaining_data, socket = parse_item(remaining_data)
+            item.sockets[i] = socket
+
+    return remaining_data, item
 
 
 def parse_item_graphic(bits: bitarray) -> tuple[int | None, int]:
@@ -319,10 +333,6 @@ def parse_basetype_data(bits: bitarray, item: Item) -> tuple[Item, int]:
         item.quantity = ba2int(bits[ptr : ptr + 9])
         ptr += 9
 
-    if item.is_socketed:
-        item.sockets = ba2int(bits[ptr : ptr + 4])
-        ptr += 4
-
     return item, ptr
 
 
diff --git a/d2warehouse/tests/test_parse_item.py b/d2warehouse/tests/test_parse_item.py
index 279d3f7..5d0e327 100644
--- a/d2warehouse/tests/test_parse_item.py
+++ b/d2warehouse/tests/test_parse_item.py
@@ -68,7 +68,7 @@ class ParseItemTest(unittest.TestCase):
         self.assertEqual(data, b"")
         self.assertEqual(item.quality, Quality.HIGH)
         self.assertEqual(len(item.stats), 2)
-        self.assertEqual(item.sockets, 2)
+        self.assertEqual(len(item.sockets), 2)
 
     def test_ed_max(self):
         # test bugfix for https://gitlab.com/omicron-oss/d2warehouse/-/issues/1
@@ -81,21 +81,24 @@ class ParseItemTest(unittest.TestCase):
         self.assertEqual(str(item.stats[0]), "+13% Enhanced Damage")
         self.assertEqual(str(item.stats[2]), "+2 to Maximum Damage")
 
-    def test_runeworld_lore(self):
+    def test_runeword_lore(self):
         # Lore: Ort Sol
         data = bytes.fromhex(
             "1008800405c055c637f1073af4558697412981070881506049e87f005516fb134582ff1000a0003500e07cbb001000a0003504e07c9800"
         )
         data, item = parse_item(data)
-        # TODO: requires socket parsing
-        # self.assertEqual(data, b"")
+        self.assertEqual(data, b"")
         self.assertTrue(item.is_runeword)
-        self.assertEqual(item.sockets, 2)
+        self.assertEqual(len(item.sockets), 2)
+        item.sockets[0].print()
+        item.sockets[1].print()
+        self.assertEqual(item.sockets[0].code, "r09")
+        self.assertEqual(item.sockets[1].code, "r12")
         rw = lookup_runeword(item.runeword_id)
         self.assertEqual(rw["name"], "Lore")
         for stat in item.stats:
             print(str(stat))
         self.assertEqual(str(item.stats[4]), "+1 to All Skills")  # runeword stat
         self.assertEqual(str(item.stats[2]), "+10 to Energy")  # runeword stat
-        # TODO: requires socket parsing
-        # self.assertEqual(str(item.stats[1]), "Lightning Resist 30%")  # sol rune stat
+        # TODO: sol rune stat -- should it be in representation?
+        # self.assertEqual(str(item.sockets[1].stats[0]), "Lightning Resist 30%")