From 88bf83802d935ab27adb40b7428512a7b7914fe3 Mon Sep 17 00:00:00 2001
From: lifegpc <root@lifegpc.com>
Date: Sat, 19 Jul 2025 22:39:46 +0800
Subject: [PATCH] add option to support extract original layer image for PIMG

---
 src/args.rs                        |  4 ++
 src/main.rs                        |  2 +
 src/scripts/kirikiri/image/pimg.rs | 67 +++++++++++++++++++++++++++++-
 src/types.rs                       |  2 +
 4 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/src/args.rs b/src/args.rs
index a33fe66..de79817 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -153,6 +153,10 @@ pub struct Arg {
     #[arg(long, global = true, action = ArgAction::SetTrue)]
     /// Whether to compress files in BGI archive when packing BGI archive.
     pub bgi_compress_file: bool,
+    #[cfg(feature = "kirikiri-img")]
+    #[arg(long, global = true)]
+    /// Whether to overlay PIMG images. (By default, true if all layers are not group layers.)
+    pub kirikiri_pimg_overlay: Option<bool>,
     #[command(subcommand)]
     /// Command
     pub command: Command,
diff --git a/src/main.rs b/src/main.rs
index d0aebef..b1ecd71 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1413,6 +1413,8 @@ fn main() {
         )),
         #[cfg(feature = "bgi-arc")]
         bgi_compress_file: arg.bgi_compress_file,
+        #[cfg(feature = "kirikiri-img")]
+        kirikiri_pimg_overlay: arg.kirikiri_pimg_overlay,
     };
     match &arg.command {
         args::Command::Export { input, output } => {
diff --git a/src/scripts/kirikiri/image/pimg.rs b/src/scripts/kirikiri/image/pimg.rs
index 5ff3a10..975967a 100644
--- a/src/scripts/kirikiri/image/pimg.rs
+++ b/src/scripts/kirikiri/image/pimg.rs
@@ -93,17 +93,21 @@ impl ScriptBuilder for PImgBuilder {
 #[derive(Debug)]
 pub struct PImg {
     psb: VirtualPsbFixed,
+    overlay: Option<bool>,
 }
 
 impl PImg {
-    pub fn new<R: Read + Seek>(reader: R, filename: &str, _config: &ExtraConfig) -> Result<Self> {
+    pub fn new<R: Read + Seek>(reader: R, filename: &str, config: &ExtraConfig) -> Result<Self> {
         let mut psb = PsbReader::open_psb(reader)
             .map_err(|e| anyhow::anyhow!("Failed to open PSB from {}: {:?}", filename, e))?;
         let psb = psb
             .load()
             .map_err(|e| anyhow::anyhow!("Failed to load PSB from {}: {:?}", filename, e))?
             .to_psb_fixed();
-        Ok(Self { psb })
+        Ok(Self {
+            psb,
+            overlay: config.kirikiri_pimg_overlay,
+        })
     }
 
     fn load_img(&self, layer_id: i64) -> Result<Tlg> {
@@ -145,6 +149,17 @@ impl Script for PImg {
         &'a self,
     ) -> Result<Box<dyn Iterator<Item = Result<ImageDataWithName>> + 'a>> {
         let psb = self.psb.root();
+        let overlay = self.overlay.unwrap_or_else(|| {
+            psb["layers"]
+                .members()
+                .all(|layer| layer["group_layer_id"].is_none())
+        });
+        if !overlay {
+            return Ok(Box::new(PImgIter2 {
+                pimg: self,
+                layers: psb.iter(),
+            }));
+        }
         let width = psb["width"]
             .as_u32()
             .ok_or(anyhow::anyhow!("missing width"))?;
@@ -332,3 +347,51 @@ impl<'a> Iterator for PImgIter<'a> {
         }
     }
 }
+
+struct PImgIter2<'a> {
+    pimg: &'a PImg,
+    layers: ObjectIter<'a>,
+}
+
+impl<'a> Iterator for PImgIter2<'a> {
+    type Item = Result<ImageDataWithName>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.layers.next() {
+            Some((k, v)) => {
+                if !k.ends_with(".tlg") {
+                    return self.next();
+                }
+                let resource_id = try_option!(
+                    v.resource_id()
+                        .ok_or_else(|| anyhow::anyhow!("Layer {} does not have a resource ID", k))
+                ) as usize;
+                let name = k.trim_end_matches(".tlg").to_string();
+                if resource_id >= self.pimg.psb.resources().len() {
+                    return Some(Err(anyhow::anyhow!(
+                        "Resource ID {} for layer {} is out of bounds",
+                        resource_id,
+                        k
+                    )));
+                }
+                let resource = &self.pimg.psb.resources()[resource_id];
+                let tlg = try_option!(load_tlg(MemReaderRef::new(&resource)));
+                Some(Ok(ImageDataWithName {
+                    name,
+                    data: ImageData {
+                        width: tlg.width,
+                        height: tlg.height,
+                        color_type: match tlg.color {
+                            TlgColorType::Bgr24 => ImageColorType::Bgr,
+                            TlgColorType::Bgra32 => ImageColorType::Bgra,
+                            TlgColorType::Grayscale8 => ImageColorType::Grayscale,
+                        },
+                        depth: 8,
+                        data: tlg.data.clone(),
+                    },
+                }))
+            }
+            None => None,
+        }
+    }
+}
diff --git a/src/types.rs b/src/types.rs
index 1b7dc50..f9342f0 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -221,6 +221,8 @@ pub struct ExtraConfig {
     pub kirikiri_message_commands: std::sync::Arc<std::collections::HashSet<String>>,
     #[cfg(feature = "bgi-arc")]
     pub bgi_compress_file: bool,
+    #[cfg(feature = "kirikiri-img")]
+    pub kirikiri_pimg_overlay: Option<bool>,
 }
 
 #[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq, PartialOrd, Ord)]