git-heatmap/src/main.rs

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

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

use anyhow::{Context, Result};
use chrono::{DateTime, Duration, Local, NaiveDate, TimeZone};
use clap::Parser;
use gix::{bstr::ByteSlice, ObjectId, Repository};
use heatmap::HeatmapColors;
use itertools::Itertools;
use rgb::Rgb;

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

mod cli;
mod heatmap;
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();
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(PartialEq, Eq, Hash)]
struct Commit {
	id: ObjectId,
	title: String,
	author: String,
	time: DateTime<Local>,
}

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

	let args = CliArgs::parse();

	CHAR.set(args.char).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 repo = gix::open(&args.input).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 commits = get_commits(repo, args, since).with_context(|| "Could not fetch commit list")?;

	let heatmap = Heatmap::new(since, until, commits);

	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 {
	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,
	}
}

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 get_commits(repo: Repository, args: CliArgs, start_date: NaiveDate) -> Result<Vec<Commit>> {
	let mut commits: HashSet<Commit> = HashSet::new();

	let branches = match args.branches {
		Some(b) => b,
		None => repo
			.branch_names()
			.into_iter()
			.map(|b| b.to_string())
			.collect_vec(),
	};

	for branch in branches {
		let branch_commits = repo
			.rev_parse(&*branch)?
			.single()
			.unwrap()
			.ancestors()
			.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();

				let author = c.author().ok()?.name.to_string();
				if author != args.author {
					return None;
				}

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

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

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

	Ok(commits)
}