tabletop-frog/src/ttfrog/db/schema/inventory.py

from sqlalchemy import ForeignKey, UniqueConstraint
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.orm import Mapped, mapped_column, relationship
from ttfrog.db.base import BaseObject, EnumField
from ttfrog.db.schema.item import Item, ItemType


class InventoryType(EnumField):
    EQUIPMENT = "EQUIPMENT"
    SPELL = "SPELL"


inventory_type_map = {
    InventoryType.EQUIPMENT: [
        ItemType.ITEM,
        ItemType.SCROLL,
    ],
    InventoryType.SPELL: [
        ItemType.SPELL
    ]
}


def inventory_map_creator(fields):
    if isinstance(fields, InventoryMap):
        return fields
    return InventoryMap(**fields)


class Inventory(BaseObject):
    __tablename__ = "inventory"
    __table_args__ = (UniqueConstraint("character_id", "inventory_type"), )
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    character_id: Mapped[int] = mapped_column(ForeignKey("character.id"))
    inventory_type: Mapped[InventoryType] = mapped_column(nullable=False)

    _inventory_map = relationship("InventoryMap", lazy="immediate", uselist=True, cascade="all,delete,delete-orphan")
    inventory_map = association_proxy("_inventory_map", "id", creator=inventory_map_creator)

    def get(self, item):
        return [mapping for mapping in self._inventory_map if mapping.item == item]

    def add(self, item):
        if item.item_type not in inventory_type_map[self.inventory_type]:
            return False
        mapping = InventoryMap(inventory_id=self.id, item_id=item.id)
        if item.consumable:
            mapping.count = item.count
        self.inventory_map.append(mapping)
        return mapping

    def remove(self, mapping):
        if mapping.id not in self.inventory_map:
            return False
        self.inventory_map.remove(mapping.id)
        return True

    def equip(self, mapping):
        if mapping.equipped:
            return False
        mapping.equipped = True
        return True

    def unequip(self, mapping):
        if not mapping.equipped:
            return False
        mapping.equipped = False
        return True

    def __contains__(self, obj):
        for mapping in self._inventory_map:
            if mapping.item == obj:
                return True
        return False

    def __iter__(self):
        yield from self._inventory_map


class InventoryMap(BaseObject):
    __tablename__ = "inventory_map"
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    inventory_id: Mapped[int] = mapped_column(ForeignKey("inventory.id"))
    item_id: Mapped[int] = mapped_column(ForeignKey("item.id"))
    item: Mapped["Item"] = relationship(uselist=False, lazy="immediate", viewonly=True, init=False)
    inventory: Mapped["Inventory"] = relationship(uselist=False, viewonly=True, init=False)

    equipped: Mapped[bool] = mapped_column(default=False)
    count: Mapped[int] = mapped_column(nullable=False, default=1)

    def use(self, count=1):
        if count < 0:
            return False
        if not self.item.consumable:
            return False
        if self.count < count:
            return False
        self.count -= count
        if self.count == 0:
            self.inventory.remove(self)
            return 0
        return self.count