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

from collections import defaultdict

from sqlalchemy import ForeignKey, String, Text, UniqueConstraint
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.orm import Mapped, mapped_column, relationship

from ttfrog.db.base import BaseObject, SlugMixin
from ttfrog.db.schema.classes import CharacterClass, ClassFeature
from ttfrog.db.schema.constants import DamageType, Defenses
from ttfrog.db.schema.modifiers import Condition, Modifier, ModifierMixin, Stat
from ttfrog.db.schema.skill import Skill

__all__ = [
    "Ancestry",
    "AncestryTrait",
    "AncestryTraitMap",
    "CharacterClassMap",
    "CharacterClassFeatureMap",
    "Character",
    "Modifier",
]


def class_map_creator(fields):
    if isinstance(fields, CharacterClassMap):
        return fields
    return CharacterClassMap(**fields)


def skill_creator(fields):
    if isinstance(fields, CharacterSkillMap):
        return fields
    return CharacterSkillMap(**fields)

def condition_creator(fields):
    if isinstance(fields, CharacterConditionMap):
        return fields
    return CharacterConditionMap(**fields)


def attr_map_creator(fields):
    if isinstance(fields, CharacterClassFeatureMap):
        return fields
    return CharacterClassFeatureMap(**fields)


class HitDie(BaseObject):
    __tablename__ = "hit_die"
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    character_id: Mapped[int] = mapped_column(ForeignKey("character.id"))
    character_class_id: Mapped[int] = mapped_column(ForeignKey("character_class.id"))
    character_class = relationship("CharacterClass", lazy="immediate")

    spent: Mapped[bool] = mapped_column(nullable=False, default=False)

    @property
    def name(self):
        return self.character_class.hit_die_name

    @property
    def stat(self):
        return self.character_class.hit_die_stat_name


class AncestryTraitMap(BaseObject):
    __tablename__ = "trait_map"
    __table_args__ = (UniqueConstraint("ancestry_id", "ancestry_trait_id"),)
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    ancestry_id: Mapped[int] = mapped_column(ForeignKey("ancestry.id"))
    ancestry_trait_id: Mapped[int] = mapped_column(ForeignKey("ancestry_trait.id"), init=False)
    trait: Mapped["AncestryTrait"] = relationship(uselist=False, lazy="immediate")
    level: Mapped[int] = mapped_column(nullable=False, info={"min": 1, "max": 20})


class Ancestry(BaseObject, ModifierMixin):
    """
    A character ancestry ("race"), which has zero or more AncestryTraits and Modifiers.
    """

    __tablename__ = "ancestry"
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    name: Mapped[str] = mapped_column(String(collation="NOCASE"), nullable=False, unique=True)

    creature_type: Mapped[str] = mapped_column(nullable=False, default="humanoid")
    size: Mapped[str] = mapped_column(nullable=False, default="medium")
    speed: Mapped[int] = mapped_column(nullable=False, default=30, info={"min": 0, "max": 99})

    fly_speed: Mapped[int] = mapped_column(init=False, nullable=True, info={"min": 0, "max": 99})
    climb_speed: Mapped[int] = mapped_column(init=False, nullable=True, info={"min": 0, "max": 99})
    swim_speed: Mapped[int] = mapped_column(init=False, nullable=True, info={"min": 0, "max": 99})

    _traits = relationship(
        "AncestryTraitMap", init=False, uselist=True, cascade="all,delete,delete-orphan", lazy="immediate"
    )

    @property
    def traits(self):
        return [mapping.trait for mapping in self._traits]

    def add_trait(self, trait, level=1):
        if trait not in self.traits:
            mapping = AncestryTraitMap(ancestry_id=self.id, trait=trait, level=level)
            if not self._traits:
                self._traits = [mapping]
            else:
                self._traits.append(mapping)
            return True
        return False


class AncestryTrait(BaseObject, ModifierMixin):
    """
    A trait granted to a character via its Ancestry.
    """

    __tablename__ = "ancestry_trait"
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    name: Mapped[str] = mapped_column(String(collation="NOCASE"), nullable=False, unique=True)
    description: Mapped[Text] = mapped_column(Text, default="")


class CharacterSkillMap(BaseObject):
    __tablename__ = "character_skill_map"
    __table_args__ = (UniqueConstraint("skill_id", "character_id", "character_class_id"),)
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    skill_id: Mapped[int] = mapped_column(ForeignKey("skill.id"))
    character_id: Mapped[int] = mapped_column(ForeignKey("character.id"), nullable=True, default=None)
    character_class_id: Mapped[int] = mapped_column(ForeignKey("character_class.id"), nullable=True, default=None)
    proficient: Mapped[bool] = mapped_column(default=True)
    expert: Mapped[bool] = mapped_column(default=False)

    skill = relationship("Skill", lazy="immediate")


class CharacterClassMap(BaseObject):
    __tablename__ = "class_map"
    __table_args__ = (UniqueConstraint("character_id", "character_class_id"),)
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    character: Mapped["Character"] = relationship(uselist=False, viewonly=True)
    character_class: Mapped["CharacterClass"] = relationship(lazy="immediate")

    character_id: Mapped[int] = mapped_column(ForeignKey("character.id"), init=False, nullable=False)
    character_class_id: Mapped[int] = mapped_column(ForeignKey("character_class.id"), init=False, nullable=False)

    level: Mapped[int] = mapped_column(nullable=False, info={"min": 1, "max": 20}, default=1)


class CharacterClassFeatureMap(BaseObject):
    __tablename__ = "character_class_feature_map"
    __table_args__ = (UniqueConstraint("character_id", "class_feature_id"),)
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    character_id: Mapped[int] = mapped_column(ForeignKey("character.id"), nullable=False)
    class_feature_id: Mapped[int] = mapped_column(ForeignKey("class_feature.id"), nullable=False)
    option_id: Mapped[int] = mapped_column(ForeignKey("class_feature_option.id"), nullable=False)

    class_feature: Mapped["ClassFeature"] = relationship(lazy="immediate")
    option = relationship("ClassFeatureOption", lazy="immediate")

    character_class = relationship(
        "CharacterClass",
        secondary="class_map",
        primaryjoin="CharacterClassFeatureMap.character_id == CharacterClassMap.character_id",
        secondaryjoin="CharacterClass.id == CharacterClassMap.character_class_id",
        viewonly=True,
        uselist=False,
    )

class CharacterConditionMap(BaseObject):
    __tablename__ = "character_condition_map"
    __table_args__ = (UniqueConstraint("condition_id", "character_id"), )
    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
    condition_id: Mapped[int] = mapped_column(ForeignKey("condition.id"))
    character_id: Mapped[int] = mapped_column(ForeignKey("character.id"), nullable=True, default=None)
    condition = relationship("Condition", lazy="immediate")

class Character(BaseObject, SlugMixin, ModifierMixin):
    __tablename__ = "character"

    id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)

    name: Mapped[str] = mapped_column(String(collation="NOCASE"), nullable=False, default="New Character")
    hit_points: Mapped[int] = mapped_column(default=10, nullable=False, info={"min": 0, "max": 999})
    temp_hit_points: Mapped[int] = mapped_column(default=0, nullable=False, info={"min": 0, "max": 999})

    _max_hit_points: Mapped[int] = mapped_column(
        default=10, nullable=False, info={"min": 0, "max": 999, "modifiable": True}
    )
    _armor_class: Mapped[int] = mapped_column(
        default=10, nullable=False, info={"min": 1, "max": 99, "modifiable": True}
    )
    _strength: Mapped[int] = mapped_column(
        nullable=False, default=10, info={"min": 0, "max": 30, "modifiable_class": Stat}
    )
    _dexterity: Mapped[int] = mapped_column(
        nullable=False, default=10, info={"min": 0, "max": 30, "modifiable_class": Stat}
    )
    _constitution: Mapped[int] = mapped_column(
        nullable=False, default=10, info={"min": 0, "max": 30, "modifiable_class": Stat}
    )
    _intelligence: Mapped[int] = mapped_column(
        nullable=False, default=10, info={"min": 0, "max": 30, "modifiable_class": Stat}
    )
    _wisdom: Mapped[int] = mapped_column(
        nullable=False, default=10, info={"min": 0, "max": 30, "modifiable_class": Stat}
    )
    _charisma: Mapped[int] = mapped_column(
        nullable=False, default=10, info={"min": 0, "max": 30, "modifiable_class": Stat}
    )

    _actions_per_turn: Mapped[int] = mapped_column(nullable=False, default=1, info={"min": 0, "max": 99, "modifiable": True})
    _bonus_actions_per_turn: Mapped[int] = mapped_column(nullable=False, default=1, info={"min": 0, "max": 99, "modifiable": True})
    _reactions_per_turn: Mapped[int] = mapped_column(nullable=False, default=1, info={"min": 0, "max": 99, "modifiable": True})

    vision: Mapped[int] = mapped_column(default=None, nullable=True, info={"min": 0, "modifiable": True})
    exhaustion: Mapped[int] = mapped_column(nullable=False, default=0, info={"min": 0, "max": 5})

    class_map = relationship("CharacterClassMap", cascade="all,delete,delete-orphan")
    class_list = association_proxy("class_map", "id", creator=class_map_creator)

    _skills = relationship("CharacterSkillMap", uselist=True, cascade="all,delete,delete-orphan", lazy="immediate")
    skills = association_proxy("_skills", "skill", creator=skill_creator)

    _conditions = relationship("CharacterConditionMap", uselist=True, cascade="all,delete,delete-orphan", lazy="immediate")
    conditions = association_proxy("_conditions", "condition", creator=condition_creator)

    character_class_feature_map = relationship("CharacterClassFeatureMap", cascade="all,delete,delete-orphan")
    feature_list = association_proxy("character_class_feature_map", "id", creator=attr_map_creator)

    ancestry_id: Mapped[int] = mapped_column(ForeignKey("ancestry.id"), nullable=False, default="1")
    ancestry: Mapped["Ancestry"] = relationship(uselist=False, default=None)

    _hit_dice = relationship("HitDie", uselist=True, cascade="all,delete,delete-orphan", lazy="immediate")

    @property
    def hit_dice(self):
        pool = defaultdict(list)
        for die in self._hit_dice:
            pool[die.character_class.name].append(die)
        return pool

    @property
    def hit_dice_available(self):
        return [die for die in self._hit_dice if die.spent is False]

    @property
    def proficiency_bonus(self):
        return 1 + int(0.5 + self.level / 4)

    @property
    def expertise_bonus(self):
        return 2 * self.proficiency_bonus

    @property
    def modifiers(self):
        unified = {}
        unified.update(**self.ancestry.modifiers)
        for trait in self.traits:
            unified.update(**trait.modifiers)
        for condition in self.conditions:
            unified.update(**condition.modifiers)
        unified.update(**super().modifiers)
        return unified

    @property
    def classes(self):
        return dict([(mapping.character_class.name, mapping.character_class) for mapping in self.class_map])

    @property
    def traits(self):
        return self.ancestry.traits

    @property
    def initiative(self):
        return self._apply_modifiers("initiative", self.dexterity.bonus)

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

    @property
    def climb_speed(self):
        return self._apply_modifiers("climb_speed", self.ancestry.climb_speed or int(self.speed / 2))

    @property
    def swim_speed(self):
        return self._apply_modifiers("swim_speed", self.ancestry.swim_speed or int(self.speed / 2))

    @property
    def fly_speed(self):
        modified = self._apply_modifiers("fly_speed", self.ancestry.fly_speed or 0)
        if self.ancestry.fly_speed is None and not modified:
            return None
        return self._apply_modifiers("speed", modified)

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

    @property
    def vision_in_darkness(self):
        return self._apply_modifiers("vision_in_darkness", self.vision if self.vision is not None else 0)

    @property
    def level(self):
        return sum(mapping.level for mapping in self.class_map)

    @property
    def levels(self):
        return dict([(mapping.character_class.name, mapping.level) for mapping in self.class_map])

    @property
    def class_features(self):
        return dict([(mapping.class_feature.name, mapping.option) for mapping in self.character_class_feature_map])

    def level_in_class(self, charclass):
        mapping = [mapping for mapping in self.class_map if mapping.character_class_id == charclass.id]
        if not mapping:
            return None
        return mapping[0]

    def immune(self, damage_type: DamageType):
        return self.defense(damage_type) == Defenses.immune

    def resistant(self, damage_type: DamageType):
        return self.defense(damage_type) == Defenses.resistant.value

    def vulnerable(self, damage_type: DamageType):
        return self.defense(damage_type) == Defenses.vulnerable

    def absorbs(self, damage_type: DamageType):
        return self.defense(damage_type) == Defenses.absorbs

    def defense(self, damage_type: DamageType):
        return self._apply_modifiers(damage_type, None)

    def check_modifier(self, skill: Skill, save: bool = False):
        # if the skill is not assigned, but we have modifiers, apply them to zero.
        if skill not in self.skills:
            target = f"{skill.name.lower()}_{'save' if save else 'check'}"
            if self.has_modifier(target):
                modified = self._apply_modifiers(target, 0)
                return modified

        # if the skill is a stat, start with the bonus value
        attr = skill.name.lower()
        stat = getattr(self, attr, None)
        initial = getattr(stat, "bonus", None)

        # if the skill isn't a stat, try the parent.
        if initial is None and skill.parent:
            stat = getattr(self, skill.parent.name.lower(), None)
            initial = getattr(stat, "bonus", initial)

        # if the skill is a proficiency, apply the bonus to the initial value
        if skill in self.skills:
            mapping = [mapping for mapping in self._skills if mapping.skill_id == skill.id][-1]
            if mapping.expert and not save:
                initial += 2 * self.proficiency_bonus
            elif mapping.proficient:
                initial += self.proficiency_bonus

        # return the initial value plus any modifiers.
        return self._apply_modifiers(f"{attr}_{'save' if save else 'check'}", initial)

    def add_class(self, newclass, level=1):
        if level == 0:
            return self.remove_class(newclass)

        # add the class mapping and/or set the character's level in the class
        mapping = self.level_in_class(newclass)
        if not mapping:
            self.class_list.append(CharacterClassMap(character=self, character_class=newclass, level=level))
        else:
            mapping.level = level

        # add class features with default values
        for lvl in range(1, level + 1):
            for attr in newclass.features_at_level(lvl):
                self.add_class_feature(newclass, attr, attr.options[0])

        # add default class skills
        for skill in newclass.skills[: newclass.starting_skills]:
            self.add_skill(skill, proficient=True, character_class=newclass)

        # add hit dice
        existing = len([die for die in self._hit_dice if die.character_class_id == newclass.id])
        for lvl in range(level - existing):
            self._hit_dice.append(HitDie(character_id=self.id, character_class_id=newclass.id))

    def remove_class(self, target):
        self.class_map = [m for m in self.class_map if m.character_class != target]
        for mapping in self.character_class_feature_map:
            if mapping.character_class == target:
                self.remove_class_feature(mapping.class_feature)
            for skill in target.skills:
                self.remove_skill(skill, proficient=True, expert=False, character_class=target)
        self._hit_dice = [die for die in self._hit_dice if die.character_class != target]

    def remove_class_feature(self, feature):
        self.character_class_feature_map = [
            m for m in self.character_class_feature_map if m.class_feature.id != feature.id
        ]

    def has_class_feature(self, feature):
        return feature in [m.class_feature for m in self.character_class_feature_map]

    def add_class_feature(self, character_class, feature, option):
        if self.has_class_feature(feature):
            return False
        mapping = self.level_in_class(character_class)
        if not mapping:
            return False
        if feature not in mapping.character_class.features_at_level(mapping.level):
            return False
        self.feature_list.append(
            CharacterClassFeatureMap(
                character_id=self.id,
                class_feature_id=feature.id,
                option_id=option.id,
                class_feature=feature,
            )
        )
        return True

    def add_modifier(self, modifier):
        if not super().add_modifier(modifier):
            return False

        if modifier.new_value != Defenses.immune:
            return True

        modified_condition = None
        for cond in self.conditions:
            if modifier.target == cond.name:
                modified_condition = cond
                break
        if not modified_condition:
            return True

        return self.remove_condition(modified_condition)

    def has_condition(self, condition):
        return condition in self.conditions

    def add_condition(self, condition):
        if self.immune(condition.name):
            return False
        if self.has_condition(condition):
            return False
        self._conditions.append(CharacterConditionMap(condition_id=condition.id, character_id=self.id))
        return True

    def remove_condition(self, condition):
        if not self.has_condition(condition):
            return False
        mappings = [mapping for mapping in self._conditions if mapping.condition_id != condition.id]
        self._conditions = mappings
        return True

    def add_skill(self, skill, proficient=False, expert=False, character_class=None):
        skillmap = None
        exists = False
        if skill in self.skills:
            for mapping in self._skills:
                if mapping.skill_id != skill.id:
                    continue
                if character_class is None and mapping.character_class_id:
                    continue
                if (character_class is None and mapping.character_class_id is None) or (
                    mapping.character_class_id == character_class.id
                ):
                    skillmap = mapping
                    exists = True
                    break

        if not skillmap:
            skillmap = CharacterSkillMap(skill_id=skill.id, character_id=self.id)

        skillmap.proficient = proficient
        skillmap.expert = expert
        if character_class:
            skillmap.character_class_id = character_class.id

        if not exists:
            self._skills.append(skillmap)
            return True
        return False

    def remove_skill(self, skill, proficient, expert, character_class):
        to_delete = [
            mapping
            for mapping in self._skills
            if (
                mapping.skill_id == skill.id
                and mapping.proficient == proficient
                and mapping.expert == expert
                and (
                    (mapping.character_class_id is None and character_class is None)
                    or (character_class and mapping.character_class_id == character_class.id)
                )
            )
        ]
        if not to_delete:
            return False
        self._skills = [m for m in self._skills if m not in to_delete]
        return True

    def apply_healing(self, value: int):
        self.hit_points = min(self.hit_points + value, self._max_hit_points)

    def apply_damage(self, value: int, damage_type: DamageType):
        total = value
        if self.absorbs(damage_type):
            return self.apply_healing(total)
        if self.immune(damage_type):
            return
        if self.resistant(damage_type):
            total = int(value / 2)
        elif self.vulnerable(damage_type):
            total = value * 2

        if total <= self.temp_hit_points:
            self.temp_hit_points -= total
            return

        self.hit_points = max(0, self.hit_points - (total - self.temp_hit_points))
        self.temp_hit_points = 0

    def spend_hit_die(self, die):
        die.spent = True

    def __after_insert__(self, session):
        """
        Called by the session after_flush event listener to add default joins in other tables.
        """
        for skill in session.query(Skill).filter(
            Skill.name.in_(("strength", "dexterity", "constitution", "intelligence", "wisdom", "charisma"))
        ):
            self.add_skill(skill, proficient=False, expert=False)