git-heatmap/src/heatmap.rs

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

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

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

pub struct Heatmap {
	data: [Vec<i32>; 7],
	since: NaiveDate,
	until: NaiveDate,
	commits: Vec<Commit>,
	months: Vec<(usize, String)>,
	highest_count: i32,
}

impl Heatmap {
	pub fn new(since: NaiveDate, until: NaiveDate, commits: Vec<Commit>) -> Self {
		let mut heatmap = Self {
			data: [vec![], vec![], vec![], vec![], vec![], vec![], vec![]],
			since,
			until,
			commits,
			months: vec![],
			highest_count: 0,
		};

		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;

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

		while current_day <= until {
			let month_name = current_day.format("%b").to_string();

			if current_day == since {
				heatmap.months.push((0, month_name));
			}
			else if current_day.day0() == 0 {
				heatmap
					.months
					.push((heatmap.data[day_of_week as usize].len(), month_name));
			}

			let value = grouped_commits.get(&current_day);
			match value {
				Some(val) => {
					heatmap.data[day_of_week as usize].push(*val);
					if *val > heatmap.highest_count {
						heatmap.highest_count = *val;
					}
				}
				None => heatmap.data[day_of_week as usize].push(0),
			}
			// println!("{} {value:?}", current_day);
			current_day += Duration::days(1);
			day_of_week = current_day.weekday().num_days_from_monday() % 7;
		}

		heatmap
	}

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

		let mut last_index = 0;

		for (index, month) in &self.months {
			let range_size = (index * 2).saturating_sub(last_index * 2).saturating_sub(3);
			for _i in 0..range_size {
				row.push(' ');
			}
			last_index = *index;
			row.push_str(month);
		}

		row
	}
}

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().to_string();

		write!(f, "{} - {}\n", start_date, end_date).unwrap();
		write!(f, "{} commits\n\n", commits).unwrap();
		write!(f, "{}\n", self.months_row()).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, self.highest_count)];
						write!(f, "{color}{}{RESET} ", get_char()).unwrap();
					}
					x if *x < 0 => {
						write!(f, "{RESET}  ").unwrap();
					}
					_ => {}
				}
			}
			write!(f, "\n").unwrap();
		}

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

		Ok(())
	}
}

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