use anyhow::{anyhow, Result};
use bytesize::ByteSize;
use clap_config_file::ClapConfigFile;
use sha2::{Digest, Sha256};
use std::io::IsTerminal;
use std::{fs, path::Path, str::FromStr, time::UNIX_EPOCH};
use crate::{
defaults::{BINARY_FILE_EXTENSIONS, DEFAULT_IGNORE_PATTERNS, DEFAULT_OUTPUT_TEMPLATE},
priority::PriorityRule,
};
#[derive(Clone, Debug, Default, clap::ValueEnum, serde::Serialize, serde::Deserialize)]
pub enum ConfigFormat {
#[default]
Toml,
Yaml,
Json,
}
#[derive(ClapConfigFile, Clone)]
#[config_file_name = "yek"]
#[config_file_formats = "toml,yaml,json"]
pub struct YekConfig {
#[config_arg(positional)]
pub input_dirs: Vec<String>,
#[config_arg(default_value = "10MB")]
pub max_size: String,
#[config_arg()]
pub tokens: String,
#[config_arg()]
pub json: bool,
#[config_arg()]
pub debug: bool,
#[config_arg()]
pub output_dir: Option<String>,
#[config_arg(default_value = ">>>> FILE_PATH\nFILE_CONTENT")]
pub output_template: String,
#[config_arg(long = "ignore-patterns", multi_value_behavior = "extend")]
pub ignore_patterns: Vec<String>,
#[config_arg(long = "unignore-patterns", multi_value_behavior = "extend")]
pub unignore_patterns: Vec<String>,
#[config_arg(accept_from = "config_only")]
pub priority_rules: Vec<PriorityRule>,
#[config_arg(accept_from = "config_only", default_value = BINARY_FILE_EXTENSIONS)]
pub binary_extensions: Vec<String>,
#[config_arg(accept_from = "config_only")]
pub git_boost_max: Option<i32>,
}
#[derive(Clone, serde::Serialize)]
pub struct FullYekConfig {
#[serde(flatten)]
pub base: YekConfig,
pub stream: bool,
pub token_mode: bool,
pub output_file_full_path: String,
}
impl std::ops::Deref for FullYekConfig {
type Target = YekConfig;
fn deref(&self) -> &Self::Target {
&self.base
}
}
impl std::ops::DerefMut for FullYekConfig {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.base
}
}
impl FullYekConfig {
pub fn extend_config_with_defaults(
input_dirs: Vec<String>,
output_dir: String,
) -> FullYekConfig {
let token_mode = false;
let checksum = get_checksum(&input_dirs);
let extension = if token_mode { "tok" } else { "txt" };
let output_file_full_path = Path::new(&output_dir)
.join(format!("yek-output-{}.{}", checksum, extension))
.to_string_lossy()
.to_string();
FullYekConfig {
base: YekConfig {
input_dirs,
output_dir: Some(output_dir),
max_size: "10MB".to_string(),
tokens: String::new(),
json: false,
debug: false,
output_template: DEFAULT_OUTPUT_TEMPLATE.to_string(),
ignore_patterns: Vec::new(),
unignore_patterns: Vec::new(),
priority_rules: Vec::new(),
binary_extensions: Vec::new(),
git_boost_max: Some(100),
},
stream: false,
token_mode: false,
output_file_full_path,
}
}
}
fn get_checksum(input_dirs: &Vec<String>) -> String {
let mut hasher = Sha256::new();
for dir in input_dirs {
let base_path = Path::new(dir);
if !base_path.exists() {
continue;
}
let entries = match fs::read_dir(base_path) {
Ok(iter) => iter.filter_map(|e| e.ok()).collect::<Vec<_>>(),
Err(_) => continue,
};
let mut sorted = entries;
sorted.sort_by_key(|a| a.path());
for entry in sorted {
let p = entry.path();
if let Ok(meta) = fs::metadata(&p) {
let path_str = p.to_string_lossy();
hasher.update(path_str.as_bytes());
hasher.update(meta.len().to_le_bytes());
if let Ok(mod_time) = meta.modified() {
if let Ok(dur) = mod_time.duration_since(UNIX_EPOCH) {
hasher.update(dur.as_secs().to_le_bytes());
}
}
}
}
}
let result = hasher.finalize();
format!("{:x}", result)
}
impl YekConfig {
pub fn init_config() -> FullYekConfig {
let mut config = YekConfig::parse();
let token_mode = !config.tokens.is_empty();
let stream = !std::io::stdout().is_terminal();
if config.input_dirs.is_empty() {
config.input_dirs.push(".".to_string());
}
config
.binary_extensions
.extend(BINARY_FILE_EXTENSIONS.iter().map(|s| s.to_string()));
let mut ignore_patterns = DEFAULT_IGNORE_PATTERNS
.iter()
.map(|s| s.to_string())
.collect::<Vec<_>>();
ignore_patterns.extend(config.ignore_patterns);
config.ignore_patterns = ignore_patterns;
config
.ignore_patterns
.extend(config.unignore_patterns.iter().map(|s|
format!("!{}", s)));
let output_dir = if let Some(dir) = config.output_dir {
dir
} else {
let temp_dir = std::env::temp_dir();
let output_dir = temp_dir.join("yek-output");
std::fs::create_dir_all(&output_dir).unwrap();
output_dir.to_string_lossy().to_string()
};
let checksum = get_checksum(&config.input_dirs);
let extension = if config.json { "json" } else { "txt" };
let output_file_full_path = Path::new(&output_dir)
.join(format!("yek-output-{}.{}", checksum, extension))
.to_string_lossy()
.to_string();
let final_config = FullYekConfig {
base: YekConfig {
input_dirs: config.input_dirs.clone(),
output_dir: Some(output_dir),
max_size: config.max_size.clone(),
tokens: config.tokens.clone(),
json: config.json,
debug: config.debug,
output_template: config.output_template.clone(),
ignore_patterns: config.ignore_patterns.clone(),
unignore_patterns: config.unignore_patterns.clone(),
priority_rules: config.priority_rules.clone(),
binary_extensions: config.binary_extensions.clone(),
git_boost_max: config.git_boost_max,
},
stream,
token_mode,
output_file_full_path,
};
if let Err(e) = validate_config(&final_config) {
eprintln!("Error: {}", e);
std::process::exit(1);
}
final_config
}
}
pub struct ConfigError {
pub field: String,
pub message: String,
}
pub fn validate_config(config: &FullYekConfig) -> Result<()> {
if !config.base.output_template.contains("FILE_PATH")
|| !config.base.output_template.contains("FILE_CONTENT")
{
return Err(anyhow!(
"output_template: Output template must contain FILE_PATH and FILE_CONTENT"
));
}
if config.base.max_size == "0" {
return Err(anyhow!("max_size: Max size cannot be 0"));
}
if !config.token_mode {
ByteSize::from_str(&config.base.max_size)
.map_err(|e| anyhow!("max_size: Invalid size format: {}", e))?;
} else if config.base.tokens.to_lowercase().ends_with('k') {
let val = config.base.tokens[..config.base.tokens.len() - 1]
.trim()
.parse::<usize>()
.map_err(|e| anyhow!("tokens: Invalid token size: {}", e))?;
if val == 0 {
return Err(anyhow!("tokens: Token size cannot be 0"));
}
} else {
let val = config
.base
.tokens
.parse::<usize>()
.map_err(|e| anyhow!("tokens: Invalid token size: {}", e))?;
if val == 0 {
return Err(anyhow!("tokens: Token size cannot be 0"));
}
}
let output_dir = config
.base
.output_dir
.as_ref()
.ok_or_else(|| anyhow!("output_dir: Output directory must be specified"))?;
let path = Path::new(output_dir);
if path.exists() && !path.is_dir() {
return Err(anyhow!(
"output_dir: Output path '{}' exists but is not a directory",
output_dir
));
}
for pattern in &config.base.ignore_patterns {
glob::Pattern::new(pattern)
.map_err(|e| anyhow!("ignore_patterns: Invalid pattern '{}': {}", pattern, e))?;
}
for rule in &config.base.priority_rules {
if rule.score < 0 || rule.score > 1000 {
return Err(anyhow!(
"priority_rules: Priority score {} must be between 0 and 1000",
rule.score
));
}
glob::Pattern::new(&rule.pattern)
.map_err(|e| anyhow!("priority_rules: Invalid pattern '{}': {}", rule.pattern, e))?;
}
if let Err(e) = std::fs::create_dir_all(path) {
return Err(anyhow!(
"output_dir: Cannot create output directory '{}': {}",
output_dir,
e
));
}
Ok(())
}