flq 1.0.0

use clap::Parser;
use flq::extract::Filter;
use flq::{crawl, process, Opts, Page};
use std::fs;
use std::io::{BufRead, Write};
use std::process::exit;

/// fetch the web as lean, token-efficient .flq
#[derive(Parser)]
#[command(name = "flq", version, about)]
struct Args {
	/// URLs to fetch (- for HTML on stdin)
	urls: Vec<String>,

	/// write output to file
	#[arg(short, long)]
	output: Option<String>,

	/// custom header(s): -H "Key: Value"
	#[arg(short = 'H', long = "header")]
	headers: Vec<String>,

	/// user-agent
	#[arg(
		short = 'A',
		long = "user-agent",
		default_value = "Mozilla/5.0 (compatible; flq/0.2)"
	)]
	agent: String,

	/// max time per request in seconds (accepts k/m suffix, e.g. 1k)
	#[arg(short = 'm', long = "max-time", default_value_t = 30, value_parser = human_u64)]
	max_time: u64,

	/// follow redirects
	#[arg(short = 'L', long, default_value_t = true)]
	location: bool,

	/// response headers only (like curl -I)
	#[arg(short = 'I', long)]
	head: bool,

	/// send cookies: -b "k=v; k2=v2"
	#[arg(short = 'b', long)]
	cookie: Option<String>,

	/// silent mode
	#[arg(short, long)]
	silent: bool,

	/// emit raw HTML
	#[arg(long)]
	raw: bool,

	/// filter blocks: headings,code,tables,links,img,text (comma-separated)
	#[arg(short = 'f', long)]
	filter: Option<String>,

	/// grep body lines by pattern (case-insensitive)
	#[arg(short = 'g', long)]
	grep: Option<String>,

	/// CSS selector to scope extraction
	#[arg(long)]
	select: Option<String>,

	/// JSON output
	#[arg(short = 'j', long)]
	json: bool,

	/// omit frontmatter
	#[arg(long)]
	no_meta: bool,

	/// approximate token budget for output (accepts k/m suffix, e.g. 2k)
	#[arg(long, value_parser = human_usize)]
	max_tokens: Option<usize>,

	/// crawl same-host links to depth N
	#[arg(long, default_value_t = 0)]
	depth: u8,

	/// max pages when crawling (accepts k/m suffix)
	#[arg(long, default_value_t = 20, value_parser = human_usize)]
	max_pages: usize,

	/// read URL list from stdin, emit NDJSON
	#[arg(long)]
	pipe: bool,

	/// cache fetched pages locally
	#[arg(long)]
	cache: bool,

	/// cache TTL in seconds (accepts k/m suffix, e.g. 86k)
	#[arg(long, default_value_t = 3600, value_parser = human_u64)]
	cache_ttl: u64,

	/// diff against cached version
	#[arg(long)]
	diff: bool,

	/// run as MCP server (stdio JSON-RPC)
	#[arg(long)]
	mcp: bool,
}

/// Parse a human-friendly count: a plain integer, or a number with a
/// `k`/`m`/`g` suffix (case-insensitive, decimals allowed). e.g. `2k` -> 2000,
/// `1.5m` -> 1_500_000.
fn human_u64(s: &str) -> Result<u64, String> {
	let s = s.trim();
	if s.is_empty() {
		return Err("empty number".into());
	}
	let (num, mult): (&str, u64) = match s.as_bytes()[s.len() - 1] {
		b'k' | b'K' => (&s[..s.len() - 1], 1_000),
		b'm' | b'M' => (&s[..s.len() - 1], 1_000_000),
		b'g' | b'G' => (&s[..s.len() - 1], 1_000_000_000),
		_ => (s, 1),
	};
	let num = num.trim();
	if let Ok(n) = num.parse::<u64>() {
		return Ok(n.saturating_mul(mult));
	}
	let f: f64 = num.parse().map_err(|_| format!("invalid number: {s}"))?;
	if !f.is_finite() || f < 0.0 {
		return Err(format!("invalid number: {s}"));
	}
	Ok((f * mult as f64) as u64)
}

fn human_usize(s: &str) -> Result<usize, String> {
	human_u64(s).map(|n| n as usize)
}

fn main() {
	let args = Args::parse();

	if args.mcp {
		tokio::runtime::Runtime::new()
			.expect("failed to start tokio")
			.block_on(async {
				if let Err(e) = flq::mcp::run().await {
					eprintln!("flq mcp: {e}");
					exit(1);
				}
			});
		return;
	}

	let filters = match &args.filter {
		Some(s) => match Filter::parse(s) {
			Ok(f) => f,
			Err(e) => {
				eprintln!("flq: {e}");
				exit(1);
			}
		},
		None => Vec::new(),
	};

	let no_meta = if args.filter.is_some() && !filters.contains(&Filter::Meta) {
		true
	} else {
		args.no_meta
	};
	let filters: Vec<Filter> = filters.into_iter().filter(|f| *f != Filter::Meta).collect();

	let opts = Opts {
		headers: args.headers,
		agent: args.agent,
		max_time: args.max_time,
		cookie: args.cookie,
		select: args.select,
		filters,
		grep: args.grep,
		no_meta,
		max_tokens: args.max_tokens,
		raw: args.raw,
		cache: args.cache,
		cache_ttl: args.cache_ttl,
		diff: args.diff,
	};

	if args.pipe {
		run_pipe(&opts);
		return;
	}

	if args.urls.is_empty() {
		eprintln!("flq: no URL given");
		exit(1);
	}

	if args.head {
		run_head(&args.urls, &opts, args.silent);
		return;
	}

	let results: Vec<Result<Page, String>> = if args.depth > 0 {
		args.urls
			.iter()
			.flat_map(|u| crawl(u, args.depth, args.max_pages, &opts))
			.collect()
	} else {
		args.urls.iter().map(|u| process(u, &opts)).collect()
	};

	let out = if args.json {
		render_json(&results)
	} else {
		render_flq(&results, &opts, args.silent)
	};

	match args.output {
		Some(path) => {
			if let Err(e) = fs::write(&path, &out) {
				eprintln!("flq: write {path}: {e}");
				exit(1);
			}
		}
		None => print!("{out}"),
	}

	if results.iter().all(|r| r.is_err()) {
		exit(1);
	}
}

fn run_pipe(opts: &Opts) {
	let stdin = std::io::stdin();
	let mut stdout = std::io::stdout();
	for line in stdin.lock().lines() {
		let Ok(url) = line else { break };
		let url = url.trim();
		if url.is_empty() {
			continue;
		}
		let obj = match process(url, opts) {
			Ok(page) => serde_json::to_string(&page).unwrap_or_default(),
			Err(e) => format!(
				r#"{{"url":{},"error":{}}}"#,
				serde_json::to_string(url).unwrap_or_default(),
				serde_json::to_string(&e).unwrap_or_default()
			),
		};
		if writeln!(stdout, "{obj}").is_err() {
			break;
		}
	}
}

fn run_head(urls: &[String], opts: &Opts, silent: bool) {
	for url in urls {
		match flq::fetch::get(url, opts) {
			Ok(res) => {
				println!("status: {}", res.status);
				for (k, v) in &res.headers {
					println!("{k}: {v}");
				}
			}
			Err(e) => {
				if !silent {
					eprintln!("flq: {e}");
				}
			}
		}
	}
}

fn render_json(results: &[Result<Page, String>]) -> String {
	let pages: Vec<&Page> = results.iter().filter_map(|r| r.as_ref().ok()).collect();
	if pages.len() == 1 {
		serde_json::to_string(&pages[0]).unwrap_or_default()
	} else {
		serde_json::to_string(&pages).unwrap_or_default()
	}
}

fn render_flq(results: &[Result<Page, String>], opts: &Opts, silent: bool) -> String {
	let mut out = String::new();
	let multi = results.len() > 1;

	for (i, r) in results.iter().enumerate() {
		match r {
			Ok(page) => {
				if !opts.no_meta {
					head_line(&mut out, "title", &page.title);
					head_line(&mut out, "url", &page.url);
					head_line(&mut out, "lang", &page.lang);
					head_line(&mut out, "desc", &page.desc);
					out.push_str("---\n");
				}
				out.push_str(&page.body);
				if page.truncated {
					out.push_str("\n...(truncated)\n");
				}
				if multi && i + 1 < results.len() {
					out.push_str("\n===\n");
				}
			}
			Err(e) => {
				if !silent {
					eprintln!("flq: {e}");
				}
			}
		}
	}
	out
}

fn head_line(s: &mut String, key: &str, val: &str) {
	if !val.is_empty() {
		s.push_str(key);
		s.push_str(": ");
		s.push_str(val);
		s.push('\n');
	}
}

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

	#[test]
	fn plain_integer() {
		assert_eq!(human_u64("500").unwrap(), 500);
		assert_eq!(human_u64("0").unwrap(), 0);
	}

	#[test]
	fn k_m_g_suffix_case_insensitive() {
		assert_eq!(human_u64("2k").unwrap(), 2_000);
		assert_eq!(human_u64("1K").unwrap(), 1_000);
		assert_eq!(human_u64("1m").unwrap(), 1_000_000);
		assert_eq!(human_u64("3M").unwrap(), 3_000_000);
		assert_eq!(human_u64("1g").unwrap(), 1_000_000_000);
	}

	#[test]
	fn decimals_allowed() {
		assert_eq!(human_u64("1.5k").unwrap(), 1_500);
		assert_eq!(human_u64("2.5m").unwrap(), 2_500_000);
	}

	#[test]
	fn surrounding_whitespace_trimmed() {
		assert_eq!(human_u64("  2k ").unwrap(), 2_000);
	}

	#[test]
	fn invalid_inputs_error() {
		assert!(human_u64("").is_err());
		assert!(human_u64("abc").is_err());
		assert!(human_u64("1x").is_err());
		assert!(human_u64("k").is_err());
		assert!(human_u64("-5").is_err());
	}

	#[test]
	fn usize_wrapper() {
		assert_eq!(human_usize("2k").unwrap(), 2_000usize);
	}
}