AnyBot/src/attach.rs

// Files and pictures the person gives a bot with a message (the original: 附件 — drag files in,
// paste a picture, or pick files; pictures go to the model, documents are read for it).
// Each one is copied into AnyBot's data folder, so the chat keeps working after the original moves.

use base64::Engine;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

/// Largest file accepted (estimated).
pub const MAX_FILE_BYTES: u64 = 52_428_800;
/// Text read from one file for the model (estimated; the original reads 20,000 characters of mail
/// and notes).
pub const MAX_TEXT_CHARS: usize = 20_000;

#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
pub struct Attachment {
    pub name: String,
    /// The copy in AnyBot's data folder.
    pub path: String,
    #[serde(default)]
    pub image: bool,
    #[serde(default)]
    pub bytes: u64,
    /// Where it was also put on the bot's own computer (empty when not).
    #[serde(default)]
    pub on_box: String,
}

fn dir_for(bot_id: u64) -> PathBuf {
    crate::model::State::path().parent().map(|p| p.to_path_buf()).unwrap_or_else(std::env::temp_dir).join("attachments").join(format!("bot-{bot_id}"))
}

pub fn is_image_name(name: &str) -> bool {
    let n = name.to_lowercase();
    [".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".heic", ".tif", ".tiff"].iter().any(|e| n.ends_with(e))
}

/// Copy a file in as an attachment.
pub fn store(bot_id: u64, src: &Path) -> Result<Attachment, String> {
    let meta = std::fs::metadata(src).map_err(|e| format!("讀不到 {}：{e}", src.display()))?;
    if meta.is_dir() {
        return Err(format!("「{}」是資料夾，請壓縮後再附加，或請 Bot 直接處理那個資料夾。", src.display()));
    }
    if meta.len() > MAX_FILE_BYTES {
        return Err(format!("「{}」太大了（最多 50 MB）", src.display()));
    }
    let name = src.file_name().map(|n| n.to_string_lossy().to_string()).unwrap_or_else(|| "附件".into());
    let dir = dir_for(bot_id);
    std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
    let dst = dir.join(format!("{}-{name}", chrono::Local::now().format("%Y%m%d%H%M%S%3f")));
    std::fs::copy(src, &dst).map_err(|e| format!("複製失敗：{e}"))?;
    Ok(Attachment { image: is_image_name(&name), name, path: dst.display().to_string(), bytes: meta.len(), on_box: String::new() })
}

/// A picture or file the bot hands over (a screen it shows, a file fetched from another computer).
pub fn store_bytes(bot_id: u64, name: &str, bytes: &[u8]) -> Result<Attachment, String> {
    if bytes.len() as u64 > MAX_FILE_BYTES {
        return Err(format!("「{name}」太大了（最多 50 MB）"));
    }
    let dir = dir_for(bot_id);
    std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
    let dst = dir.join(format!("{}-{name}", chrono::Local::now().format("%Y%m%d%H%M%S%3f")));
    std::fs::write(&dst, bytes).map_err(|e| format!("存不下來：{e}"))?;
    Ok(Attachment { image: is_image_name(name), name: name.to_string(), path: dst.display().to_string(), bytes: bytes.len() as u64, on_box: String::new() })
}

/// A picture on the clipboard (e.g. a screenshot copied with ⌃⇧⌘4) becomes an attachment.
pub fn paste_image(bot_id: u64) -> Option<Attachment> {
    let dir = dir_for(bot_id);
    std::fs::create_dir_all(&dir).ok()?;
    let name = format!("貼上的圖片 {}.png", chrono::Local::now().format("%H%M%S"));
    let dst = dir.join(format!("{}-{name}", chrono::Local::now().format("%Y%m%d%H%M%S%3f")));
    if !cfg!(target_os = "macos") {
        let png = crate::platform::clipboard_png_linux()?;
        std::fs::write(&dst, &png).ok()?;
        return Some(Attachment { name, path: dst.display().to_string(), image: true, bytes: png.len() as u64, on_box: String::new() });
    }
    let script = r#"ObjC.import('AppKit');
function run(argv) {
  const img = $.NSImage.alloc.initWithPasteboard($.NSPasteboard.generalPasteboard);
  if (img.isNil()) return 'none';
  const rep = $.NSBitmapImageRep.imageRepWithData(img.TIFFRepresentation);
  const png = rep.representationUsingTypeProperties($.NSBitmapImageFileTypePNG, $({}));
  return png.writeToFileAtomically(argv[0], true) ? 'ok' : 'fail';
}"#;
    let out = std::process::Command::new("osascript").args(["-l", "JavaScript", "-e", script, "--"]).arg(&dst).output().ok()?;
    if String::from_utf8_lossy(&out.stdout).trim() != "ok" {
        return None;
    }
    let bytes = std::fs::metadata(&dst).map(|m| m.len()).unwrap_or(0);
    Some(Attachment { name, path: dst.display().to_string(), image: true, bytes, on_box: String::new() })
}

/// The Finder file picker (several files at once).
pub fn pick_files() -> Vec<PathBuf> {
    if !cfg!(target_os = "macos") {
        return crate::platform::pick_files_linux(&crate::i18n::t("選擇要附加的檔案"));
    }
    let script = r#"function run() {
  const app = Application.currentApplication(); app.includeStandardAdditions = true;
  try { return app.chooseFile({withPrompt: '選擇要附加的檔案', multipleSelectionsAllowed: true}).map(f => f.toString()).join('\n'); } catch (e) { return ''; }
}"#;
    std::process::Command::new("osascript")
        .args(["-l", "JavaScript", "-e", script])
        .output()
        .map(|o| String::from_utf8_lossy(&o.stdout).lines().map(str::trim).filter(|l| !l.is_empty()).map(PathBuf::from).collect())
        .unwrap_or_default()
}

pub fn human(b: u64) -> String {
    crate::cloudbox::human(b)
}

/// A picture, sized like a screenshot, as base64 JPEG.
pub fn image_b64(path: &str) -> Option<String> {
    // The viewer's decoder: every common format, SVG, and HEIC/AVIF through sips or ffmpeg.
    let img = image::DynamicImage::ImageRgba8(crate::preview::decode_image(std::path::Path::new(path))?);
    let img = img.to_rgb8();
    let (w, h) = img.dimensions();
    let (fw, fh, resize) = crate::computer::fit(w, h);
    let img = if resize { image::imageops::resize(&img, fw, fh, image::imageops::FilterType::Triangle) } else { img };
    crate::computer::encode_jpeg(&img).map(|j| base64::engine::general_purpose::STANDARD.encode(j))
}

fn clip(s: &str) -> String {
    let n = s.chars().count();
    if n <= MAX_TEXT_CHARS {
        s.to_string()
    } else {
        format!("{}\n…（後面省略 {} 個字）", s.chars().take(MAX_TEXT_CHARS).collect::<String>(), n - MAX_TEXT_CHARS)
    }
}

/// The text inside a document, when it can be read here.
pub fn read_text(path: &str, name: &str) -> Option<String> {
    let n = name.to_lowercase();
    let ext = n.rsplit('.').next().unwrap_or("");
    // Word and web pages: read in AnyBot itself (same on every system).
    if ext == "docx" {
        if let Ok(md) = crate::preview::docx_markdown(std::path::Path::new(path)) {
            return Some(md).filter(|t| !t.trim().is_empty());
        }
    }
    if ext == "html" || ext == "htm" {
        let bytes = std::fs::read(path).ok()?;
        return Some(crate::preview::html_markdown(&String::from_utf8_lossy(&bytes))).filter(|t| !t.trim().is_empty());
    }
    if ["docx", "doc", "rtf", "rtfd", "odt", "webarchive", "pages"].contains(&ext) && cfg!(target_os = "macos") {
        let out = std::process::Command::new("textutil").args(["-convert", "txt", "-stdout", path]).output().ok()?;
        return out.status.success().then(|| String::from_utf8_lossy(&out.stdout).to_string()).filter(|t| !t.trim().is_empty());
    }
    if ext == "pdf" {
        // pdftotext keeps the layout best; without it, read the text in AnyBot (pure Rust).
        if let Some(t) = std::process::Command::new("pdftotext").env("PATH", crate::cloudbox::tool_path()).args(["-layout", path, "-"]).output().ok().filter(|o| o.status.success()).map(|o| String::from_utf8_lossy(&o.stdout).to_string()).filter(|t| !t.trim().is_empty()) {
            return Some(t);
        }
        let bytes = std::fs::read(path).ok()?;
        // NFKC: PDFs often carry CJK radicals ("⽂") for the ordinary characters ("文").
        use unicode_normalization::UnicodeNormalization;
        return std::panic::catch_unwind(|| pdf_extract::extract_text_from_mem(&bytes).ok()).ok().flatten().map(|t| t.nfkc().collect::<String>()).filter(|t| !t.trim().is_empty());
    }
    let bytes = std::fs::read(path).ok()?;
    if bytes.len() > 2 * 1024 * 1024 || bytes.iter().take(8000).any(|b| *b == 0) {
        return None;
    }
    String::from_utf8(bytes).ok()
}

/// What the model gets for a message's attachments: extra text and pictures.
pub fn for_model(atts: &[Attachment]) -> (String, Vec<String>) {
    let mut text = String::new();
    let mut images = vec![];
    for a in atts {
        let where_ = if a.on_box.is_empty() { format!("檔案存在 AnyBot 所在的電腦：{}", a.path) } else { format!("檔案已放進你自己的電腦：{}", a.on_box) };
        if a.image {
            match image_b64(&a.path) {
                Some(b) => {
                    images.push(b);
                    text.push_str(&format!("\n\n[附件圖片「{}」，就是下面這張。{where_}]", a.name));
                }
                None => text.push_str(&format!("\n\n[附件圖片「{}」讀不出來。{where_}]", a.name)),
            }
        } else if let Some(t) = read_text(&a.path, &a.name) {
            text.push_str(&format!("\n\n[附件「{}」（{}）。{where_}]\n```\n{}\n```", a.name, human(a.bytes), clip(&t)));
        } else {
            text.push_str(&format!("\n\n[附件「{}」（{}），這裡讀不出內容，可以用工具處理。{where_}]", a.name, human(a.bytes)));
        }
    }
    (text, images)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn text_and_images_reach_the_model() {
        let dir = std::env::temp_dir().join(format!("anybot-attach-test-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let txt = dir.join("筆記.md");
        std::fs::write(&txt, "# 標題\n內容 hello").unwrap();
        let png = dir.join("圖.png");
        image::RgbImage::from_pixel(2000, 1000, image::Rgb([10, 20, 30])).save(&png).unwrap();
        let bin = dir.join("data.bin");
        std::fs::write(&bin, [0u8, 1, 2, 0, 5]).unwrap();
        let atts: Vec<Attachment> = [&txt, &png, &bin].iter().map(|p| store(99_999, p).unwrap()).collect();
        assert!(atts[1].image && !atts[0].image);
        let (text, images) = for_model(&atts);
        assert!(text.contains("內容 hello") && text.contains("附件圖片「圖.png」") && text.contains("讀不出內容"), "{text}");
        assert_eq!(images.len(), 1);
        for a in &atts {
            let _ = std::fs::remove_file(&a.path);
        }
        let _ = std::fs::remove_dir_all(&dir);
        assert!(store(1, Path::new("/nonexistent/x")).is_err());
    }
}