1use crate::ext::io::*;
3use crate::ext::psb::*;
4use crate::scripts::base::*;
5use crate::try_option;
6use crate::types::*;
7use crate::utils::img::*;
8use crate::utils::psd::*;
9use crate::utils::struct_pack::*;
10use anyhow::Result;
11use emote_psb::PsbReader;
12use libtlg_rs::*;
13use std::collections::HashMap;
14use std::io::{Read, Seek};
15use std::path::Path;
16
17#[derive(Debug)]
18pub struct PImgBuilder {}
20
21impl PImgBuilder {
22 pub const fn new() -> Self {
24 Self {}
25 }
26}
27
28impl ScriptBuilder for PImgBuilder {
29 fn default_encoding(&self) -> Encoding {
30 Encoding::Utf8
31 }
32
33 fn build_script(
34 &self,
35 buf: Vec<u8>,
36 filename: &str,
37 _encoding: Encoding,
38 _archive_encoding: Encoding,
39 config: &ExtraConfig,
40 _archive: Option<&Box<dyn Script>>,
41 ) -> Result<Box<dyn Script>> {
42 Ok(Box::new(PImg::new(MemReader::new(buf), filename, config)?))
43 }
44
45 fn build_script_from_file(
46 &self,
47 filename: &str,
48 _encoding: Encoding,
49 _archive_encoding: Encoding,
50 config: &ExtraConfig,
51 _archive: Option<&Box<dyn Script>>,
52 ) -> Result<Box<dyn Script>> {
53 if filename == "-" {
54 let data = crate::utils::files::read_file(filename)?;
55 Ok(Box::new(PImg::new(MemReader::new(data), filename, config)?))
56 } else {
57 let f = std::fs::File::open(filename)?;
58 let reader = std::io::BufReader::new(f);
59 Ok(Box::new(PImg::new(reader, filename, config)?))
60 }
61 }
62
63 fn build_script_from_reader(
64 &self,
65 reader: Box<dyn ReadSeek>,
66 filename: &str,
67 _encoding: Encoding,
68 _archive_encoding: Encoding,
69 config: &ExtraConfig,
70 _archive: Option<&Box<dyn Script>>,
71 ) -> Result<Box<dyn Script>> {
72 Ok(Box::new(PImg::new(reader, filename, config)?))
73 }
74
75 fn extensions(&self) -> &'static [&'static str] {
76 &["pimg"]
77 }
78
79 fn script_type(&self) -> &'static ScriptType {
80 &ScriptType::EmotePimg
81 }
82
83 fn is_this_format(&self, filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
84 if Path::new(filename)
85 .extension()
86 .map(|ext| ext.to_ascii_lowercase() == "pimg")
87 .unwrap_or(false)
88 && buf_len >= 4
89 && buf.starts_with(b"PSB\0")
90 {
91 return Some(255);
92 }
93 None
94 }
95
96 fn is_image(&self) -> bool {
97 true
98 }
99}
100
101struct PImgLayer<'a> {
102 data: &'a PsbValueFixed,
103 name: &'a str,
104 layer_id: i64,
105 layer_type: i64,
107 left: u32,
108 top: u32,
109 width: u32,
110 height: u32,
111 opacity: u8,
112 visible: bool,
113 type_: i64,
114 children: Vec<PImgLayer<'a>>,
115}
116
117impl std::fmt::Debug for PImgLayer<'_> {
118 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119 f.debug_struct("PImgLayer")
120 .field("layer_id", &self.layer_id)
121 .field("layer_type", &self.layer_type)
122 .field("name", &self.name)
123 .field("left", &self.left)
124 .field("top", &self.top)
125 .field("width", &self.width)
126 .field("height", &self.height)
127 .field("opacity", &self.opacity)
128 .field("visible", &self.visible)
129 .field("type", &self.type_)
130 .field("children", &self.children)
131 .finish()
132 }
133}
134
135impl<'a> PImgLayer<'a> {
136 pub fn new(data: &'a PsbValueFixed, layers: &'a PsbValueFixed) -> Result<Self> {
137 let layer_id = data["layer_id"]
138 .as_i64()
139 .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid layer_id"))?;
140 let layer_type = data["layer_type"]
141 .as_i64()
142 .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid layer_type"))?;
143 let name = data["name"]
144 .as_str()
145 .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid name"))?;
146 let left = data["left"].as_u32();
147 let top = data["top"].as_u32();
148 let width = data["width"].as_u32();
149 let height = data["height"].as_u32();
150 let (left, top, width, height) = if layer_type != 0 {
151 (
152 left.unwrap_or(0),
153 top.unwrap_or(0),
154 width.unwrap_or(0),
155 height.unwrap_or(0),
156 )
157 } else {
158 (
159 left.ok_or_else(|| anyhow::anyhow!("Layer does not have a valid left"))?,
160 top.ok_or_else(|| anyhow::anyhow!("Layer does not have a valid top"))?,
161 width.ok_or_else(|| anyhow::anyhow!("Layer does not have a valid width"))?,
162 height.ok_or_else(|| anyhow::anyhow!("Layer does not have a valid height"))?,
163 )
164 };
165 let opacity = data["opacity"]
166 .as_u8()
167 .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid opacity"))?;
168 let visible = data["visible"].as_i64().unwrap_or(1) != 0;
169 let type_ = data["type"]
170 .as_i64()
171 .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid type"))?;
172 let mut children = Vec::new();
173 for layer in layers.members() {
174 if layer_type == 2 || layer_type == 1 {
175 if let Some(parent_id) = layer["group_layer_id"].as_i64() {
176 if parent_id == layer_id {
177 children.push(PImgLayer::new(layer, layers)?);
178 }
179 }
180 } else if layer_type == 0 {
181 if let Some(base_id) = layer["diff_id"].as_i64() {
182 if base_id == layer_id {
183 children.push(PImgLayer::new(layer, layers)?);
184 }
185 }
186 }
187 }
188 Ok(Self {
189 data,
190 layer_id,
191 layer_type,
192 name,
193 left,
194 top,
195 width,
196 height,
197 opacity,
198 visible,
199 type_,
200 children,
201 })
202 }
203
204 fn len(&self) -> usize {
205 1 + self.children.iter().map(|c| c.len()).sum::<usize>()
206 }
207
208 fn load_img(&self, img: &PImg) -> Result<ImageData> {
209 if self.layer_type == 2 || self.layer_type == 1 {
210 anyhow::bail!("Group layers do not have image data");
211 }
212 if self.layer_id == -1 {
213 Ok(ImageData {
215 width: self.width,
216 height: self.height,
217 color_type: ImageColorType::Rgba,
218 depth: 8,
219 data: vec![0u8; (self.width * self.height * 4) as usize],
220 })
221 } else {
222 let tlg = img.load_img(self.layer_id).map_err(|e| {
223 anyhow::anyhow!("Failed to load image for layer_id {}: {}", self.layer_id, e)
224 })?;
225 let mut img = ImageData {
226 width: tlg.width,
227 height: tlg.height,
228 color_type: match tlg.color {
229 TlgColorType::Bgr24 => ImageColorType::Bgr,
230 TlgColorType::Bgra32 => ImageColorType::Bgra,
231 TlgColorType::Grayscale8 => ImageColorType::Grayscale,
232 },
233 depth: 8,
234 data: tlg.data.clone(),
235 };
236 convert_to_rgba(&mut img)?;
237 Ok(img)
238 }
239 }
240
241 fn save_to_psd(&self, img: &PImg, psd: &mut PsdWriter, base: &mut ImageData) -> Result<()> {
242 if self.children.is_empty() {
243 let img = self.load_img(img)?;
244 let mut visible = self.visible;
245 if !self.data["diff_id"].is_none() {
246 visible = false; }
248 if visible {
249 draw_on_img_with_opacity(base, &img, self.left, self.top, self.opacity)?;
250 }
251 let layer_name_source_setting = LayerNameSourceSetting {
252 id: self.layer_id as i32,
253 };
254 let mut packed = Vec::new();
255 layer_name_source_setting.pack(&mut packed, true, Encoding::Utf8, &None)?;
256 let additional_info = vec![AdditionalLayerInfo {
257 signature: *IMAGE_RESOURCE_SIGNATURE,
258 key: *LAYER_NAME_SOURCE_SETTING_KEY,
259 data: packed,
260 }];
261 let option = PsdLayerOption {
262 visible,
263 opacity: self.opacity,
264 additional_info,
265 };
266 psd.add_layer(self.name, self.left, self.top, img, Some(option))?;
267 } else {
268 psd.add_layer_group_end()?;
269 if self.layer_type == 0 {
270 let img = self.load_img(img)?;
271 let visible = self.visible;
272 if visible {
273 draw_on_img_with_opacity(base, &img, self.left, self.top, self.opacity)?;
274 }
275 let layer_name_source_setting = LayerNameSourceSetting {
276 id: self.layer_id as i32,
277 };
278 let mut packed = Vec::new();
279 layer_name_source_setting.pack(&mut packed, true, Encoding::Utf8, &None)?;
280 let additional_info = vec![AdditionalLayerInfo {
281 signature: *IMAGE_RESOURCE_SIGNATURE,
282 key: *LAYER_NAME_SOURCE_SETTING_KEY,
283 data: packed,
284 }];
285 let option = PsdLayerOption {
286 visible,
287 opacity: self.opacity,
288 additional_info,
289 };
290 psd.add_layer(self.name, self.left, self.top, img, Some(option))?;
291 }
292 for child in &self.children {
293 child.save_to_psd(img, psd, base)?;
294 }
295 let layer_name_source_setting = LayerNameSourceSetting {
296 id: self.layer_id as i32,
297 };
298 let mut packed = Vec::new();
299 layer_name_source_setting.pack(&mut packed, true, Encoding::Utf8, &None)?;
300 let additional_info = vec![AdditionalLayerInfo {
301 signature: *IMAGE_RESOURCE_SIGNATURE,
302 key: *LAYER_NAME_SOURCE_SETTING_KEY,
303 data: packed,
304 }];
305 let option = if self.layer_type == 0 {
306 Some(PsdLayerOption {
307 additional_info,
308 ..Default::default()
309 })
310 } else {
311 Some(PsdLayerOption {
312 visible: self.visible,
313 opacity: self.opacity,
314 additional_info,
315 })
316 };
317 psd.add_layer_group(self.name, self.layer_type == 2, option)?;
318 }
319 Ok(())
320 }
321}
322
323#[derive(Debug)]
324struct PImgLayerRoot<'a> {
325 layers: Vec<PImgLayer<'a>>,
326}
327
328impl<'a> PImgLayerRoot<'a> {
329 pub fn new(layers: &'a PsbValueFixed) -> Result<Self> {
330 let mut root_layers = Vec::new();
331 for layer in layers.members() {
332 if layer["group_layer_id"].is_none() && layer["diff_id"].is_none() {
333 root_layers.push(PImgLayer::new(layer, layers)?);
334 }
335 }
336 Ok(Self {
337 layers: root_layers,
338 })
339 }
340
341 fn len(&self) -> usize {
342 self.layers.iter().map(|l| l.len()).sum()
343 }
344
345 fn save_to_psd(&self, img: &PImg, psd: &mut PsdWriter, base: &mut ImageData) -> Result<()> {
346 for layer in &self.layers {
347 layer.save_to_psd(img, psd, base)?;
348 }
349 Ok(())
350 }
351}
352
353#[derive(Debug)]
354pub struct PImg {
356 psb: VirtualPsbFixed,
357 overlay: Option<bool>,
358 psd: bool,
359 psd_compress: bool,
360 zlib_compression_level: u32,
361}
362
363impl PImg {
364 pub fn new<R: Read + Seek>(reader: R, _filename: &str, config: &ExtraConfig) -> Result<Self> {
370 let psb = PsbReader::open_psb_v2(reader)?.to_psb_fixed();
371 Ok(Self {
372 psb,
373 overlay: config.emote_pimg_overlay,
374 psd: config.emote_pimg_psd,
375 psd_compress: config.psd_compress,
376 zlib_compression_level: config.zlib_compression_level,
377 })
378 }
379
380 fn load_img(&self, layer_id: i64) -> Result<Tlg> {
381 let layer_id = layer_id as usize;
382 let psb = self.psb.root();
383 let reference = &psb[format!("{layer_id}.tlg")];
384 let resource_id = reference
385 .resource_id()
386 .ok_or_else(|| anyhow::anyhow!("Layer {layer_id} does not have a resource ID"))?
387 as usize;
388 if resource_id >= self.psb.resources().len() {
389 return Err(anyhow::anyhow!(
390 "Resource ID {resource_id} for layer {layer_id} is out of bounds"
391 ));
392 }
393 let resource = &self.psb.resources()[resource_id];
394 Ok(load_tlg(MemReaderRef::new(&resource))?)
395 }
396}
397
398impl Script for PImg {
399 fn default_output_script_type(&self) -> OutputScriptType {
400 OutputScriptType::Custom
401 }
402
403 fn is_output_supported(&self, output: OutputScriptType) -> bool {
404 matches!(output, OutputScriptType::Custom)
405 }
406
407 fn custom_output_extension<'a>(&'a self) -> &'a str {
408 "psd"
409 }
410
411 fn default_format_type(&self) -> FormatOptions {
412 FormatOptions::None
413 }
414
415 fn is_image(&self) -> bool {
416 !self.psd
417 }
418
419 fn is_multi_image(&self) -> bool {
420 !self.psd
421 }
422
423 fn export_multi_image<'a>(
424 &'a self,
425 ) -> Result<Box<dyn Iterator<Item = Result<ImageDataWithName>> + 'a>> {
426 let psb = self.psb.root();
427 let overlay = self.overlay.unwrap_or_else(|| {
428 psb["layers"]
429 .members()
430 .all(|layer| layer["group_layer_id"].is_none())
431 });
432 if !overlay {
433 return Ok(Box::new(PImgIter2 {
434 pimg: self,
435 layers: psb.iter(),
436 }));
437 }
438 let width = psb["width"]
439 .as_u32()
440 .ok_or(anyhow::anyhow!("missing width"))?;
441 let height = psb["height"]
442 .as_u32()
443 .ok_or(anyhow::anyhow!("missing height"))?;
444 if !psb["layers"].is_list() {
445 return Err(anyhow::anyhow!("layers is not a list"));
446 }
447 if psb["layers"].len() == 0 {
448 return Ok(Box::new(std::iter::empty()));
449 }
450 let mut bases = HashMap::new();
451 for i in psb["layers"].members() {
452 if !i["diff_id"].is_none() {
453 continue; }
455 let layer_id = i["layer_id"]
456 .as_i64()
457 .ok_or(anyhow::anyhow!("missing layer_id"))?;
458 let top = i["top"].as_u32().ok_or(anyhow::anyhow!("missing top"))?;
459 let left = i["left"].as_u32().ok_or(anyhow::anyhow!("missing left"))?;
460 let opacity = i["opacity"]
461 .as_u8()
462 .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid opacity"))?;
463 bases.insert(layer_id, (self.load_img(layer_id)?, top, left, opacity));
464 }
465 Ok(Box::new(PImgIter {
466 pimg: self,
467 width,
468 height,
469 layers: psb["layers"].members(),
470 bases,
471 }))
472 }
473
474 fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
475 let psb = self.psb.root();
476 let width = psb["width"]
477 .as_u32()
478 .ok_or(anyhow::anyhow!("missing width"))?;
479 let height = psb["height"]
480 .as_u32()
481 .ok_or(anyhow::anyhow!("missing height"))?;
482 let mut psd = PsdWriter::new(width, height, ImageColorType::Rgba, 8, encoding)?
483 .compress(self.psd_compress)
484 .zlib_compression_level(self.zlib_compression_level);
485 let mut base = ImageData {
486 width,
487 height,
488 color_type: ImageColorType::Rgba,
489 depth: 8,
490 data: vec![0u8; (width * height * 4) as usize],
491 };
492 let layers = PImgLayerRoot::new(&psb["layers"])?;
493 if layers.len() != psb["layers"].len() {
494 return Err(anyhow::anyhow!("Layer hierarchy is invalid"));
495 }
496 layers.save_to_psd(self, &mut psd, &mut base)?;
497 let file = std::fs::File::create(filename)?;
498 let mut writer = std::io::BufWriter::new(file);
499 psd.save(base, &mut writer)?;
500 Ok(())
501 }
502}
503
504struct PImgIter<'a> {
505 pimg: &'a PImg,
506 width: u32,
507 height: u32,
508 layers: ListIter<'a>,
509 bases: HashMap<i64, (Tlg, u32, u32, u8)>,
510}
511
512impl<'a> Iterator for PImgIter<'a> {
513 type Item = Result<ImageDataWithName>;
514
515 fn next(&mut self) -> Option<Self::Item> {
516 match self.layers.next() {
517 Some(layer) => {
518 let layer_id =
519 try_option!(layer["layer_id"].as_i64().ok_or_else(|| {
520 anyhow::anyhow!("Layer does not have a valid layer_id")
521 }));
522 let layer_name = try_option!(
523 layer["name"]
524 .as_str()
525 .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid name") })
526 );
527 let width = try_option!(
528 layer["width"]
529 .as_u32()
530 .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid width") })
531 );
532 let height = try_option!(
533 layer["height"]
534 .as_u32()
535 .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid height") })
536 );
537 let top = try_option!(
538 layer["top"]
539 .as_u32()
540 .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid top") })
541 );
542 let left = try_option!(
543 layer["left"]
544 .as_u32()
545 .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid left") })
546 );
547 let opacity = try_option!(
548 layer["opacity"]
549 .as_u8()
550 .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid opacity") })
551 );
552 if layer["diff_id"].is_none() {
553 let base = &try_option!(self.bases.get(&layer_id).ok_or(anyhow::anyhow!(
554 "Base image for layer_id {} not found",
555 layer_id
556 )))
557 .0;
558 let mut data = ImageData {
559 width: self.width,
560 height: self.height,
561 color_type: match base.color {
562 TlgColorType::Bgr24 => ImageColorType::Bgr,
563 TlgColorType::Bgra32 => ImageColorType::Bgra,
564 TlgColorType::Grayscale8 => ImageColorType::Grayscale,
565 },
566 depth: 8,
567 data: base.data.clone(),
568 };
569 if opacity != 255 {
570 try_option!(apply_opacity(&mut data, opacity));
571 }
572 if self.width != width || self.height != height || top != 0 || left != 0 {
573 data =
574 try_option!(draw_on_canvas(data, self.width, self.height, left, top));
575 }
576 return Some(Ok(ImageDataWithName {
577 name: layer_name.to_string(),
578 data,
579 }));
580 } else {
581 let diff_id =
582 try_option!(layer["diff_id"].as_i64().ok_or_else(|| {
583 anyhow::anyhow!("Layer does not have a valid diff_id")
584 }));
585 let (base, base_top, base_left, base_opacity) = try_option!(
586 self.bases
587 .get(&diff_id)
588 .ok_or(anyhow::anyhow!("Base image layer {} not found", diff_id))
589 );
590 let diff = try_option!(self.pimg.load_img(layer_id));
591 if base.color != diff.color {
592 return Some(Err(anyhow::anyhow!(
593 "Color type mismatch for layer_id {}: base color {:?}, diff color {:?}",
594 layer_id,
595 base.color,
596 diff.color
597 )));
598 }
599 let mut base_img = ImageData {
600 width: base.width,
601 height: base.height,
602 color_type: match base.color {
603 TlgColorType::Bgr24 => ImageColorType::Bgr,
604 TlgColorType::Bgra32 => ImageColorType::Bgra,
605 TlgColorType::Grayscale8 => ImageColorType::Grayscale,
606 },
607 depth: 8,
608 data: base.data.clone(),
609 };
610 if base.width != self.width
611 || base.height != self.height
612 || *base_top != 0
613 || *base_left != 0
614 {
615 base_img = try_option!(draw_on_canvas(
616 base_img,
617 self.width,
618 self.height,
619 *base_left,
620 *base_top
621 ));
622 }
623 if *base_opacity != 255 {
624 try_option!(apply_opacity(&mut base_img, *base_opacity));
625 }
626 let diff = ImageData {
627 width: diff.width,
628 height: diff.height,
629 color_type: match diff.color {
630 TlgColorType::Bgr24 => ImageColorType::Bgr,
631 TlgColorType::Bgra32 => ImageColorType::Bgra,
632 TlgColorType::Grayscale8 => ImageColorType::Grayscale,
633 },
634 depth: 8,
635 data: diff.data.clone(),
636 };
637 try_option!(draw_on_img_with_opacity(
638 &mut base_img,
639 &diff,
640 left,
641 top,
642 opacity
643 ));
644 Some(Ok(ImageDataWithName {
645 name: layer_name.to_string(),
646 data: base_img,
647 }))
648 }
649 }
650 None => None,
651 }
652 }
653}
654
655struct PImgIter2<'a> {
656 pimg: &'a PImg,
657 layers: ObjectIter<'a>,
658}
659
660impl<'a> Iterator for PImgIter2<'a> {
661 type Item = Result<ImageDataWithName>;
662
663 fn next(&mut self) -> Option<Self::Item> {
664 match self.layers.next() {
665 Some((k, v)) => {
666 if !k.ends_with(".tlg") {
667 return self.next();
668 }
669 let resource_id = try_option!(
670 v.resource_id()
671 .ok_or_else(|| anyhow::anyhow!("Layer {} does not have a resource ID", k))
672 ) as usize;
673 let name = k.trim_end_matches(".tlg").to_string();
674 if resource_id >= self.pimg.psb.resources().len() {
675 return Some(Err(anyhow::anyhow!(
676 "Resource ID {} for layer {} is out of bounds",
677 resource_id,
678 k
679 )));
680 }
681 let resource = &self.pimg.psb.resources()[resource_id];
682 let tlg = try_option!(load_tlg(MemReaderRef::new(&resource)));
683 Some(Ok(ImageDataWithName {
684 name,
685 data: ImageData {
686 width: tlg.width,
687 height: tlg.height,
688 color_type: match tlg.color {
689 TlgColorType::Bgr24 => ImageColorType::Bgr,
690 TlgColorType::Bgra32 => ImageColorType::Bgra,
691 TlgColorType::Grayscale8 => ImageColorType::Grayscale,
692 },
693 depth: 8,
694 data: tlg.data.clone(),
695 },
696 }))
697 }
698 None => None,
699 }
700 }
701}