git-heatmap/src/heatmap.rs

use std::{
	collections::BTreeMap,
	fmt::{Display, Write},
};

use chrono::{Datelike, Duration, NaiveDate};
use clap::ValueEnum;
use itertools::Itertools;

use crate::{get_char, get_color, get_color_map, Commit, DAYS, RESET};

pub struct Heatmap {
	since: NaiveDate,
	until: NaiveDate,
	commits: Vec<Commit>,
	highest_count: i32,
	branches: usize,
	repos: usize,
	chunks: Vec<Chunk>,

	format: Format,
}

impl Heatmap {
	pub fn new(
		since: NaiveDate,
		until: NaiveDate,
		repos: usize,
		branches: usize,
		commits: Vec<Commit>,
		split_months: bool,
		months_per_row: u16,
		format: Format,
	) -> Self {
		let mut heatmap = Self {
			since,
			until,
			commits,
			highest_count: 0,
			branches,
			repos,
			chunks: vec![],
			format,
		};

		let mut grouped_commits = BTreeMap::new();

		for commit in &heatmap.commits {
			let commit_day = commit.time.date_naive();
			let record = grouped_commits.entry(commit_day).or_insert(0);
			*record += 1;
		}

		let mut current_day = since;
		let mut day_of_week = current_day.weekday().num_days_from_monday() % 7;

		let mut chunk = Chunk::new(day_of_week);
		let mut chunk_idx = 0;

		// Track the very first day of the heatmap, as we don't want the extra spacing in front of
		// those.
		let mut first_day = true;

		while current_day <= until {
			if split_months {
				// If current day is the first of the month, but not the first day of the heatmap
				// we add 2 weeks worth of empty space so months are more visible
				if !first_day && current_day.day0() == 0 {
					for i in 0..14 {
						chunk.data[(i as usize) % 7].push(-1);
					}
				}
				first_day = false;
			}

			let month_name = current_day.format("%b").to_string();

			if current_day == since {
				chunk.months.push((0, month_name));
			}
			else if current_day.day0() == 0 {
				chunk_idx += 1;

				if chunk_idx > months_per_row - 1 {
					heatmap.chunks.push(chunk);
					chunk = Chunk::new(day_of_week);
					chunk_idx = 0;
				}

				chunk
					.months
					.push((chunk.data[day_of_week as usize].len(), month_name));
			}

			let value = grouped_commits.get(&current_day);
			match value {
				Some(val) => {
					chunk.data[day_of_week as usize].push(*val);
					if *val > heatmap.highest_count {
						heatmap.highest_count = *val;
					}
				}
				None => {
					chunk.data[day_of_week as usize].push(0);
				}
			}

			current_day += Duration::days(1);
			day_of_week = current_day.weekday().num_days_from_monday() % 7;
		}

		if chunk_idx <= months_per_row {
			heatmap.chunks.push(chunk);
		}

		heatmap
	}
}

impl Display for Heatmap {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		let start_date = self.since.format("%Y-%b-%d").to_string();
		let end_date = self.until.format("%Y-%b-%d").to_string();
		let commits = self.commits.len();
		let authors = self.commits.iter().unique_by(|c| &c.author.name).count();

		let repos_label = if self.repos == 1 { "repo" } else { "repos" };
		let branches_label = if self.branches == 1 {
			"branch"
		}
		else {
			"branches"
		};
		let authors_label = if authors == 1 { "author" } else { "authors" };
		let commits_label = if commits == 1 { "commit" } else { "commits" };

		writeln!(f, "{} - {}", start_date, end_date).unwrap();
		writeln!(
			f,
			"{} {} | {} {}",
			self.repos, repos_label, self.branches, branches_label
		)
		.unwrap();
		writeln!(f, "{} {}", authors, authors_label).unwrap();
		writeln!(f, "{} {}\n", commits, commits_label).unwrap();

		for chunk in &self.chunks {
			chunk.display(self, f);
			writeln!(f).unwrap();
		}

		write!(f, "\nLess ").unwrap();
		for color in get_color_map() {
			write!(f, "{color}{}{RESET} ", get_char()).unwrap();
		}
		writeln!(f, " More").unwrap();

		Ok(())
	}
}

struct Chunk {
	data: [Vec<i32>; 7],
	months: Vec<(usize, String)>,
}

impl Chunk {
	pub fn new(day_of_week: u32) -> Self {
		let mut chunk = Self {
			data: [vec![], vec![], vec![], vec![], vec![], vec![], vec![]],
			months: vec![],
		};

		if day_of_week != 0 {
			for i in 0..day_of_week {
				chunk.data[i as usize].push(-1);
			}
		}

		chunk
	}

	pub fn display<T: Write>(&self, heatmap: &Heatmap, f: &mut T) {
		writeln!(f, "{}", self.months_row(heatmap)).unwrap();

		for (day, row) in DAYS.iter().zip(&self.data) {
			write!(f, "{day} ").unwrap();
			for val in row {
				match val {
					x if *x >= 0 => {
						let color = &get_color_map()[get_color(*val, heatmap.highest_count)];
						match heatmap.format {
							Format::Chars => write!(f, "{color}{}{RESET} ", get_char()).unwrap(),
							Format::Numbers => {
								let val = val.min(&99);
								write!(f, "{color}{:0>2}{RESET} ", val).unwrap();
							}
						}
					}
					x if *x < 0 => match heatmap.format {
						Format::Chars => write!(f, "{RESET}  ").unwrap(),
						Format::Numbers => write!(f, "{RESET}   ").unwrap(),
					},
					_ => {}
				}
			}
			writeln!(f).unwrap();
		}
	}

	fn months_row(&self, heatmap: &Heatmap) -> String {
		let mut row = "   ".to_string();

		let mut last_index = 0;
		let mul = match heatmap.format {
			Format::Chars => 2,
			Format::Numbers => 3,
		};

		for (index, month) in &self.months {
			let range_size = (index * mul)
				.saturating_sub(last_index * mul)
				.saturating_sub(3);

			for _i in 0..range_size {
				row.push(' ');
			}

			last_index = *index;
			row.push_str(month);
		}

		row
	}
}

#[derive(Clone, Debug, PartialEq, Eq, ValueEnum)]
pub enum HeatmapColors {
	Green,
	Red,
}

#[derive(Clone, Debug, PartialEq, Eq, ValueEnum)]
pub enum ColorLogic {
	ByAmount,
	ByWeight,
}

#[derive(Clone, Debug, PartialEq, Eq, ValueEnum)]
pub enum Format {
	Chars,
	Numbers,
}