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

from collections import defaultdict

from sqlalchemy import ForeignKey, UniqueConstraint
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy.orm import Mapped, mapped_column, relationship

from ttfrog.db.base import BaseObject


class ModifierMap(BaseObject):
    """
    Creates a many-to-many between Modifier and any model inheriting from the ModifierMixin.
    """

    __tablename__ = "modifier_map"
    __table_args__ = (UniqueConstraint("primary_table_name", "primary_table_id", "modifier_id"),)
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    modifier_id: Mapped[int] = mapped_column(ForeignKey("modifier.id"), init=False)
    modifier: Mapped["Modifier"] = relationship(uselist=False, lazy="immediate")

    primary_table_name: Mapped[str] = mapped_column(nullable=False)
    primary_table_id: Mapped[int] = mapped_column(nullable=False)


class Modifier(BaseObject):
    """
    Modifiers modify the base value of an existing attribute on another table.

    Modifiers are applied by the Character class, but may be associated with any model via the
    ModifierMixIn model; refer to the Ancestry class for an example.
    """

    __tablename__ = "modifier"
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    name: Mapped[str] = mapped_column(nullable=False)
    target: Mapped[str] = mapped_column(nullable=False)
    absolute_value: Mapped[int] = mapped_column(nullable=True, default=None)
    relative_value: Mapped[int] = mapped_column(nullable=True, default=None)
    multiply_value: Mapped[float] = mapped_column(nullable=True, default=None)
    new_value: Mapped[str] = mapped_column(nullable=True, default=None)
    description: Mapped[str] = mapped_column(default="")


class ModifierMixin:
    """
    Add modifiers to an existing class.

    Attributes:
         modifier_map   - get/set a list of Modifier records associated with the parent
         modifiers      - read-only dict of lists of modifiers keyed on Modifier.target

    Methods:
        add_modifier     - Add a Modifier association to the modifier_map
        remove_modifier  - Remove a modifier association from the modifier_map

    Example:

        >>> class Item(BaseObject, ModifierMixin):
                id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
                name: Mapped[str] = mapped_column(nullable=False)
        >>> dwarven_belt = Item(name="Dwarven Belt")
        >>> dwarven_belt.add_modifier(Modifier(name="STR+1", target="strength", relative_value=1))
        >>> dwarven_belt.modifiers
        {'strength': [Modifier(id=1, target='strength', name='STR+1', relative_value=1 ... ]}
    """

    _modifiable_attributes = dict()

    @declared_attr
    def modifier_map(cls):
        return relationship(
            "ModifierMap",
            primaryjoin=(
                "and_("
                f"foreign(ModifierMap.primary_table_name)=='{cls.__tablename__}', "
                f"foreign(ModifierMap.primary_table_id)=={cls.__name__}.id"
                ")"
            ),
            cascade="all,delete,delete-orphan",
            overlaps="modifier_map,modifier_map",
            single_parent=True,
            uselist=True,
            lazy="immediate",
        )

    @property
    def modifiers(self):
        all_modifiers = defaultdict(list)
        for mapping in self.modifier_map:
            all_modifiers[mapping.modifier.target].append(mapping.modifier)
        return all_modifiers

    def add_modifier(self, modifier):
        if modifier.absolute_value is not None and modifier.relative_value is not None and modifier.multiple_value:
            raise AttributeError(f"You must provide only one of absolute, relative, and multiple values {modifier}.")

        if [mod for mod in self.modifier_map if mod.modifier == modifier]:
            return False
        self.modifier_map.append(
            ModifierMap(
                primary_table_name=self.__tablename__,
                primary_table_id=self.id,
                modifier=modifier,
            )
        )
        return True

    def remove_modifier(self, modifier):
        if modifier not in self.modifiers[modifier.target]:
            return False
        self.modifier_map = [mapping for mapping in self.modifier_map if mapping.modifier != modifier]
        return True

    def apply_modifiers(self, target, initial):
        if not self._modifiable_attributes:
            raise NotImplementedError(
                f"You must define the '_modifiable_attributes' property on {self.__class__.__name__}."
            )

        modifiers = list(reversed(self.modifiers.get(target, [])))
        if initial is None:
            return initial
        if isinstance(initial, int):
            absolute = [mod for mod in modifiers if mod.absolute_value is not None]
            if absolute:
                return absolute[0].absolute_value
            multiple = [mod for mod in modifiers if mod.multiply_value is not None]
            if multiple:
                return int(initial * multiple[0].multiply_value + 0.5)
            return initial + sum(mod.relative_value for mod in modifiers if mod.relative_value is not None)

        new = [mod for mod in modifiers if mod.new_value is not None]
        if new:
            return new[0].new_value
        return initial

    def __getattr__(self, attr_name):
        prop = self._modifiable_attributes.get(attr_name, None)
        if not prop:
            raise AttributeError(f"Attribute not found: {attr_name}")
        return self.apply_modifiers(attr_name, getattr(self, prop))