msg_tool\scripts\emote/
dref.rs

1//! Emote DPAK-referenced Image File (.dref)
2use crate::ext::io::*;
3use crate::ext::psb::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::encoding::*;
7use crate::utils::img::*;
8use anyhow::Result;
9use emote_psb::PsbReader;
10use std::collections::HashMap;
11use std::io::Read;
12use std::path::{Path, PathBuf};
13use url::Url;
14
15#[derive(Debug)]
16/// Emote DREF Script Builder
17pub struct DrefBuilder {}
18
19impl DrefBuilder {
20    /// Creates a new instance of `DrefBuilder`
21    pub fn new() -> Self {
22        Self {}
23    }
24}
25
26impl ScriptBuilder for DrefBuilder {
27    fn default_encoding(&self) -> Encoding {
28        Encoding::Cp932
29    }
30
31    fn build_script(
32        &self,
33        buf: Vec<u8>,
34        filename: &str,
35        encoding: Encoding,
36        _archive_encoding: Encoding,
37        config: &ExtraConfig,
38        archive: Option<&Box<dyn Script>>,
39    ) -> Result<Box<dyn Script>> {
40        Ok(Box::new(Dref::new(
41            buf, encoding, filename, config, archive,
42        )?))
43    }
44
45    fn extensions(&self) -> &'static [&'static str] {
46        &["dref"]
47    }
48
49    fn script_type(&self) -> &'static ScriptType {
50        &ScriptType::EmoteDref
51    }
52
53    fn is_image(&self) -> bool {
54        true
55    }
56}
57
58struct Dpak {
59    psb: VirtualPsbFixed,
60}
61
62struct OffsetData {
63    left: u32,
64    top: u32,
65}
66
67impl Dpak {
68    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
69        let f = std::fs::File::open(path)?;
70        let mut f = std::io::BufReader::new(f);
71        let mut psb = PsbReader::open_psb(&mut f)
72            .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK: {:?}", e))?;
73        let psb = psb
74            .load()
75            .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK: {:?}", e))?;
76        let psb = psb.to_psb_fixed();
77        Ok(Self { psb })
78    }
79
80    pub fn load_from_data(data: &[u8]) -> Result<Self> {
81        let mut psb = PsbReader::open_psb(MemReaderRef::new(data))
82            .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK data: {:?}", e))?;
83        let psb = psb
84            .load()
85            .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK data: {:?}", e))?;
86        let psb = psb.to_psb_fixed();
87        Ok(Self { psb })
88    }
89
90    pub fn load_image(&self, name: &str) -> Result<(ImageData, Option<OffsetData>)> {
91        let root = self.psb.root();
92        let rid = root[name]
93            .resource_id()
94            .ok_or_else(|| anyhow::anyhow!("Resource ID for image '{}' not found in DPAK", name))?
95            as usize;
96        if rid >= self.psb.resources().len() {
97            return Err(anyhow::anyhow!(
98                "Resource ID {} out of bounds for DPAK with {} resources",
99                rid,
100                self.psb.resources().len()
101            ));
102        }
103        let resource = &self.psb.resources()[rid];
104        Self::load_png(&resource)
105    }
106
107    fn load_png(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
108        let mut img = load_png(MemReaderRef::new(&data))?;
109        match img.color_type {
110            ImageColorType::Rgb => {
111                convert_rgb_to_rgba(&mut img)?;
112            }
113            _ => {}
114        }
115        Ok((
116            img,
117            Self::try_read_offset_from_png(MemReaderRef::new(&data))?,
118        ))
119    }
120
121    fn try_read_offset_from_png(mut data: MemReaderRef) -> Result<Option<OffsetData>> {
122        data.pos = 8; // Skip PNG signature
123        data.pos += 8; // Skip chunk size, type
124        data.pos += 17; // Skip IHDR chunk (length + type + width + height + bit depth + color type + compression method + filter method + interlace method)
125        loop {
126            let chunk_size = data.read_u32_be()?;
127            let mut chunk_type = [0u8; 4];
128            data.read_exact(&mut chunk_type)?;
129            if &chunk_type == b"IDAT" || &chunk_type == b"IEND" {
130                break;
131            }
132            if &chunk_type == b"oFFs" {
133                let x = data.read_u32_be()?;
134                let y = data.read_u32_be()?;
135                if data.read_u8()? == 0 {
136                    return Ok(Some(OffsetData { left: x, top: y }));
137                }
138            }
139            data.pos += chunk_size as usize + 4; // Skip chunk data and CRC
140        }
141        Ok(None)
142    }
143}
144
145#[derive(Default)]
146struct DpakLoader {
147    map: HashMap<String, Dpak>,
148}
149
150impl DpakLoader {
151    pub fn load_image(
152        &mut self,
153        dir: &Path,
154        dpak: &str,
155        filename: &str,
156    ) -> Result<(ImageData, Option<OffsetData>)> {
157        let dpak = match self.map.get(dpak) {
158            Some(d) => d,
159            None => {
160                let path = dir.join(dpak);
161                let ndpak = Dpak::new(&path)?;
162                self.map.insert(dpak.to_string(), ndpak);
163                self.map.get(dpak).unwrap()
164            }
165        };
166        dpak.load_image(filename)
167    }
168
169    pub fn load_archives(&mut self, in_archives: &HashMap<String, Vec<u8>>) -> Result<()> {
170        for (name, data) in in_archives.iter() {
171            if !self.map.contains_key(name) {
172                let dpak = Dpak::load_from_data(data)?;
173                self.map.insert(name.clone(), dpak);
174            }
175        }
176        Ok(())
177    }
178}
179
180/// Emote DREF Script
181pub struct Dref {
182    urls: Vec<Url>,
183    dir: PathBuf,
184    in_archives: HashMap<String, Vec<u8>>,
185}
186
187impl std::fmt::Debug for Dref {
188    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189        f.debug_struct("Dref")
190            .field("urls", &self.urls)
191            .field("dir", &self.dir)
192            .finish()
193    }
194}
195
196impl Dref {
197    /// Create a new dref script
198    ///
199    /// * `buf` - The buffer containing the dref script data
200    /// * `encoding` - The encoding of the script
201    /// * `filename` - The name of the file
202    /// * `config` - Extra configuration options
203    /// * `archive` - Optional archive containing additional resources
204    pub fn new(
205        buf: Vec<u8>,
206        encoding: Encoding,
207        filename: &str,
208        _config: &ExtraConfig,
209        archive: Option<&Box<dyn Script>>,
210    ) -> Result<Self> {
211        let text = decode_with_bom_detect(encoding, &buf, true)?.0;
212        let mut urls = Vec::new();
213        for text in text.lines() {
214            let text = text.trim();
215            if text.is_empty() {
216                continue;
217            }
218            urls.push(Url::parse(text)?);
219        }
220        let path = Path::new(filename);
221        let dir = if let Some(parent) = path.parent() {
222            parent.to_path_buf()
223        } else {
224            PathBuf::from(".")
225        };
226        if urls.is_empty() {
227            return Err(anyhow::anyhow!("No URLs found in DREF file: {}", filename));
228        }
229        for u in urls.iter() {
230            if u.scheme() != "psb" {
231                return Err(anyhow::anyhow!(
232                    "Invalid URL scheme in DREF file: {} (expected 'psb')",
233                    u
234                ));
235            }
236        }
237        let mut in_archives = HashMap::new();
238        if let Some(archive) = archive {
239            if archive.is_archive() {
240                for url in urls.iter() {
241                    let filename = url.domain().ok_or(anyhow::anyhow!(
242                        "Invalid URL in DREF file: {} (missing domain)",
243                        url
244                    ))?;
245                    if let Ok(mut content) = archive.open_file_by_name(filename, true) {
246                        in_archives.insert(filename.to_string(), content.data()?);
247                    }
248                }
249            }
250        }
251        Ok(Self {
252            urls,
253            dir,
254            in_archives,
255        })
256    }
257}
258
259impl Script for Dref {
260    fn default_output_script_type(&self) -> OutputScriptType {
261        OutputScriptType::Json
262    }
263
264    fn default_format_type(&self) -> FormatOptions {
265        FormatOptions::None
266    }
267
268    fn is_image(&self) -> bool {
269        true
270    }
271
272    fn export_image(&self) -> Result<ImageData> {
273        let mut loader = DpakLoader::default();
274        loader.load_archives(&self.in_archives)?;
275        let base_url = &self.urls[0];
276        let dpak = base_url.domain().ok_or(anyhow::anyhow!(
277            "Invalid URL in DREF file: {} (missing domain)",
278            base_url
279        ))?;
280        let (mut base_img, base_offset) =
281            loader.load_image(&self.dir, dpak, base_url.path().trim_start_matches("/"))?;
282        if let Some(o) = base_offset {
283            eprintln!("WARN: Base image offset: left={}, top={}", o.left, o.top);
284            crate::COUNTER.inc_warning();
285        }
286        for url in &self.urls[1..] {
287            let dpak = url.domain().ok_or(anyhow::anyhow!(
288                "Invalid URL in DREF file: {} (missing domain)",
289                url
290            ))?;
291            let (img, img_offset) =
292                loader.load_image(&self.dir, dpak, url.path().trim_start_matches("/"))?;
293            let (top, left) = match img_offset {
294                Some(o) => (o.top, o.left),
295                None => (0, 0),
296            };
297            draw_on_img_with_opacity(&mut base_img, &img, left, top, 0xff)?;
298        }
299        Ok(base_img)
300    }
301}