git-heatmap/src/heatmap.rs

use std::collections::BTreeMap;

use chrono::{DateTime, Datelike, Duration, Local};
use clap::ValueEnum;

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

pub struct Heatmap {
	data: [Vec<i32>; 7],
	start_date: DateTime<Local>,
	end_date: DateTime<Local>,
	commits: Vec<Commit>,
	months: Vec<(usize, String)>,
	highest_count: i32,
}

impl Heatmap {
	pub fn new(
		start_date: DateTime<Local>,
		end_date: DateTime<Local>,
		commits: Vec<Commit>,
	) -> Self {
		let mut heatmap = Self {
			data: [vec![], vec![], vec![], vec![], vec![], vec![], vec![]],
			start_date,
			end_date,
			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 = start_date;
		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 <= end_date {
			let month_name = current_day.format("%b").to_string();

			if current_day == start_date {
				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.date_naive());
			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:?} {highest_count}", current_day.date_naive());
			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
	}

	pub fn print(&self) {
		println!(
			"{} - {}",
			self.start_date.format("%Y-%b-%d"),
			self.end_date.format("%Y-%b-%d")
		);
		println!("{} commits\n", &self.commits.len());
		println!("{}", self.months_row());
		for (day, row) in DAYS.iter().zip(&self.data) {
			print!("{day} ");
			for val in row {
				match val {
					x if *x >= 0 => {
						let color = &get_color_map()[get_color(*val, self.highest_count)];
						print!("{color}{}{RESET} ", get_char());
					}
					x if *x < 0 => {
						print!("{RESET}  ");
					}
					_ => {}
				}
			}
			print!("\n");
		}
		print!("\nLess ");
		for color in get_color_map() {
			print!("{color}{}{RESET} ", get_char())
		}
		print!(" More\n");
	}
}

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