1use agent_diva_files::handle::FileMetadata;
37use agent_diva_files::FileHandle;
38use chrono::{DateTime, Utc};
39use serde::{Deserialize, Serialize};
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct FileAttachment {
48 pub file_id: String,
50
51 pub filename: String,
53
54 pub size: u64,
56
57 #[serde(skip_serializing_if = "Option::is_none")]
59 pub mime_type: Option<String>,
60
61 pub channel: String,
63
64 #[serde(skip_serializing_if = "Option::is_none")]
66 pub message_id: Option<String>,
67
68 #[serde(skip_serializing_if = "Option::is_none")]
70 pub uploaded_by: Option<String>,
71
72 pub stored_at: DateTime<Utc>,
74
75 pub ref_count: usize,
77}
78
79impl FileAttachment {
80 pub fn from_handle(handle: FileHandle, channel: &str, message_id: Option<&str>) -> Self {
109 let ref_count = handle.ref_count();
111 let file_id = handle.id;
112 let metadata = &handle.metadata;
113
114 Self {
115 file_id,
116 filename: metadata.name.clone(),
117 size: metadata.size,
118 mime_type: metadata.mime_type.clone(),
119 channel: channel.to_string(),
120 message_id: message_id.map(String::from),
121 uploaded_by: metadata.source.clone(),
122 stored_at: metadata.created_at,
123 ref_count,
124 }
125 }
126
127 pub fn from_metadata(
132 file_id: &str,
133 metadata: &FileMetadata,
134 channel: &str,
135 message_id: Option<&str>,
136 ref_count: usize,
137 ) -> Self {
138 Self {
139 file_id: file_id.to_string(),
140 filename: metadata.name.clone(),
141 size: metadata.size,
142 mime_type: metadata.mime_type.clone(),
143 channel: channel.to_string(),
144 message_id: message_id.map(String::from),
145 uploaded_by: metadata.source.clone(),
146 stored_at: metadata.created_at,
147 ref_count,
148 }
149 }
150
151 pub fn display(&self) -> String {
177 let size_str = Self::format_size(self.size);
178 format!("{} ({}) from {}", self.filename, size_str, self.channel)
179 }
180
181 fn format_size(bytes: u64) -> String {
183 const KB: u64 = 1024;
184 const MB: u64 = KB * 1024;
185 const GB: u64 = MB * 1024;
186
187 if bytes >= GB {
188 format!("{:.1} GB", bytes as f64 / GB as f64)
189 } else if bytes >= MB {
190 format!("{:.1} MB", bytes as f64 / MB as f64)
191 } else if bytes >= KB {
192 format!("{:.1} KB", bytes as f64 / KB as f64)
193 } else {
194 format!("{} B", bytes)
195 }
196 }
197
198 pub fn is_image(&self) -> bool {
200 self.mime_type
201 .as_ref()
202 .map(|m| m.starts_with("image/"))
203 .unwrap_or(false)
204 }
205
206 pub fn is_video(&self) -> bool {
208 self.mime_type
209 .as_ref()
210 .map(|m| m.starts_with("video/"))
211 .unwrap_or(false)
212 }
213
214 pub fn is_audio(&self) -> bool {
216 self.mime_type
217 .as_ref()
218 .map(|m| m.starts_with("audio/"))
219 .unwrap_or(false)
220 }
221
222 pub fn is_document(&self) -> bool {
224 self.mime_type
225 .as_ref()
226 .map(|m| {
227 m.starts_with("application/pdf")
228 || m.starts_with("application/")
229 || m.starts_with("text/")
230 })
231 .unwrap_or(false)
232 }
233}
234
235impl std::fmt::Display for FileAttachment {
236 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
237 write!(f, "{}", self.display())
238 }
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244 use std::path::PathBuf;
245
246 fn create_test_metadata() -> FileMetadata {
247 FileMetadata {
248 name: "test_document.pdf".to_string(),
249 size: 1024 * 1024, mime_type: Some("application/pdf".to_string()),
251 source: Some("telegram".to_string()),
252 created_at: Utc::now(),
253 last_accessed_at: None,
254 preview: None,
255 }
256 }
257
258 fn create_test_handle() -> FileHandle {
259 let metadata = create_test_metadata();
260 FileHandle::new(
261 "sha256:abc123def456".to_string(),
262 PathBuf::from("ab/c123def456"),
263 metadata,
264 )
265 }
266
267 #[test]
268 fn test_from_handle() {
269 let handle = create_test_handle();
270 let attachment = FileAttachment::from_handle(handle, "telegram", Some("msg_789"));
271
272 assert_eq!(attachment.file_id, "sha256:abc123def456");
273 assert_eq!(attachment.filename, "test_document.pdf");
274 assert_eq!(attachment.size, 1024 * 1024);
275 assert_eq!(attachment.mime_type, Some("application/pdf".to_string()));
276 assert_eq!(attachment.channel, "telegram");
277 assert_eq!(attachment.message_id, Some("msg_789".to_string()));
278 assert_eq!(attachment.uploaded_by, Some("telegram".to_string()));
279 }
280
281 #[test]
282 fn test_display() {
283 let handle = create_test_handle();
284 let attachment = FileAttachment::from_handle(handle, "discord", None);
285
286 let display = attachment.display();
287 assert!(display.contains("test_document.pdf"));
288 assert!(display.contains("discord"));
289 assert!(display.contains("MB")); }
291
292 #[test]
293 fn test_is_image() {
294 let mut handle = create_test_handle();
295 handle.metadata.mime_type = Some("image/png".to_string());
296 let attachment = FileAttachment::from_handle(handle, "telegram", None);
297
298 assert!(attachment.is_image());
299 assert!(!attachment.is_video());
300 assert!(!attachment.is_audio());
301 assert!(!attachment.is_document());
302 }
303
304 #[test]
305 fn test_is_video() {
306 let mut handle = create_test_handle();
307 handle.metadata.mime_type = Some("video/mp4".to_string());
308 let attachment = FileAttachment::from_handle(handle, "discord", None);
309
310 assert!(!attachment.is_image());
311 assert!(attachment.is_video());
312 }
313
314 #[test]
315 fn test_format_size() {
316 assert_eq!(FileAttachment::format_size(500), "500 B");
317 assert_eq!(FileAttachment::format_size(1024), "1.0 KB");
318 assert_eq!(FileAttachment::format_size(1024 * 512), "512.0 KB");
319 assert_eq!(FileAttachment::format_size(1024 * 1024), "1.0 MB");
320 assert_eq!(FileAttachment::format_size(1024 * 1024 * 50), "50.0 MB");
321 assert_eq!(FileAttachment::format_size(1024 * 1024 * 1024), "1.0 GB");
322 }
323
324 #[test]
325 fn test_serialize() {
326 let handle = create_test_handle();
327 let attachment = FileAttachment::from_handle(handle, "slack", Some("ts_123"));
328
329 let json = serde_json::to_string(&attachment).unwrap();
330 assert!(json.contains("test_document.pdf"));
331 assert!(json.contains("slack"));
332 assert!(json.contains("sha256:abc123def456"));
333 }
334
335 #[test]
336 fn test_deserialize() {
337 let json = r#"{
338 "file_id": "sha256:test123",
339 "filename": "doc.pdf",
340 "size": 2048,
341 "mime_type": "application/pdf",
342 "channel": "telegram",
343 "message_id": "msg_456",
344 "uploaded_by": "user123",
345 "stored_at": "2024-01-15T10:30:00Z",
346 "ref_count": 3
347 }"#;
348
349 let attachment: FileAttachment = serde_json::from_str(json).unwrap();
350 assert_eq!(attachment.file_id, "sha256:test123");
351 assert_eq!(attachment.filename, "doc.pdf");
352 assert_eq!(attachment.size, 2048);
353 assert_eq!(attachment.channel, "telegram");
354 }
355}