scripts/item_generator.py

Part of Dungeon Explorer.

  1"""Loot generator: roll random items from TOML data files.
  2
  3Usage:
  4    gen = ItemGenerator("data/items.toml", "data/affixes.toml")
  5    item = gen.roll_item(ilvl=15, rarity=Rarity.RARE)
  6"""
  7
  8import random
  9from pathlib import Path
 10
 11try:
 12    import tomllib
 13except ImportError:
 14    import tomli as tomllib  # type: ignore[no-redef]
 15
 16from .item import Affix, Item, Rarity
 17
 18
 19def _lerp(a: float, b: float, t: float) -> float:
 20    return a + (b - a) * max(0.0, min(1.0, t))
 21
 22
 23class ItemGenerator:
 24    """Generates random items from base templates and affix tables."""
 25
 26    def __init__(self, items_path: str | Path, affixes_path: str | Path):
 27        with open(items_path, "rb") as f:
 28            self._templates = tomllib.load(f)
 29        with open(affixes_path, "rb") as f:
 30            self._affixes = tomllib.load(f)
 31
 32    def roll_item(
 33        self,
 34        template_key: str | None = None,
 35        ilvl: int = 1,
 36        rarity: Rarity = Rarity.COMMON,
 37        rng: random.Random | None = None,
 38    ) -> Item:
 39        """Roll a random item.
 40
 41        Args:
 42            template_key: Specific template (e.g. "sword"). None = random.
 43            ilvl: Item level for stat scaling.
 44            rarity: Rarity tier.
 45            rng: Random instance for reproducibility.
 46        """
 47        rng = rng or random.Random()
 48
 49        # Pick template
 50        if template_key is None:
 51            template_key = rng.choice(list(self._templates.keys()))
 52        template = self._templates[template_key]
 53
 54        # Base stats
 55        base_stats: dict[str, float] = {}
 56        for key in ("damage", "defence", "speed", "heal", "speed_bonus", "block_chance"):
 57            val = template.get(key)
 58            if val is not None:
 59                if isinstance(val, list) and len(val) == 2:
 60                    base_stats[key] = rng.uniform(val[0], val[1])
 61                else:
 62                    base_stats[key] = float(val)
 63
 64        # Max affixes for this rarity
 65        max_affixes_table = template.get("max_affixes", [0, 1, 2, 3, 4])
 66        max_affixes = max_affixes_table[min(int(rarity), len(max_affixes_table) - 1)]
 67
 68        # Roll affixes
 69        slot = template.get("slot", "consumable")
 70        prefixes = self._roll_affixes("prefix", slot, ilvl, max_affixes, rng)
 71        suffixes = self._roll_affixes("suffix", slot, ilvl, max_affixes, rng)
 72
 73        # Limit total affixes
 74        total = prefixes + suffixes
 75        if len(total) > max_affixes:
 76            total = rng.sample(total, max_affixes)
 77            prefixes = [a for a in total if a.affix_type == "prefix"]
 78            suffixes = [a for a in total if a.affix_type == "suffix"]
 79
 80        return Item(
 81            template_key=template_key,
 82            display_base=template.get("display", template_key.title()),
 83            slot=slot,
 84            rarity=rarity,
 85            ilvl=ilvl,
 86            prefixes=prefixes,
 87            suffixes=suffixes,
 88            base_stats=base_stats,
 89            stack_size=template.get("stack_size", 1),
 90        )
 91
 92    def _roll_affixes(
 93        self,
 94        affix_type: str,
 95        slot: str,
 96        ilvl: int,
 97        max_count: int,
 98        rng: random.Random,
 99    ) -> list[Affix]:
100        """Roll affixes of a given type for a slot and ilvl."""
101        candidates = []
102        for key, data in self._affixes.items():
103            if data.get("type") != affix_type:
104                continue
105            if slot not in data.get("slots", []):
106                continue
107            ilvl_min = data.get("ilvl_min", 1)
108            ilvl_max = data.get("ilvl_max", 100)
109            if ilvl < ilvl_min:
110                continue
111            candidates.append((key, data, ilvl_min, ilvl_max))
112
113        if not candidates:
114            return []
115
116        # Weighted random selection (up to max_count // 2 of this type)
117        count = min(len(candidates), max(1, max_count // 2))
118        weights = [c[1].get("weight", 10) for c in candidates]
119        selected = _weighted_sample(candidates, weights, count, rng)
120
121        affixes = []
122        for key, data, ilvl_min, ilvl_max in selected:
123            stat_range = data.get("range", [0, 0])
124            t = (ilvl - ilvl_min) / max(1, ilvl_max - ilvl_min)
125            value = _lerp(stat_range[0], stat_range[1], t)
126            # Add some randomness (+-15%)
127            value *= rng.uniform(0.85, 1.15)
128            affixes.append(
129                Affix(
130                    name=data.get("display", key.title()),
131                    affix_type=affix_type,
132                    stat=data["stat"],
133                    value=round(value, 3),
134                )
135            )
136
137        return affixes
138
139    def roll_from_loot_table(
140        self,
141        loot_table: dict,
142        dungeon_level: int,
143        rng: random.Random | None = None,
144    ) -> list[Item]:
145        """Roll drops from a loot table entry.
146
147        Returns list of dropped items (may be empty).
148        """
149        rng = rng or random.Random()
150        drops = []
151        for entry in loot_table.get("drops", []):
152            if rng.random() > entry.get("chance", 0.1):
153                continue
154            # Roll rarity
155            rarity_weights = entry.get("rarity_weights", [60, 25, 10, 4, 1])
156            rarity = Rarity(_weighted_index(rarity_weights, rng))
157            item = self.roll_item(
158                template_key=entry["item"],
159                ilvl=dungeon_level,
160                rarity=rarity,
161                rng=rng,
162            )
163            drops.append(item)
164        return drops
165
166
167def _weighted_sample(items: list, weights: list[float], count: int, rng: random.Random) -> list:
168    """Weighted random sample without replacement."""
169    pool = list(zip(items, weights, strict=True))
170    result = []
171    for _ in range(min(count, len(pool))):
172        total = sum(w for _, w in pool)
173        if total <= 0:
174            break
175        r = rng.uniform(0, total)
176        cumulative = 0.0
177        for i, (item, w) in enumerate(pool):
178            cumulative += w
179            if cumulative >= r:
180                result.append(item)
181                pool.pop(i)
182                break
183    return result
184
185
186def _weighted_index(weights: list[float], rng: random.Random) -> int:
187    """Return a random index weighted by the given weights."""
188    total = sum(weights)
189    if total <= 0:
190        return 0
191    r = rng.uniform(0, total)
192    cumulative = 0.0
193    for i, w in enumerate(weights):
194        cumulative += w
195        if cumulative >= r:
196            return i
197    return len(weights) - 1