Add functionality to read prompt metadata from PNG files and include it in map listings

src-tauri/src/main.rs

@@ -2,8 +2,8 @@
 
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-use std::fs;
+use std::fs::{self, File};
-use std::io::{BufRead, BufReader, Read};
+use std::io::{BufRead, BufReader, Read, Seek, SeekFrom};
 use std::path::{Path, PathBuf};
 use std::process::{Command, Stdio};
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -22,6 +22,7 @@ struct MapInfo {
     path: String,
     filename: String,
     modified: i64,
+    prompt: Option<String>,
 }
 
 #[derive(Serialize)]
@@ -78,6 +79,89 @@ fn python_binary(root: &Path) -> PathBuf {
     PathBuf::from("python3")
 }
 
+fn decode_latin1(bytes: &[u8]) -> String {
+    bytes.iter().map(|&byte| byte as char).collect()
+}
+
+fn prompt_from_png_text_chunk(chunk_type: &[u8; 4], data: &[u8]) -> Option<String> {
+    match chunk_type {
+        b"tEXt" => {
+            let separator = data.iter().position(|&byte| byte == 0)?;
+            let key = decode_latin1(&data[..separator]);
+            if !key.eq_ignore_ascii_case("prompt") {
+                return None;
+            }
+            Some(decode_latin1(&data[separator + 1..]))
+        }
+        b"iTXt" => {
+            let keyword_end = data.iter().position(|&byte| byte == 0)?;
+            let key = decode_latin1(&data[..keyword_end]);
+            if !key.eq_ignore_ascii_case("prompt") {
+                return None;
+            }
+
+            let compression_flag_index = keyword_end + 1;
+            let compression_method_index = compression_flag_index + 1;
+            let after_compression = compression_method_index + 1;
+            if data.len() < after_compression || data[compression_flag_index] != 0 {
+                return None;
+            }
+
+            let rest = &data[after_compression..];
+            let language_end = rest.iter().position(|&byte| byte == 0)?;
+            let translated_keyword_start = language_end + 1;
+            let translated_keyword_end = rest[translated_keyword_start..]
+                .iter()
+                .position(|&byte| byte == 0)?
+                + translated_keyword_start;
+            let text_start = translated_keyword_end + 1;
+            let text = &rest[text_start..];
+            match String::from_utf8(text.to_vec()) {
+                Ok(value) => Some(value),
+                Err(_) => Some(String::from_utf8_lossy(text).to_string()),
+            }
+        }
+        _ => None,
+    }
+}
+
+fn read_png_prompt(path: &Path) -> Option<String> {
+    let mut file = File::open(path).ok()?;
+    let mut signature = [0_u8; 8];
+    file.read_exact(&mut signature).ok()?;
+    if signature != [137, 80, 78, 71, 13, 10, 26, 10] {
+        return None;
+    }
+
+    loop {
+        let mut length_buf = [0_u8; 4];
+        if file.read_exact(&mut length_buf).is_err() {
+            break;
+        }
+
+        let length = u32::from_be_bytes(length_buf) as usize;
+        let mut chunk_type = [0_u8; 4];
+        file.read_exact(&mut chunk_type).ok()?;
+
+        if &chunk_type == b"tEXt" || &chunk_type == b"iTXt" {
+            let mut data = vec![0_u8; length];
+            file.read_exact(&mut data).ok()?;
+            if let Some(prompt) = prompt_from_png_text_chunk(&chunk_type, &data) {
+                return Some(prompt);
+            }
+            file.seek(SeekFrom::Current(4)).ok()?;
+        } else {
+            file.seek(SeekFrom::Current(length as i64 + 4)).ok()?;
+        }
+
+        if &chunk_type == b"IEND" {
+            break;
+        }
+    }
+
+    None
+}
+
 #[tauri::command]
 fn list_maps(state: State<Paths>) -> Result<Vec<MapInfo>, String> {
     let dir = output_dir(&state.project_root)?;
@@ -116,6 +200,7 @@ fn list_maps(state: State<Paths>) -> Result<Vec<MapInfo>, String> {
             path: path.to_string_lossy().to_string(),
             filename,
             modified,
+            prompt: read_png_prompt(&path),
         });
     }
     maps.sort_by(|a, b| b.modified.cmp(&a.modified));