git-heatmap/src/main.rs

#![feature(byte_slice_trim_ascii)]
#![feature(let_chains)]
#![allow(dead_code)]

use std::{
	cmp::Reverse,
	collections::HashSet,
	path::{self, PathBuf},
	sync::OnceLock,
};

use anyhow::{anyhow, Context, Result};
use chrono::{DateTime, Duration, Local, NaiveDate, TimeZone};
use clap::Parser;
use gix::{bstr::ByteSlice, traverse::commit::simple::Sorting, ObjectId};
use heatmap::{ColorLogic, HeatmapColors};
use itertools::Itertools;
use mailmap::Mailmap;
use rgb::Rgb;

use crate::{cli::CliArgs, heatmap::Heatmap};

mod cli;
mod heatmap;
mod mailmap;
mod rgb;

pub const ESCAPE: &str = "\x1B";
pub const RESET: &str = "\x1B[0m";
pub const DAYS: [&str; 7] = ["Mo", "Tu", "We", "Th", "Fr", "Sa", "Su"];
pub static CHAR: OnceLock<char> = OnceLock::new();
static COLOR_LOGIC: OnceLock<ColorLogic> = OnceLock::new();
pub static COLOR_MAP: OnceLock<Vec<String>> = OnceLock::new();

const GREEN_COLOR_MAP: [Rgb; 5] = [
	Rgb(0, 0, 0),
	Rgb(14, 68, 41),
	Rgb(0, 109, 50),
	Rgb(38, 166, 65),
	Rgb(25, 255, 64),
];

const RED_COLOR_MAP: [Rgb; 5] = [
	Rgb(0, 0, 0),
	Rgb(208, 169, 35),
	Rgb(208, 128, 35),
	Rgb(208, 78, 35),
	Rgb(255, 0, 0),
];

#[derive(Debug, PartialEq, Eq, Hash)]
struct Commit {
	id: ObjectId,
	title: String,
	author: Author,
	time: DateTime<Local>,
}

#[derive(Debug, PartialEq, Eq, Hash, Clone)]
struct Author {
	name: String,
	email: String,
}

fn main() -> Result<()> {
	// clear_screen();

	let args = CliArgs::parse();

	// dbg!(&args);

	CHAR.set(args.char).unwrap();
	COLOR_LOGIC.set(args.counting.clone()).unwrap();
	let color_map = match args.color_scheme {
		HeatmapColors::Green => GREEN_COLOR_MAP,
		HeatmapColors::Red => RED_COLOR_MAP,
	};
	let color_map = color_map
		.into_iter()
		.map(|c| c.to_ansi())
		.collect::<Vec<_>>();
	COLOR_MAP.set(color_map).unwrap();

	let since = NaiveDate::parse_from_str(&args.since, "%Y-%m-%d").unwrap();

	let until = args
		.until
		.clone()
		.unwrap_or_else(|| get_default_until(since));
	let until = NaiveDate::parse_from_str(&until, "%Y-%m-%d").unwrap();

	let split_months = args.split_months;

	let commits = get_commits(args, since, until).with_context(|| "Could not fetch commit list")?;

	let heatmap = Heatmap::new(since, until, commits.0, commits.1, split_months);

	println!("{heatmap}");

	Ok(())
}

fn get_default_until(since: NaiveDate) -> String {
	let mut until = Local::now().date_naive();
	if since + Duration::days(365) < until {
		until = since + Duration::days(365);
	}
	until.format("%Y-%m-%d").to_string()
}

fn get_color(val: i32, high: i32) -> usize {
	match COLOR_LOGIC.get() {
		Some(logic) => match logic {
			ColorLogic::ByAmount => match val {
				0 => 0,
				x if x < 2 => 1,
				x if x < 4 => 2,
				x if x < 6 => 3,
				x if x >= 6 => 4,
				_ => 0,
			},
			ColorLogic::ByWeight => {
				let color = val as f32 / high as f32;
				match color {
					0.0 => 0,
					x if x <= 0.2 => 1,
					x if x <= 0.4 => 2,
					x if x <= 0.8 => 3,
					x if x > 0.8 => 4,
					_ => 0,
				}
			}
		},
		None => 0,
	}
}

fn clear_screen() {
	print!("\x1b[2J\x1b[1;1H");
}

fn get_char() -> char {
	*CHAR.get_or_init(|| '▩')
}

fn get_color_map() -> Vec<String> {
	COLOR_MAP
		.get_or_init(|| {
			GREEN_COLOR_MAP
				.into_iter()
				.map(|c| c.to_ansi())
				.collect::<Vec<_>>()
		})
		.to_vec()
}

fn find_git_repos(
	scan_path: &path::Path,
	repos: &mut Vec<PathBuf>,
	ignored_repos: &Vec<String>,
	_args: &CliArgs,
) {
	let Ok(dirs) = scan_path.read_dir()
	else {
		return;
	};

	let dirs: Vec<_> = dirs
		.filter_map(|d| d.ok())
		.filter(|d| {
			let dir_name = d.file_name().to_string_lossy().to_string();
			!ignored_repos.contains(&dir_name)
		})
		.filter(|d| d.file_type().is_ok_and(|t| t.is_dir()))
		.collect_vec();

	let dirs = dirs.iter().map(|d| d.path());

	for dir in dirs {
		let filename = dir.file_name().unwrap_or_default().to_string_lossy();
		match filename.as_ref() {
			".git" => repos.push(dir),
			_ => find_git_repos(&dir, repos, ignored_repos, _args),
		}
	}
}

fn get_commits(
	args: CliArgs,
	start_date: NaiveDate,
	end_date: NaiveDate,
) -> Result<(usize, Vec<Commit>)> {
	let mut commits: HashSet<Commit> = HashSet::new();

	let ignored_repos = args.ignored_repos.as_ref().unwrap_or(&vec![]).to_owned();

	let (repos, branches) = match &args.root_dir {
		Some(root) => {
			let mut repos: Vec<PathBuf> = vec![];
			find_git_repos(root, &mut repos, &ignored_repos, &args);
			let branches = vec!["".to_string(); repos.len()];
			(repos, branches)
		}
		None => {
			let repos = match args.repos {
				Some(r) => r,
				None => vec![PathBuf::from(".")],
			};

			let branches = args.branches.unwrap_or_else(|| vec!["".to_string()]);

			if repos.len() > 1 && repos.len() != branches.len() {
				return Err(anyhow!(
					"Number of repos ({}) needs to match the number of branch lists ({})!",
					repos.len(),
					branches.len()
				));
			}
			(repos, branches)
		}
	};

	let current_time = Local::now().time();
	let start_date = start_date.and_time(current_time);
	let start_date = Local.from_local_datetime(&start_date).unwrap();

	let end_date = end_date.and_time(current_time);
	let end_date = Local.from_local_datetime(&end_date).unwrap();

	let authors = args.authors.unwrap_or_default();
	let mut repos_count = 0;

	for (i, repo_path) in repos.iter().enumerate() {
		let repo = gix::open(repo_path).unwrap();

		let branch_names = &*branches[i];
		let mut branches = vec![];
		if branch_names.is_empty() {
			branches.extend(repo.branch_names().iter());
		}
		else {
			let branch_names = branch_names.split(' ');
			branches.extend(branch_names);
		}

		let mailmap = Mailmap::new(repo_path);
		let mut has_commits = false;

		for branch in branches {
			let branch_commits = repo
				.rev_parse(branch)?
				.single()
				.unwrap()
				.ancestors()
				.sorting(Sorting::ByCommitTimeNewestFirstCutoffOlderThan {
					seconds: start_date.timestamp(),
				})
				.all()?;

			branch_commits
				.filter_map(|c| c.ok())
				.filter_map(|c| c.object().ok())
				.filter_map(|c| {
					let title = c
						.message()
						.ok()?
						.title
						.trim_ascii()
						.to_str()
						.ok()?
						.to_string();

					if args.no_merges {
						let is_merge = c.parent_ids().count() > 1;
						if is_merge {
							return None;
						}
					}

					let author = c.author().ok()?;

					let email = author.email.to_string();
					let name = author.name.to_string();

					let author = Author { name, email };
					let author = mailmap.resolve(author);

					if !authors.is_empty() && !authors.contains(&author.name) {
						return None;
					}

					let time = c.time().ok()?;
					let time =
						DateTime::from_timestamp_millis(time.seconds * 1000)?.with_timezone(&Local);
					if time <= start_date + Duration::days(1) || time > end_date {
						return None;
					}

					has_commits = true;

					Some(Commit {
						id: c.id,
						title,
						author,
						time,
					})
				})
				.for_each(|c| {
					commits.insert(c);
				});
		}

		if has_commits {
			repos_count += 1;
		}
	}

	let commits = commits
		.into_iter()
		.sorted_by_cached_key(|a| Reverse(a.time))
		.collect_vec();

	Ok((repos_count, commits))
}