d2warehouse/d2warehouse/parser.py

# Copyright 2023 <omicron.me@protonmail.com>
# Copyright 2023 <andreasruden91@gmail.com>
#
# This file is part of d2warehouse.
#
# d2warehouse is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# d2warehouse is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# Mercator. If not, see <https://www.gnu.org/licenses/>.
import struct
from bitarray import bitarray
from bitarray.util import ba2int
from d2warehouse.stash import Stash, StashTab
from d2warehouse.item import (
    Item,
    LowQualityType,
    Quality,
    Stat,
    lookup_basetype,
    lookup_stat,
)
import d2warehouse.huffman as huffman
from d2warehouse.fileformat import STASH_TAB_MAGIC, STASH_TAB_VERSION, ITEM_DATA_MAGIC


class ParseError(RuntimeError):
    pass


class UnsupportedItemError(ParseError):
    pass


def ensure_length(data, n):
    if len(data) < n:
        nearby = data[:10].hex()
        raise ParseError(
            f"Expected {n} bytes but only {len(data)} bytes are available near {nearby}"
        )


def parse_fixed(data: bytes, prefix: bytes) -> bytes:
    if data.startswith(prefix):
        return data[len(prefix) :]
    raise ParseError(f"Expected {prefix} near {data[:10].hex()}")


def parse_bytes(data: bytes, n: int) -> tuple[bytes, bytes]:
    ensure_length(data, n)
    return data[n:], data[:n]


def parse_u32(data: bytes) -> tuple[bytes, int]:
    ensure_length(data, 4)
    return data[4:], struct.unpack("<I", data[:4])[0]


def parse_u16(data: bytes) -> tuple[bytes, int]:
    ensure_length(data, 2)
    return data[2:], struct.unpack("<H", data[:2])[0]


def parse_stash(data) -> Stash:
    stash = Stash()
    while len(data) > 0:
        data, tab = parse_stash_tab(data)
        stash.tabs.append(tab)
    return stash


def parse_stash_tab(data: bytes) -> tuple[bytes, StashTab]:
    data = parse_fixed(data, STASH_TAB_MAGIC)
    data, unknown = parse_u32(data)
    data, version = parse_u32(data)

    if unknown != 1:
        ParseError("Unknown stash tab field is not 1")
    if version != STASH_TAB_VERSION:
        ParseError(f"Unsupported stash tab version ({version} instead of 99)")

    tab = StashTab()
    data, gold = parse_u32(data)
    tab.gold = gold

    # Length is the total length of the tab data in bytes. This includes all
    # fields that are already parsed above.
    data, length = parse_u32(data)

    # Skip what is probably zero padding (?) but check just in case it contains data
    data, pad = parse_bytes(data, 44)
    if pad != b"\x00" * 44:
        raise ParseError(f"Unexpecteed data in zero padding: {pad}")

    # Separate out the item data so we can easilly verify items parsed the correct
    # number of bytes
    data, item_data = parse_bytes(data, length - 64)

    tab.item_data = item_data
    tab.items = parse_items(item_data)

    return data, tab


def parse_item(data: bytes) -> tuple[bytes, Item]:
    bits = bitarray(endian="little")
    bits.frombytes(data)

    is_identified = bool(bits[4])
    is_socketed = bool(bits[11])
    is_ear = bool(bits[16])
    is_beginner = bool(bits[17])
    is_simple = bool(bits[21])
    is_ethereal = bool(bits[22])
    is_personalized = bool(bits[24])
    is_runeword = bool(bits[26])
    pos_x = ba2int(bits[42:46])
    pos_y = ba2int(bits[46:50])
    code, code_end = huffman.decode(bits[53:], 3)
    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
    filled_sockets = ba2int(bits[code_end:sockets_end])
    if is_ear:
        raise UnsupportedItemError("Ear items are not supported")

    simple_byte_sz = int((sockets_end + 7) / 8)
    item = Item(
        data[:simple_byte_sz],
        STASH_TAB_VERSION,
        is_identified,
        is_socketed,
        is_beginner,
        is_simple,
        is_ethereal,
        is_personalized,
        is_runeword,
        pos_x,
        pos_y,
        code,
    )
    if is_simple:
        return data[simple_byte_sz:], item

    item.uid, item.lvl, item.quality = parse_extended_item(bits[sockets_end:])

    extended_end = sockets_end + 32 + 7 + 4

    item.graphic, graphic_end = parse_item_graphic(bits[extended_end:])
    graphic_end += extended_end

    item.implicit, inherent_end = parse_implicit_mod(bits[graphic_end:])
    inherent_end += graphic_end

    item, quality_end = parse_quality_data(bits[inherent_end:], item)
    quality_end += inherent_end

    if item.is_runeword:
        item.runeword_id, runeword_end = parse_runeword(bits[quality_end:])
        runeword_end += quality_end
    else:
        runeword_end = quality_end

    if item.is_personalized:
        item.personal_name, personalized_end = parse_personalization(
            bits[runeword_end:], item
        )
        personalized_end += runeword_end
    else:
        personalized_end = runeword_end

    item, itembase_end = parse_basetype_data(bits[personalized_end:], item)
    itembase_end += personalized_end

    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 += sockets_end

    extended_byte_size = int((enchantments_end + 7) / 8)

    item.raw_data = data[:extended_byte_size]
    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]:
    if not bits[0]:
        return None, 1
    else:
        return ba2int(bits[1:4]), 4


def parse_implicit_mod(bits: bitarray) -> tuple[int | None, int]:
    # References a row in automagic.txt which adds an implicit mod to the item.
    # However, that mod seems to also be present in the stats at the end.
    if not bits[0]:
        return None, 1
    else:
        return ba2int(bits[1:11]), 12


def parse_extended_item(bits: bitarray) -> tuple[int, int, Quality]:
    uid = ba2int(bits[:32])
    lvl = ba2int(bits[32:39])
    quality = ba2int(bits[39:43])
    return uid, lvl, Quality(quality)


def parse_quality_data(bits: bitarray, item: Item) -> tuple[Item, int]:
    if item.quality == Quality.LOW:
        return parse_low_quality_data(bits, item)
    elif item.quality == item.quality.NORMAL:
        return item, 0
    elif item.quality == item.quality.HIGH:
        return parse_high_quality_data(bits, item)
    elif item.quality == item.quality.MAGIC:
        return parse_magic_data(bits, item)
    elif item.quality == item.quality.SET:
        return parse_set_data(bits, item)
    elif item.quality == item.quality.RARE:
        return parse_rare_data(bits, item)
    elif item.quality == item.quality.UNIQUE:
        return parse_unique_data(bits, item)
    elif item.quality == item.quality.CRAFTED:
        return parse_rare_data(bits, item)  # crafted & rare are the same


def parse_low_quality_data(bits: bitarray, item: Item) -> tuple[Item, int]:
    item.low_quality = LowQualityType(ba2int(bits[0:3]))
    return item, 3


def parse_high_quality_data(bits: bitarray, item: Item) -> tuple[Item, int]:
    # The data for superior item is unknown
    return item, 3


def parse_magic_data(bits: bitarray, item: Item) -> tuple[Item, int]:
    item.prefixes = [ba2int(bits[0:11])]
    item.suffixes = [ba2int(bits[11:22])]
    return item, 22


def parse_set_data(bits: bitarray, item: Item) -> tuple[Item, int]:
    item.set_id = ba2int(bits[0:12])
    return item, 12


def parse_rare_data(bits: bitarray, item: Item) -> tuple[Item, int]:
    item.nameword1 = ba2int(bits[0:8])
    item.nameword2 = ba2int(bits[8:16])
    affixes = []
    ptr = 16
    for _ in range(0, 6):
        (affix, sz) = parse_affix(bits[ptr:])
        ptr += sz
        affixes.append(affix)
    item.prefixes = [affix for affix in affixes[0:3] if affix is not None]
    item.suffixes = [affix for affix in affixes[3:6] if affix is not None]

    return item, ptr


def parse_unique_data(bits: bitarray, item: Item) -> tuple[Item, int]:
    item.unique_id = ba2int(bits[0:12])
    return item, 12


def parse_affix(bits: bitarray) -> tuple[int | None, int]:
    if not bits[0]:
        return None, 1
    else:
        return ba2int(bits[1:12]), 12


def parse_runeword(bits: bitarray) -> tuple[int, int]:
    id = ba2int(bits[0:12])
    # + 4 unknown bits
    # FIXME: Could this be the same field that set has? I.e. extra enchantment lists
    return id, 16


def parse_personalization(bits: bitarray) -> tuple[str, int]:
    output = ""
    ptr = 0
    ascii = ba2int(bits[0:7])
    while ascii:
        output += chr(ascii)
        ascii = ba2int(bits[ptr : ptr + 7])
        ptr += 7
    return output, ptr + 1


def parse_basetype_data(bits: bitarray, item: Item) -> tuple[Item, int]:
    basetype = lookup_basetype(item.code)
    cls = basetype["class"]
    stackable = basetype["stackable"]
    ptr = 0

    if cls == "tome":
        ptr += 5  # unknown field

    # Unknown bit here
    ptr += 1

    if cls == "armor":
        item.defense = ba2int(bits[ptr : ptr + 11]) - 10
        ptr += 11

    if cls in ["armor", "weapon"]:
        item.max_durability = ba2int(bits[ptr : ptr + 8])
        ptr += 8
        if item.max_durability > 0:
            item.durability = ba2int(bits[ptr : ptr + 8])
            ptr += 8 + 1  # uknown bit after durability

    if stackable:
        item.quantity = ba2int(bits[ptr : ptr + 9])
        ptr += 9

    return item, ptr


def parse_enchantments(bits: bitarray, item: Item) -> tuple[Item, int]:
    ptr = 0
    if item.quality == Quality.SET:
        val = ba2int(bits[0:5])
        # FIXME: how exactly this value works is unconfirmed
        num_lists = 1 + val.bit_count()
        ptr += 5
    elif item.is_runeword:
        # Runewords have one empty dummy list of only 0x1ff followed by the runeword stats
        num_lists = 2
    else:
        num_lists = 1

    # TODO: One extra list for runewords?

    if num_lists > 0:
        item.stats = []

    while num_lists > 0:
        stats, sz = parse_enchantment_list(bits[ptr:])
        ptr += sz
        num_lists -= 1
        item.stats.extend(stats)

    return item, ptr


def parse_enchantment_list(bits: bitarray) -> tuple[list[Stat], int]:
    ptr = 0
    stats = []

    while True:
        id = ba2int(bits[ptr : ptr + 9])
        ptr += 9
        if id == 0x1FF:
            break
        print("stat id", id)

        stat = lookup_stat(id)

        values = []
        param = None
        if stat["save_param_bits"]:
            param = ba2int(bits[ptr : ptr + stat["save_param_bits"]])
            ptr += stat["save_param_bits"]
        for b in stat["save_bits"]:
            values.append(ba2int(bits[ptr : ptr + b]))
            ptr += b
        values = [v - stat["save_add"] for v in values]
        text = stat["text"]

        stat = Stat(
            id=id,
            values=values,
            parameter=param,
            text=text,
        )

        print("Stat:")
        stat.print()

        stats.append(stat)

    return stats, ptr


def parse_items(data: bytes) -> list[Item]:
    data = parse_fixed(data, ITEM_DATA_MAGIC)
    data, num = parse_u16(data)

    items: list[Item] = []
    while data:
        data, item = parse_item(data)
        items.append(item)

    # TODO: check if num == len(items)

    return items