from __future__ import annotations import os, glob, random, uuid, json from dataclasses import asdict from typing import Any, Dict, List, Tuple, Optional import yaml from app.game.models.weapons import Weapon WEAPON_TEMPLATE_PATH = os.path.join("app","game","templates","weapons") # ---------------- Loader ---------------- def load_config_dir(dir_path: str = WEAPON_TEMPLATE_PATH) -> Dict[str, Any]: """ Load and merge the split YAML files from templates/weapons/. Expected keys across files: rarities, bases, affixes, implicits, grammar, balance, economy. """ if not os.path.isdir(dir_path): raise FileNotFoundError(f"Directory not found: {dir_path}") cfg: Dict[str, Any] = {} # Deterministic load order helps avoid surprises order = [ "rarities.yaml", "bases.yaml", "affixes.yaml", "implicits.yaml", "grammar.yaml", "tuning.yaml", ] filenames = [os.path.join(dir_path, f) for f in order if os.path.exists(os.path.join(dir_path, f))] if not filenames: # Fallback to load all .yaml if custom names were used filenames = sorted(glob.glob(os.path.join(dir_path, "*.yaml"))) for path in filenames: with open(path, "r", encoding="utf-8") as f: data = yaml.safe_load(f) or {} # shallow merge by top-level keys for k, v in data.items(): if k in cfg and isinstance(cfg[k], dict) and isinstance(v, dict): cfg[k].update(v) else: cfg[k] = v # Minimal validation required = ["rarities", "bases", "affixes", "grammar", "balance", "economy"] missing = [k for k in required if k not in cfg] if missing: raise ValueError(f"Missing required sections in merged config: {missing}") # Normalize shapes used by the generator # affixes should contain 'prefixes' and 'suffixes' if "prefixes" not in cfg["affixes"] or "suffixes" not in cfg["affixes"]: raise ValueError("affixes.yaml must contain 'prefixes' and 'suffixes' lists.") return cfg # ---------------- Generator ---------------- class WeaponGenerator: def __init__(self, rng: Optional[random.Random] = None): self.cfg = load_config_dir() self.rng = rng or random.Random() self._rarities = list(self.cfg["rarities"].keys()) self._rarity_weights = [self.cfg["rarities"][r]["weight"] for r in self._rarities] @staticmethod def _roll_float(rng: random.Random, lo: float, hi: float) -> float: return rng.uniform(lo, hi) def _weighted_choice(self, pairs: List[Tuple[Any, int]]): items, weights = zip(*pairs) return self.rng.choices(items, weights=weights, k=1)[0] def _choose_rarity(self) -> str: return self.rng.choices(self._rarities, weights=self._rarity_weights, k=1)[0] def _choose_base(self) -> Dict[str, Any]: return self.rng.choice(self.cfg["bases"]) def _roll_item_level(self, char_level: int) -> int: off = self.rng.randint(self.cfg["balance"]["item_level_offset"]["min"], self.cfg["balance"]["item_level_offset"]["max"]) lvl = max(self.cfg["balance"]["min_item_level"], char_level + off) return max(1, int(lvl)) def _scale_base_damage(self, base_min: int, base_max: int, item_level: int) -> Tuple[float, float]: g = self.cfg["balance"]["base_damage_growth_per_level"] growth = (1.0 + g) ** max(item_level - 1, 0) return base_min * growth, base_max * growth def _roll_affix_block(self, items: List[Dict[str, Any]], count: int, item_level: int) -> Tuple[List[Dict], Dict[str, float]]: chosen, stats = [], {} pool = items[:] self.rng.shuffle(pool) for pick in pool: if len(chosen) >= count: break rolled = [] for eff in pick.get("effects", []): lo, hi = eff["min"], eff["max"] per_lv = eff.get("per_level", 0.0) val = self._roll_float(self.rng, lo, hi) + per_lv * item_level rolled.append({"stat": eff["stat"], "value": val}) stats[eff["stat"]] = stats.get(eff["stat"], 0.0) + val chosen.append({"id": pick["id"], "name": pick["name"], "effects": rolled}) return chosen, stats def _apply_implicits(self, rarity: str, item_level: int) -> Tuple[List[Dict], Dict[str, float]]: table = self.cfg.get("implicits", {}).get(rarity, []) if not table: return [], {} imp = self.rng.choice(table) rolled, stats = [], {} for eff in imp.get("effects", []): val = self._roll_float(self.rng, eff["min"], eff["max"]) + eff.get("per_level", 0.0) * item_level rolled.append({"stat": eff["stat"], "value": val}) stats[eff["stat"]] = stats.get(eff["stat"], 0.0) + val return [{"id": imp["id"], "name": imp["name"], "effects": rolled}], stats def _compose_name(self, base_name: str, rarity: str, aff_prefixes: List[Dict], aff_suffixes: List[Dict]) -> str: g = self.cfg["grammar"] prefix_primary = aff_prefixes[0]["name"] if aff_prefixes else "" suffix_primary = aff_suffixes[0]["name"] if aff_suffixes else "" rarity_title = self.cfg["rarities"][rarity]["title"] token_values = { "rarity_title_opt": (rarity_title + " ") if rarity_title else "", "prefix_primary_opt": (prefix_primary + " ") if prefix_primary else "", "suffix_primary_opt": (" " + suffix_primary) if suffix_primary else "", "base_name": base_name, "crafted_by": self.rng.choice(g["lists"]["crafted_by"]), "mythic_noun": self.rng.choice(g["lists"]["mythic_noun"]), "epithet": self.rng.choice(g["lists"]["epithet"]), } patterns = [(p["template"], p["weight"]) for p in g["patterns"]] template = self._weighted_choice(patterns) name = template for k, v in token_values.items(): name = name.replace("{" + k + "}", v) name = " ".join(name.split()).replace(" ,", ",") return name def _compute_value(self, rarity: str, item_level: int, char_level: int, dmg_min: float, dmg_max: float, stats: Dict[str, float]) -> int: avg_dmg = 0.5 * (dmg_min + dmg_max) crit = stats.get("crit_chance_pct", 0.0) * 0.5 elem = stats.get("fire_damage_flat", 0.0) * 0.8 flat = stats.get("flat_damage", 0.0) item_power = avg_dmg + crit + elem + flat rarity_mul = self.cfg["rarities"][rarity]["value_multiplier"] raw_value = item_power * rarity_mul econ = self.cfg["economy"] delta = max(char_level, 1) - max(item_level, 1) if delta > 0: mult = (max(item_level, 1) / max(char_level, 1)) ** econ["underlevel_elasticity"] adj = max(mult, econ["min_sell_floor"] / max(raw_value, econ["min_sell_floor"])) elif delta < 0: over = min(abs(delta) * 0.05, econ["overlevel_bonus_pct"]) adj = 1.0 + over else: adj = 1.0 return int(max(econ["min_sell_floor"], raw_value * adj)) def generate(self, char_level: int) -> Dict[str, Any]: rarity = self._choose_rarity() base = self._choose_base() item_level = self._roll_item_level(char_level) smin, smax = self._scale_base_damage(base["min_dmg"], base["max_dmg"], item_level) rdef = self.cfg["rarities"][rarity] pslots, sslots = rdef["prefix_slots"], rdef["suffix_slots"] prefixes, pstats = self._roll_affix_block(self.cfg["affixes"]["prefixes"], pslots, item_level) suffixes, sstats = self._roll_affix_block(self.cfg["affixes"]["suffixes"], sslots, item_level) implicits, istats = self._apply_implicits(rarity, item_level) stats: Dict[str, float] = {} for source in (pstats, sstats, istats): for k, v in source.items(): stats[k] = stats.get(k, 0.0) + v rarity_mul = rdef["dmg_multiplier"] flat = stats.get("flat_damage", 0.0) min_dmg = (smin + flat) * rarity_mul max_dmg = (smax + flat) * rarity_mul name = self._compose_name(base["name"], rarity, prefixes, suffixes) value_sell = self._compute_value(rarity, item_level, char_level, min_dmg, max_dmg, stats) w = Weapon( uid=str(uuid.uuid4()), item_level=item_level, char_level=char_level, base_id=base["id"], base_name=base["name"], rarity=rarity, name=name, min_dmg=int(round(min_dmg)), max_dmg=int(round(max_dmg)), speed=float(base["speed"]), affixes=prefixes + suffixes, implicits=implicits, stats={k: round(v, 2) for k, v in stats.items()}, value_sell=value_sell, ) return asdict(w) # ---------------- Example CLI ---------------- # if __name__ == "__main__": # cfg = load_config_dir("templates/weapons") # gen = WeaponGenerator(cfg) # for lvl in (1, 5, 15, 30): # item = gen.generate(char_level=lvl) # print(json.dumps(item, indent=2))