From e99d93a6a02971621004b064deb31e1b061c334e Mon Sep 17 00:00:00 2001
From: Andreas <andreasruden91@gmail.com>
Date: Sun, 22 Oct 2023 13:03:38 +0200
Subject: [PATCH] Add item basetype parsing

---
 d2warehouse/item.py   | 21 ++++++++++++++++++++-
 d2warehouse/parser.py | 32 ++++++++++++++++++++++++++++++--
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/d2warehouse/item.py b/d2warehouse/item.py
index c10bf67..5d118e4 100644
--- a/d2warehouse/item.py
+++ b/d2warehouse/item.py
@@ -1,7 +1,12 @@
+import json
+import os
 from bitarray import bitarray
 from dataclasses import dataclass
 from enum import Enum
 
+_data_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+_basetype_map = None
+
 
 class Quality(Enum):
     LOW = 1
@@ -73,10 +78,16 @@ class Item:
     nameword2: int | None = None
     runeword_id: int | None = None
     personal_name: str | None = None
+    defense: int | None = None
+    durability: int | None = None
+    max_durability: int | None = None
+    sockets: int | None = None
+    quantity: int | None = None
 
     def print(self, indent=5, with_raw=False):
         properties = []
-        print(" " * indent, self.kind)
+        kind_name = lookup_basetype(self.kind)["name"]
+        print(" " * indent, f"{kind_name} ({self.kind})")
         if self.lvl:
             print(" " * indent, f"ilvl {self.lvl}")
         if self.is_simple:
@@ -117,3 +128,11 @@ class Item:
             print(" " * indent, txtbits(bits))
         print("")
         print("")
+
+
+def lookup_basetype(code: str) -> dict:
+    global _basetype_map
+    if _basetype_map is None:
+        with open(os.path.join(_data_path, "items.json")) as f:
+            _basetype_map = json.load(f)
+    return _basetype_map[code]
diff --git a/d2warehouse/parser.py b/d2warehouse/parser.py
index 294f2f8..818de55 100644
--- a/d2warehouse/parser.py
+++ b/d2warehouse/parser.py
@@ -2,7 +2,7 @@ import struct
 from bitarray import bitarray
 from bitarray.util import ba2int
 from d2warehouse.stash import Stash, StashTab
-from d2warehouse.item import Affix, Item, LowQualityType, Quality
+from d2warehouse.item import Affix, Item, LowQualityType, Quality, lookup_basetype
 import d2warehouse.huffman as huffman
 
 STASH_TAB_MAGIC = b"\x55\xAA\x55\xAA"
@@ -154,7 +154,9 @@ def parse_item(data: bytes) -> tuple[bytes, Item]:
     else:
         personalized_end = runeword_end
 
-    extended_byte_size = int((personalized_end + 7) / 8)
+    item, itemtype_end = parse_basetype_data(bits[personalized_end:], item)
+
+    extended_byte_size = int((itemtype_end + 7) / 8)
     print("extended size", extended_byte_size)
 
     item.raw_data = data[:]
@@ -265,6 +267,32 @@ def parse_personalization(bits: bitarray) -> tuple[str, int]:
     return output, ptr + 1
 
 
+def parse_basetype_data(bits: bitarray, item: Item) -> tuple[Item, int]:
+    cls = lookup_basetype(item.kind)["class"]
+    ptr = 0
+
+    if cls == "armor":
+        item.defense = ba2int(bits[0:10]) + 10
+        ptr += 10
+
+    if cls in ["armor", "weapon"]:
+        item.max_durability = ba2int(bits[ptr : ptr + 8])
+        item.durability = ba2int(bits[ptr : ptr + 8])
+        ptr += 16
+        if item.is_socketed:
+            item.sockets = ba2int(bits[ptr : ptr + 4])
+            ptr += 4
+
+    if cls == "tome":
+        ptr += 5  # unknown field
+
+    if cls in ["tome", "stackable"]:
+        item.quality = ba2int(bits[ptr : ptr + 9])
+        ptr += 9
+
+    return item, ptr
+
+
 def parse_items(data: bytes) -> list[Item]:
     data = parse_fixed(data, ITEM_DATA_MAGIC)
     data, num = parse_u16(data)