1use crate::{Algorithm, Grid, Rng, Tile};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::path::Path;
5
6#[derive(Debug, Clone)]
7pub struct PrefabConfig {
8 pub max_prefabs: usize,
9 pub min_spacing: usize,
10 pub allow_rotation: bool,
11 pub allow_mirroring: bool,
12 pub weighted_selection: bool,
13 pub placement_mode: PrefabPlacementMode,
14 pub tags: Option<Vec<String>>,
15}
16
17impl Default for PrefabConfig {
18 fn default() -> Self {
19 Self {
20 max_prefabs: 3,
21 min_spacing: 5,
22 allow_rotation: true,
23 allow_mirroring: false,
24 weighted_selection: true,
25 placement_mode: PrefabPlacementMode::Overwrite,
26 tags: None,
27 }
28 }
29}
30
31#[derive(Debug, Clone, Copy)]
32pub enum PrefabPlacementMode {
33 Overwrite,
34 Merge,
35 PaintFloor,
36 PaintWall,
37}
38
39#[derive(Clone, Debug, Serialize, Deserialize)]
40pub struct PrefabLegendEntry {
41 pub tile: Option<String>,
42 pub marker: Option<String>,
43 pub mask: Option<String>,
44}
45
46#[derive(Clone, Debug, Serialize, Deserialize)]
47pub struct PrefabData {
48 pub name: String,
49 pub width: usize,
50 pub height: usize,
51 pub pattern: Vec<String>,
52 pub weight: f32,
53 pub tags: Vec<String>,
54 #[serde(default)]
55 pub legend: Option<HashMap<String, PrefabLegendEntry>>,
56}
57
58#[derive(Clone, Debug, Default)]
59pub struct PrefabCell {
60 pub tile: Option<Tile>,
61 pub marker: Option<String>,
62 pub mask: Option<String>,
63}
64
65#[derive(Clone)]
66pub struct Prefab {
67 pub name: String,
68 pub width: usize,
69 pub height: usize,
70 pub cells: Vec<PrefabCell>,
71 pub symbols: Vec<char>,
72 pub legend: Option<HashMap<char, PrefabLegendEntry>>,
73 pub weight: f32,
74 pub tags: Vec<String>,
75}
76
77impl Prefab {
78 pub fn new(pattern: &[&str]) -> Self {
79 let height = pattern.len();
80 let width = pattern.first().map(|s| s.len()).unwrap_or(0);
81 let (cells, symbols) = parse_pattern_with_legend(pattern, None);
82 Self {
83 name: "unnamed".to_string(),
84 width,
85 height,
86 cells,
87 symbols,
88 legend: None,
89 weight: 1.0,
90 tags: Vec::new(),
91 }
92 }
93
94 pub fn from_data(data: PrefabData) -> Self {
95 let width = data.width;
96 let height = data.height;
97 let legend = data.legend.as_ref().map(convert_legend);
98 let (cells, symbols) = parse_pattern_with_legend(&data.pattern, legend.as_ref());
99
100 Self {
101 name: data.name,
102 width,
103 height,
104 cells,
105 symbols,
106 legend,
107 weight: data.weight,
108 tags: data.tags,
109 }
110 }
111
112 pub fn rect(w: usize, h: usize) -> Self {
113 Self {
114 name: format!("rect_{}x{}", w, h),
115 width: w,
116 height: h,
117 cells: vec![
118 PrefabCell {
119 tile: Some(Tile::Floor),
120 marker: None,
121 mask: None,
122 };
123 w * h
124 ],
125 symbols: vec!['.'; w * h],
126 legend: None,
127 weight: 1.0,
128 tags: vec!["rectangle".to_string()],
129 }
130 }
131
132 pub fn rotated(&self) -> Self {
134 let mut rotated_cells = vec![PrefabCell::default(); self.width * self.height];
135 let mut rotated_symbols = vec!['#'; self.width * self.height];
136 for y in 0..self.height {
137 for x in 0..self.width {
138 let old_idx = y * self.width + x;
139 let new_x = self.height - 1 - y;
140 let new_y = x;
141 let new_idx = new_y * self.height + new_x;
142 rotated_cells[new_idx] = self.cells[old_idx].clone();
143 rotated_symbols[new_idx] = self.symbols[old_idx];
144 }
145 }
146
147 Self {
148 name: format!("{}_rot90", self.name),
149 width: self.height,
150 height: self.width,
151 cells: rotated_cells,
152 symbols: rotated_symbols,
153 legend: self.legend.clone(),
154 weight: self.weight,
155 tags: self.tags.clone(),
156 }
157 }
158
159 pub fn mirrored_horizontal(&self) -> Self {
161 let mut mirrored_cells = vec![PrefabCell::default(); self.width * self.height];
162 let mut mirrored_symbols = vec!['#'; self.width * self.height];
163 for y in 0..self.height {
164 for x in 0..self.width {
165 let old_idx = y * self.width + x;
166 let new_x = self.width - 1 - x;
167 let new_idx = y * self.width + new_x;
168 mirrored_cells[new_idx] = self.cells[old_idx].clone();
169 mirrored_symbols[new_idx] = self.symbols[old_idx];
170 }
171 }
172
173 Self {
174 name: format!("{}_mirror_h", self.name),
175 width: self.width,
176 height: self.height,
177 cells: mirrored_cells,
178 symbols: mirrored_symbols,
179 legend: self.legend.clone(),
180 weight: self.weight,
181 tags: self.tags.clone(),
182 }
183 }
184
185 pub fn mirrored_vertical(&self) -> Self {
187 let mut mirrored_cells = vec![PrefabCell::default(); self.width * self.height];
188 let mut mirrored_symbols = vec!['#'; self.width * self.height];
189 for y in 0..self.height {
190 for x in 0..self.width {
191 let old_idx = y * self.width + x;
192 let new_y = self.height - 1 - y;
193 let new_idx = new_y * self.width + x;
194 mirrored_cells[new_idx] = self.cells[old_idx].clone();
195 mirrored_symbols[new_idx] = self.symbols[old_idx];
196 }
197 }
198
199 Self {
200 name: format!("{}_mirror_v", self.name),
201 width: self.width,
202 height: self.height,
203 cells: mirrored_cells,
204 symbols: mirrored_symbols,
205 legend: self.legend.clone(),
206 weight: self.weight,
207 tags: self.tags.clone(),
208 }
209 }
210
211 pub fn get(&self, x: usize, y: usize) -> bool {
212 self.cell_tile(x, y) == Some(Tile::Floor)
213 }
214
215 pub fn cell_tile(&self, x: usize, y: usize) -> Option<Tile> {
216 if x < self.width && y < self.height {
217 self.cells[y * self.width + x].tile
218 } else {
219 None
220 }
221 }
222
223 pub fn cell_marker(&self, x: usize, y: usize) -> Option<&str> {
224 if x < self.width && y < self.height {
225 self.cells[y * self.width + x].marker.as_deref()
226 } else {
227 None
228 }
229 }
230
231 pub fn cell_mask(&self, x: usize, y: usize) -> Option<&str> {
232 if x < self.width && y < self.height {
233 self.cells[y * self.width + x].mask.as_deref()
234 } else {
235 None
236 }
237 }
238
239 pub fn has_tag(&self, tag: &str) -> bool {
240 self.tags.contains(&tag.to_string())
241 }
242}
243
244#[derive(Debug, Clone, Serialize, Deserialize)]
245pub struct PrefabLibraryData {
246 pub prefabs: Vec<PrefabData>,
247}
248
249#[derive(Clone)]
250pub struct PrefabLibrary {
251 prefabs: Vec<Prefab>,
252 by_tag: HashMap<String, Vec<usize>>,
253}
254
255impl PrefabLibrary {
256 pub fn new() -> Self {
257 Self {
258 prefabs: Vec::new(),
259 by_tag: HashMap::new(),
260 }
261 }
262
263 pub fn add_prefab(&mut self, prefab: Prefab) {
264 let index = self.prefabs.len();
265
266 for tag in &prefab.tags {
267 self.by_tag.entry(tag.clone()).or_default().push(index);
268 }
269
270 self.prefabs.push(prefab);
271 }
272
273 pub fn load_from_json<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn std::error::Error>> {
274 let content = std::fs::read_to_string(path)?;
275 let data: PrefabLibraryData = serde_json::from_str(&content)?;
276
277 let mut library = Self::new();
278 for prefab_data in data.prefabs {
279 library.add_prefab(Prefab::from_data(prefab_data));
280 }
281
282 Ok(library)
283 }
284
285 pub fn load_from_paths<I, P>(paths: I) -> Result<Self, Box<dyn std::error::Error>>
286 where
287 I: IntoIterator<Item = P>,
288 P: AsRef<Path>,
289 {
290 let mut library = Self::new();
291 for path in paths {
292 let loaded = Self::load_from_json(path)?;
293 library.extend_from(loaded);
294 }
295 Ok(library)
296 }
297
298 pub fn load_from_dir<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn std::error::Error>> {
299 let mut entries: Vec<std::path::PathBuf> = std::fs::read_dir(path)?
300 .filter_map(|entry| entry.ok().map(|e| e.path()))
301 .filter(|p| p.extension().and_then(|e| e.to_str()) == Some("json"))
302 .collect();
303 entries.sort();
304 if entries.is_empty() {
305 return Ok(Self::new());
306 }
307 Self::load_from_paths(entries)
308 }
309
310 pub fn save_to_json<P: AsRef<Path>>(&self, path: P) -> Result<(), Box<dyn std::error::Error>> {
311 let data = PrefabLibraryData {
312 prefabs: self
313 .prefabs
314 .iter()
315 .map(|p| PrefabData {
316 name: p.name.clone(),
317 width: p.width,
318 height: p.height,
319 pattern: self.prefab_to_pattern(p),
320 weight: p.weight,
321 tags: p.tags.clone(),
322 legend: p.legend.as_ref().map(convert_legend_to_strings),
323 })
324 .collect(),
325 };
326
327 let content = serde_json::to_string_pretty(&data)?;
328 std::fs::write(path, content)?;
329 Ok(())
330 }
331
332 fn prefab_to_pattern(&self, prefab: &Prefab) -> Vec<String> {
333 let mut pattern = Vec::new();
334 for y in 0..prefab.height {
335 let mut row = String::new();
336 for x in 0..prefab.width {
337 let idx = y * prefab.width + x;
338 let symbol = prefab.symbols.get(idx).copied().unwrap_or_else(|| {
339 if prefab.get(x, y) {
340 '.'
341 } else {
342 '#'
343 }
344 });
345 row.push(symbol);
346 }
347 pattern.push(row);
348 }
349 pattern
350 }
351
352 pub fn get_prefabs(&self) -> &[Prefab] {
353 &self.prefabs
354 }
355
356 pub fn extend_from(&mut self, other: PrefabLibrary) {
357 for prefab in other.prefabs {
358 self.add_prefab(prefab);
359 }
360 }
361
362 pub fn get_by_tag(&self, tag: &str) -> Vec<&Prefab> {
363 self.by_tag
364 .get(tag)
365 .map(|indices| indices.iter().map(|&i| &self.prefabs[i]).collect())
366 .unwrap_or_default()
367 }
368
369 pub fn select_weighted(&self, rng: &mut Rng, tag: Option<&str>) -> Option<&Prefab> {
370 let tags = tag.map(|t| vec![t.to_string()]);
371 self.select_with_tags(rng, tags.as_deref(), true)
372 }
373
374 pub fn select_with_tags(
375 &self,
376 rng: &mut Rng,
377 tags: Option<&[String]>,
378 weighted: bool,
379 ) -> Option<&Prefab> {
380 let candidates = if let Some(tags) = tags {
381 self.get_by_any_tag(tags)
382 } else {
383 self.prefabs.iter().collect()
384 };
385
386 if candidates.is_empty() {
387 return None;
388 }
389
390 if !weighted {
391 return rng.pick(&candidates).copied();
392 }
393
394 let total_weight: f32 = candidates.iter().map(|p| p.weight).sum();
395 if total_weight <= 0.0 {
396 return rng.pick(&candidates).copied();
397 }
398
399 let mut target = rng.random() as f32 * total_weight;
400 for prefab in &candidates {
401 target -= prefab.weight;
402 if target <= 0.0 {
403 return Some(prefab);
404 }
405 }
406
407 candidates.last().copied()
408 }
409
410 pub fn get_by_any_tag(&self, tags: &[String]) -> Vec<&Prefab> {
411 if tags.is_empty() {
412 return Vec::new();
413 }
414 let mut indices = Vec::new();
415 for tag in tags {
416 if let Some(found) = self.by_tag.get(tag) {
417 indices.extend_from_slice(found);
418 }
419 }
420 if indices.is_empty() {
421 return Vec::new();
422 }
423 indices.sort_unstable();
424 indices.dedup();
425 indices.iter().map(|&i| &self.prefabs[i]).collect()
426 }
427
428 pub fn create_default() -> Self {
429 let mut library = Self::new();
430
431 let mut room = Prefab::rect(5, 5);
433 room.name = "small_room".to_string();
434 room.tags = vec!["room".to_string(), "small".to_string()];
435 library.add_prefab(room);
436
437 let mut corridor = Prefab::new(&[".....", "#####"]);
438 corridor.name = "corridor".to_string();
439 corridor.tags = vec!["corridor".to_string()];
440 library.add_prefab(corridor);
441
442 let mut cross = Prefab::new(&["#.#", "...", "#.#"]);
443 cross.name = "cross".to_string();
444 cross.tags = vec!["junction".to_string()];
445 library.add_prefab(cross);
446
447 library
448 }
449}
450
451impl Default for PrefabLibrary {
452 fn default() -> Self {
453 Self::create_default()
454 }
455}
456
457#[derive(Debug, Clone, Default)]
458pub struct PrefabTransform {
459 pub rotation: u8, pub mirror_h: bool,
461 pub mirror_v: bool,
462}
463
464impl PrefabTransform {
465 pub fn apply(&self, prefab: &Prefab) -> Prefab {
466 let mut result = prefab.clone();
467
468 if self.mirror_h {
470 result = result.mirrored_horizontal();
471 }
472 if self.mirror_v {
473 result = result.mirrored_vertical();
474 }
475
476 for _ in 0..self.rotation {
478 result = result.rotated();
479 }
480
481 result
482 }
483
484 pub fn random(rng: &mut Rng, allow_rotation: bool, allow_mirroring: bool) -> Self {
485 Self {
486 rotation: if allow_rotation {
487 rng.range(0, 4) as u8
488 } else {
489 0
490 },
491 mirror_h: allow_mirroring && rng.chance(0.5),
492 mirror_v: allow_mirroring && rng.chance(0.5),
493 }
494 }
495}
496pub struct PrefabPlacer {
497 config: PrefabConfig,
498 library: PrefabLibrary,
499}
500
501impl PrefabPlacer {
502 pub fn new(config: PrefabConfig, library: PrefabLibrary) -> Self {
503 Self { config, library }
504 }
505
506 pub fn with_library(library: PrefabLibrary) -> Self {
507 Self::new(PrefabConfig::default(), library)
508 }
509}
510
511impl Default for PrefabPlacer {
512 fn default() -> Self {
513 Self::with_library(PrefabLibrary::default())
514 }
515}
516
517impl Algorithm<Tile> for PrefabPlacer {
518 fn generate(&self, grid: &mut Grid<Tile>, seed: u64) {
519 self.generate_internal(grid, seed, None);
520 }
521
522 fn name(&self) -> &'static str {
523 "PrefabPlacer"
524 }
525}
526
527impl PrefabPlacer {
528 pub fn generate_with_semantic(
529 &self,
530 grid: &mut Grid<Tile>,
531 seed: u64,
532 semantic: &mut crate::semantic::SemanticLayers,
533 ) {
534 self.generate_internal(grid, seed, Some(semantic));
535 }
536
537 fn generate_internal(
538 &self,
539 grid: &mut Grid<Tile>,
540 seed: u64,
541 mut semantic: Option<&mut crate::semantic::SemanticLayers>,
542 ) {
543 let mut rng = Rng::new(seed);
544 let mut placed: Vec<(usize, usize, usize, usize)> = Vec::new();
545
546 for _ in 0..self.config.max_prefabs * 10 {
547 if placed.len() >= self.config.max_prefabs {
548 break;
549 }
550
551 let base_prefab = if let Some(prefab) = self.library.select_with_tags(
552 &mut rng,
553 self.config.tags.as_deref(),
554 self.config.weighted_selection,
555 ) {
556 prefab
557 } else {
558 continue;
559 };
560
561 let transform = PrefabTransform::random(
563 &mut rng,
564 self.config.allow_rotation,
565 self.config.allow_mirroring,
566 );
567 let prefab = transform.apply(base_prefab);
568
569 if prefab.width + 2 >= grid.width() || prefab.height + 2 >= grid.height() {
570 continue;
571 }
572
573 let x = rng.range_usize(1, grid.width() - prefab.width - 1);
574 let y = rng.range_usize(1, grid.height() - prefab.height - 1);
575
576 let overlaps = placed.iter().any(|&(px, py, pw, ph)| {
577 let s = self.config.min_spacing;
578 !(x + prefab.width + s < px
579 || px + pw + s < x
580 || y + prefab.height + s < py
581 || py + ph + s < y)
582 });
583
584 if overlaps {
585 continue;
586 }
587
588 for py in 0..prefab.height {
589 for px in 0..prefab.width {
590 let cell_tile = prefab.cell_tile(px, py);
591 let cell_marker = prefab.cell_marker(px, py);
592 let cell_mask = prefab.cell_mask(px, py);
593 let gx = (x + px) as i32;
594 let gy = (y + py) as i32;
595
596 let mut applied = false;
597 if let Some(tile) = cell_tile {
598 let current = *grid.get(gx, gy).unwrap_or(&Tile::Wall);
599 let should_place = match self.config.placement_mode {
600 PrefabPlacementMode::Overwrite => true,
601 PrefabPlacementMode::Merge => matches!(current, Tile::Wall),
602 PrefabPlacementMode::PaintFloor => matches!(current, Tile::Floor),
603 PrefabPlacementMode::PaintWall => matches!(current, Tile::Wall),
604 };
605 if should_place {
606 grid.set(gx, gy, tile);
607 applied = true;
608 }
609 }
610
611 if let Some(layers) = semantic.as_deref_mut() {
612 let marker_allowed = cell_tile.is_none() || applied;
613 if marker_allowed {
614 if let Some(tag) = cell_marker {
615 layers.markers.push(crate::semantic::Marker::with_tag(
616 gx as u32,
617 gy as u32,
618 tag.to_string(),
619 ));
620 }
621 if let Some(mask) = cell_mask {
622 apply_prefab_mask(&mut layers.masks, gx, gy, mask);
623 }
624 }
625 }
626 }
627 }
628 placed.push((x, y, prefab.width, prefab.height));
629 }
630 }
631}
632
633fn parse_pattern_with_legend(
634 pattern: &[impl AsRef<str>],
635 legend: Option<&HashMap<char, PrefabLegendEntry>>,
636) -> (Vec<PrefabCell>, Vec<char>) {
637 let mut cells = Vec::new();
638 let mut symbols = Vec::new();
639 for row in pattern {
640 for ch in row.as_ref().chars() {
641 let cell = if let Some(legend) = legend {
642 legend
643 .get(&ch)
644 .map(parse_legend_entry)
645 .unwrap_or_else(|| default_cell_from_symbol(ch))
646 } else {
647 default_cell_from_symbol(ch)
648 };
649 cells.push(cell);
650 symbols.push(ch);
651 }
652 }
653 (cells, symbols)
654}
655
656fn default_cell_from_symbol(ch: char) -> PrefabCell {
657 match ch {
658 '.' => PrefabCell {
659 tile: Some(Tile::Floor),
660 marker: None,
661 mask: None,
662 },
663 '#' => PrefabCell::default(),
664 _ => PrefabCell::default(),
665 }
666}
667
668fn parse_legend_entry(entry: &PrefabLegendEntry) -> PrefabCell {
669 PrefabCell {
670 tile: parse_tile_name(entry.tile.as_deref()),
671 marker: entry.marker.clone(),
672 mask: entry.mask.clone(),
673 }
674}
675
676fn parse_tile_name(value: Option<&str>) -> Option<Tile> {
677 let value = value?;
678 match value.trim().to_ascii_lowercase().as_str() {
679 "floor" | "f" => Some(Tile::Floor),
680 "wall" | "w" => Some(Tile::Wall),
681 "empty" | "none" | "skip" => None,
682 _ => None,
683 }
684}
685
686fn convert_legend(legend: &HashMap<String, PrefabLegendEntry>) -> HashMap<char, PrefabLegendEntry> {
687 let mut out = HashMap::new();
688 for (key, value) in legend {
689 let mut chars = key.chars();
690 if let Some(ch) = chars.next() {
691 out.insert(ch, value.clone());
692 }
693 }
694 out
695}
696
697fn convert_legend_to_strings(
698 legend: &HashMap<char, PrefabLegendEntry>,
699) -> HashMap<String, PrefabLegendEntry> {
700 legend
701 .iter()
702 .map(|(ch, entry)| (ch.to_string(), entry.clone()))
703 .collect()
704}
705
706fn apply_prefab_mask(masks: &mut crate::semantic::Masks, x: i32, y: i32, mask: &str) {
707 if x < 0 || y < 0 {
708 return;
709 }
710 let (x, y) = (x as usize, y as usize);
711 if y >= masks.height || x >= masks.width {
712 return;
713 }
714 match mask.trim().to_ascii_lowercase().as_str() {
715 "no_spawn" | "nospawn" | "reserved" => {
716 if let Some(row) = masks.no_spawn.get_mut(y) {
717 if let Some(cell) = row.get_mut(x) {
718 *cell = true;
719 }
720 }
721 }
722 _ => {}
723 }
724}