use anyhow::Result;
use clap::ArgMatches;
use image::ImageEncoder;
use image::codecs::png::PngEncoder;
use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use crate::configs::app::AppConfig;
use crate::models::action::ActionModel;
use crate::models::arg::ArgActionModel;
use crate::models::arg::ArgExpect;
use crate::validate::ValidateTrait;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FlowModel {
pub name: String,
pub about: String,
#[serde(default)]
pub check: Option<String>,
#[serde(default)]
pub clipboard: bool,
#[serde(default)]
pub notify: bool,
#[serde(default)]
pub args: Vec<ArgActionModel>,
#[serde(default)]
pub actions: Vec<ActionModel>,
#[serde(skip, default)]
pub input_tags: HashMap<String, String>,
}
impl FlowModel {
pub fn load(path: &PathBuf) -> Result<Self> {
let content = fs::read_to_string(path)?;
let flow: Self = yaml_serde::from_str(&content)
.map_err(|e| anyhow::anyhow!("Failed to parse {}: {}", path.display(), e))?;
flow.validate()?;
Ok(flow.apply_system_tags())
}
pub fn apply_args(mut self, matches: &ArgMatches) -> Self {
for arg in &self.args {
if let Some(value) = matches.get_one::<String>(&arg.name) {
self.input_tags.insert(arg.name.clone(), value.clone());
} else if let Some(default) = &arg.default {
let tag = default.trim_start_matches('{').trim_end_matches('}');
let resolved = self
.input_tags
.get(tag)
.cloned()
.unwrap_or_else(|| default.clone());
self.input_tags.insert(arg.name.clone(), resolved);
}
if arg.expect == ArgExpect::Image {
if let Some(value) = self.input_tags.get(&arg.name) {
let resolved = Self::resolve_image(value);
self.input_tags.insert(arg.name.clone(), resolved);
}
}
}
self
}
fn resolve_image(value: &str) -> String {
if value.starts_with("data:") || value.starts_with("iVBOR") || value.starts_with("/9j/") {
return value.to_string();
}
let resolved = crate::utils::path::resolve(value).unwrap_or_else(|_| PathBuf::from(value));
if resolved.exists() {
std::fs::read(&resolved)
.ok()
.map(|bytes| {
use base64::Engine;
base64::engine::general_purpose::STANDARD.encode(&bytes)
})
.unwrap_or_else(|| value.to_string())
} else {
value.to_string()
}
}
fn apply_system_tags(mut self) -> Self {
self.input_tags.insert(
"system_pwd".into(),
std::env::current_dir()
.map(|p| p.display().to_string())
.unwrap_or_default(),
);
self.input_tags
.insert("system_os".into(), std::env::consts::OS.into());
self.input_tags.insert(
"system_user".into(),
std::env::var("USER").unwrap_or_default(),
);
self.input_tags.insert(
"system_home".into(),
std::env::var("HOME").unwrap_or_default(),
);
self.input_tags.insert(
"system_date".into(),
chrono::Local::now().format("%Y-%m-%d").to_string(),
);
self.input_tags.insert(
"system_time".into(),
chrono::Local::now().format("%H:%M:%S").to_string(),
);
self.input_tags.insert(
"system_clipboard".into(),
self.get_clipboard_text().unwrap_or_default(),
);
if let Ok(image_base64) = self.get_clipboard_image() {
self.input_tags
.insert("system_clipboard_image".into(), image_base64);
}
self.input_tags
.insert("system_pid".into(), std::process::id().to_string());
self.input_tags.insert(
"system_temp".into(),
std::env::temp_dir().display().to_string(),
);
self
}
pub fn get_clipboard_text(&self) -> Result<String> {
let mut clipboard = arboard::Clipboard::new()
.map_err(|e| anyhow::anyhow!("Failed to access clipboard: {}", e))?;
let text = clipboard
.get_text()
.map_err(|_| anyhow::anyhow!("Clipboard is empty or contains non-text data."))?;
let bpe = tiktoken_rs::cl100k_base().unwrap();
let tokens = bpe.encode_with_special_tokens(&text).len();
let config = AppConfig::instance()?;
let max_ctx = config
.cluster
.iter()
.map(|c| c.num_ctx)
.max()
.unwrap_or(4096);
if tokens > max_ctx {
anyhow::bail!(
"Clipboard text is too large ({} tokens). Max context size is {} tokens.",
tokens,
max_ctx
);
}
Ok(text)
}
pub fn get_clipboard_image(&self) -> Result<String> {
let mut clipboard = arboard::Clipboard::new()
.map_err(|e| anyhow::anyhow!("Failed to access clipboard: {}", e))?;
let img = clipboard
.get_image()
.map_err(|_| anyhow::anyhow!("Clipboard does not contain an image."))?;
let mut png_bytes = Vec::new();
let encoder = PngEncoder::new(&mut png_bytes);
encoder
.write_image(
&img.bytes,
img.width as u32,
img.height as u32,
image::ExtendedColorType::Rgba8,
)
.map_err(|e| anyhow::anyhow!("Failed to encode PNG: {}", e))?;
use base64::Engine;
Ok(base64::engine::general_purpose::STANDARD.encode(&png_bytes))
}
pub fn get_images(&self) -> Option<Vec<String>> {
let images: Vec<String> = self
.args
.iter()
.filter(|a| a.expect == ArgExpect::Image)
.filter_map(|a| self.input_tags.get(&a.name))
.cloned()
.collect();
if images.is_empty() {
None
} else {
Some(images)
}
}
}