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

from typing import List

from sqlalchemy import ForeignKey, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship

from ttfrog.db.base import BaseObject, EnumField
from ttfrog.db.schema.item import Item, ItemProperty, ItemType


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


inventory_type_map = {
    InventoryType.EQUIPMENT: [
        ItemType.WEAPON,
        ItemType.ARMOR,
        ItemType.SHIELD,
        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)

    items: Mapped[List["InventoryMap"]] = relationship(
        uselist=True, cascade="all,delete,delete-orphan", lazy="immediate", default_factory=lambda: []
    )

    character = relationship("Character", init=False, viewonly=True, lazy="immediate")

    def get(self, item):
        return self.get_all(item)[0]

    def get_all(self, item):
        return [mapping for mapping in self.items 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
        if item.charges:
            mapping.charges = [Charge(inventory_map_id=mapping.id) for i in range(item.charges)]
        self.items.append(mapping)
        return mapping

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

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

    def __iter__(self):
        yield from self.items


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)
    attuned: Mapped[bool] = mapped_column(default=False)
    count: Mapped[int] = mapped_column(nullable=False, default=1)

    always_prepared: Mapped[bool] = mapped_column(default=False)

    charges: Mapped[List["Charge"]] = relationship(
        uselist=True, cascade="all,delete,delete-orphan", lazy="immediate", default_factory=lambda: []
    )

    @property
    def charges_available(self):
        return [charge for charge in self.charges if not charge.expended]

    @property
    def prepared(self):
        if self.item.item_type == ItemType.SPELL:
            return self.equipped or self.always_prepared

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

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

    def prepare(self):
        if self.item.item_type != ItemType.SPELL:
            return False
        if self.item.level > 0 and not self.inventory.character.spell_slots_by_level[self.item.level]:
            return False
        return self.equip()

    def unprepare(self):
        if self.item.item_type != ItemType.SPELL:
            return False
        return self.unequip()

    def use(self, item_property: ItemProperty, charges=None):
        if item_property.charge_cost is None:
            return True
        avail = self.charges_available
        if charges is None:
            charges = item_property.charge_cost
        if len(avail) < charges:
            return False
        for charge in avail:
            charge.expended = True
        return True

    def consume(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

    def cast(self, level=0):
        if self.item.item_type != ItemType.SPELL:
            return False

        if not self.prepared:
            return False
        if not level:
            level = self.item.level

        # cantrips
        if level == 0:
            return True

        # expend the spell slot
        avail = self.inventory.character.spell_slots_available[level]
        if not avail:
            return False
        avail[0].expended = True
        return True

    def attune(self):
        if self.attuned:
            return False
        if not self.item.requires_attunement:
            return False
        if len(self.inventory.character.attuned_items) >= 3:
            return False
        self.attuned = True
        return True

    def unattune(self):
        if not self.attuned:
            return False
        self.attuned = False
        return True


class Charge(BaseObject):
    __tablename__ = "charge"
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    inventory_map_id: Mapped[int] = mapped_column(ForeignKey("inventory_map.id"))
    expended: Mapped[bool] = mapped_column(nullable=False, default=False)