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

from collections import defaultdict
from typing import Any, Union

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 Modifiable:
    def __new__(cls, base, modified=None):
        cls.base = base
        return super().__new__(cls, modified)


class ModifiableStr(Modifiable, str):
    """
    A string that also has a '.base' property.
    """


class ModifiableInt(Modifiable, int):
    """
    An integer that also has a '.base' property
    """


class Stat(ModifiableInt):
    """
    Same as a Score except it also has a bonus for STR, DEX, CON, etc.
    """

    @property
    def bonus(self):
        return int((self - 10) / 2)


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)
    stacks: Mapped[bool] = mapped_column(nullable=False, default=False)
    absolute_value: Mapped[int] = mapped_column(nullable=True, default=None)
    multiply_value: Mapped[float] = mapped_column(nullable=True, default=None)
    multiply_attribute: Mapped[str] = mapped_column(nullable=True, default=None)
    relative_value: Mapped[int] = mapped_column(nullable=True, default=None)
    relative_attribute: Mapped[str] = 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 ... ]}
    """

    @declared_attr
    def modifier_map(cls):
        """
        Create the join between the current model and the ModifierMap table.
        """
        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):
        """
        Return all modifiers for the current instance as a dict keyed on target attribute name.
        """
        all_modifiers = defaultdict(list)
        for mapping in self.modifier_map:
            all_modifiers[mapping.modifier.target].append(mapping.modifier)
        return all_modifiers

    def has_modifier(self, name: str):
        return True if self.modifiers.get(name, None) else False

    def add_modifier(self, modifier: Modifier) -> bool:
        """
        Associate a modifier to the current instance if it isn't already.

        Returns True if the modifier was added; False if was already present.
        """
        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: Modifier) -> bool:
        """
        Remove a modifier from the map.

        Returns True if it was removed and False if it wasn't present.
        """
        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 _modifiable_column(self, attr_name: str) -> Union[mapped_column, None]:
        """
        Given an atttribute name, look for a column attribute with the same
        name but with an underscore prefix. If that column exists, and it
        has one or more of the expected "modifiable" keys in its info, the
        column is modifiable.

        Returns the matching column if it was found, or None.
        """
        col = getattr(self.__table__.columns, f"_{attr_name}", None)
        if col is None:
            return None
        for key in col.info.keys():
            if key.startswith("modifiable"):
                return col
        return None

    def _get_modifiable_base(self, attr_name: str) -> object:
        """
        Resolve a dottted string "foo.bar.baz" as its corresponding nested attribute.

        This is useful for cases where a column definition includes a modifiable_base
        that is some other attribute. For example:

            foo[int] = mapped_column(default=0, info={"modifiable_base": "ancestry.bar")

        This will create an initial value for self.foo equal to self.ancesetry.bar.
        """

        def get_attr(obj, parts):
            if parts:
                name, *parts = parts
                return get_attr(getattr(obj, name), parts)
            return obj

        return get_attr(self, attr_name.split("."))

    def _apply_one_modifier(self, modifier, initial, modified):
        print(f"Trying to apply {modifier}")
        if modifier.new_value is not None:
            return modifier.new_value
        elif modifier.absolute_value is not None:
            return modifier.absolute_value

        base_value = modified if modifier.stacks else initial

        if modifier.multiply_attribute is not None:
            return int(base_value * getattr(self, modifier.multiply_attribute) + 0.5)

        if modifier.multiply_value is not None:
            return int(base_value * modifier.multiply_value + 0.5)

        if modifier.relative_attribute is not None:
            return base_value + getattr(self, modifier.relative_attribute)

        if modifier.relative_value is not None:
            return base_value + modifier.relative_value

        raise Exception(f"Cannot apply modifier: {modifier = }")

    def _apply_modifiers(self, target: str, initial: Any, modifiable_class: type = None) -> Modifiable:
        """
        Apply all the modifiers for a given target and return the modified value.

        This is mostly called from __getattr__() below to handle cases where a
        column is named self._foo but the modified value is accessible as
        self.foo. It can also be invoked directly, as, say from a property:

            @property
            def speed(self):
                return self._apply_modifiers("speed", self.ancestry.speed)

        Args:
            target            - The name of the attribute to modify
            initial           - The initial value for the target
            modifiable_class  - The object type to return; inferred from the
                                target attribute's type if not specified.
        """
        if not modifiable_class:
            modifiable_class = globals()["ModifiableInt"] if isinstance(initial, int) else globals()["ModifiableStr"]

        modifiers = self.modifiers.get(target, [])

        nonstacking = [m for m in modifiers if not m.stacks]
        if nonstacking:
            return modifiable_class(base=initial, modified=self._apply_one_modifier(nonstacking[-1], initial, initial))

        modified = initial
        for modifier in modifiers:
            if modifier.stacks:
                modified = self._apply_one_modifier(modifier, initial, modified)
        return modifiable_class(base=initial, modified=modified)

    def __setattr__(self, attr_name, value):
        """
        Prevent callers from setting the value of a Modifiable directly.
        """
        col = self._modifiable_column(attr_name)
        if col is not None:
            raise AttributeError(f"You cannot modify .{attr_name}. Did you mean ._{attr_name}?")
        return super().__setattr__(attr_name, value)

    def __getattr__(self, attr_name):
        """
        If the instance has an attribute equal to attr_name but prefixed with an
        underscore, check to see if that attribute is a column, and modifiable.
        If it is, return a Modifiable instance corresponding to that column's value.
        """
        col = self._modifiable_column(attr_name)
        if col is not None:
            return self._apply_modifiers(
                attr_name,
                self._get_modifiable_base(col.info.get("modifiable_base", col.name)),
                modifiable_class=col.info.get("modifiable_class", None),
            )
        raise AttributeError(f"No such attribute: {attr_name}.")