import itertools from collections import defaultdict 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 from ttfrog.db.schema.skill import Skill __all__ = [ "ClassFeatureMap", "ClassFeature", "ClassFeatureOption", "ClassSpellSlotMap", "CharacterClass", "Skill", "ClassSkillMap", ] def skill_creator(fields): if isinstance(fields, ClassSkillMap): return fields return ClassSkillMap(**fields) class ClassSkillMap(BaseObject): __tablename__ = "class_skill_map" __table_args__ = (UniqueConstraint("skill_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_class_id: Mapped[int] = mapped_column(ForeignKey("character_class.id")) proficient: Mapped[bool] = mapped_column(default=True) expert: Mapped[bool] = mapped_column(default=False) skill = relationship("Skill", lazy="immediate") class ClassFeatureMap(BaseObject): __tablename__ = "class_feature_map" class_feature_id: Mapped[int] = mapped_column(ForeignKey("class_feature.id"), primary_key=True) character_class_id: Mapped[int] = mapped_column(ForeignKey("character_class.id"), primary_key=True) level: Mapped[int] = mapped_column(nullable=False, info={"min": 1, "max": 20}, default=1) feature = relationship("ClassFeature", uselist=False, viewonly=True, lazy="immediate") class ClassSpellSlotMap(BaseObject): __tablename__ = "class_spell_slot_map" id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True) character_class_id: Mapped[int] = mapped_column(ForeignKey("character_class.id")) class_level: Mapped[int] = mapped_column(nullable=False, info={"min": 1, "max": 20}, default=1) spell_level: Mapped[int] = mapped_column(nullable=False, info={"min": 1, "max": 9}, default=1) class ClassFeature(BaseObject): __tablename__ = "class_feature" id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True) name: Mapped[str] = mapped_column(nullable=False) options = relationship("ClassFeatureOption", cascade="all,delete,delete-orphan", lazy="immediate") def add_option(self, **kwargs): option = ClassFeatureOption(feature_id=self.id, **kwargs) if not self.options or option not in self.options: option.feature_id = self.id if not self.options: self.options = [option] else: self.options.append(option) return True return False def __repr__(self): return f"{self.id}: {self.name}" class ClassFeatureOption(BaseObject): __tablename__ = "class_feature_option" id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True) name: Mapped[str] = mapped_column(nullable=False) feature_id: Mapped[int] = mapped_column(ForeignKey("class_feature.id"), nullable=True) class CharacterClass(BaseObject): __tablename__ = "character_class" id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True) name: Mapped[str] = mapped_column(index=True, unique=True) hit_die_name: Mapped[str] = mapped_column(default="1d6") hit_die_stat_name: Mapped[str] = mapped_column(default="") starting_skills: int = mapped_column(nullable=False, default=0) features = relationship("ClassFeatureMap", cascade="all,delete,delete-orphan", lazy="immediate") spell_slots = relationship("ClassSpellSlotMap", cascade="all,delete,delete-orphan", lazy="immediate") _skills = relationship("ClassSkillMap", cascade="all,delete,delete-orphan", lazy="immediate") skills = association_proxy("_skills", "skill", creator=skill_creator) def add_skill(self, skill, expert=False): if not self.skills or skill not in self.skills: if not self.id: raise Exception("Cannot add a skill before the class has been persisted.") mapping = ClassSkillMap(character_class_id=self.id, skill_id=skill.id, proficient=True, expert=expert) self.skills.append(mapping) return True return False def add_feature(self, feature, level=1): if not self.features or feature not in self.features: mapping = ClassFeatureMap(character_class_id=self.id, class_feature_id=feature.id, level=level) if not self.features: self.features = [mapping] else: self.features.append(mapping) return True return False @property def spell_slots_by_level(self): by_level = defaultdict(list) for mapping in self.spell_slots: by_level[mapping.class_level].append(mapping) return by_level def spell_slots_at_level(self, level: int): return list(itertools.chain(*[mapping for lvl, mapping in self.spell_slots_by_level.items() if lvl <= level])) @property def features_by_level(self): by_level = defaultdict(list) for mapping in self.features: by_level[mapping.level].append(mapping.feature) return by_level def feature(self, name: str): for mapping in self.features: if mapping.feature.name.lower() == name.lower(): return mapping.feature return None def features_at_level(self, level: int): return list(itertools.chain(*[attrs for lvl, attrs in self.features_by_level.items() if lvl <= level]))