import subprocess
import datetime
import os
from collections import defaultdict
import json
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm
import numpy as np
import re
import matplotlib.dates as mdates
from matplotlib.ticker import MultipleLocator

def normalize_name(name):
    """标准化名字：小写+移除非字母数字字符"""
    return re.sub(r'[^a-z0-9]', '', name.lower())

def get_all_contributors(git_dir):
    """获取所有贡献者列表并合并相似名字"""
    cmd = [
        'git',
        '--git-dir', git_dir,
        'log',
        '--all',
        '--pretty=format:%an',
        '--date=short'
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True,encoding='utf-8')
    if result.returncode != 0:
        raise RuntimeError(f"Git命令执行失败: {result.stderr}")
    
    # 原始贡献者列表
    raw_contributors = result.stdout.splitlines()
    
    # 创建名字分组字典
    name_groups = defaultdict(list)
    for name in raw_contributors:
        normalized = normalize_name(name)
        name_groups[normalized].append(name)
    
    # 选择每个组中最长的名字作为代表
    merged_contributors = []
    canonical_names = {}
    
    for group in name_groups.values():
        # 选择组内最长的名字作为标准名称
        canonical_name = min(group, key=len)
        merged_contributors.append(canonical_name)
        
        # 创建映射：组内所有名字 -> 标准名称
        for name in group:
            canonical_names[name] = canonical_name
    
    return sorted(merged_contributors), canonical_names

def get_commit_counts(git_dir, author_names):
    """
    获取多个作者的提交统计数据
    :param git_dir: .git目录路径
    :param author_names: 作者名称列表
    :return: (提交计数, 最早提交日期, 最晚提交日期)
    """
    # 如果没有指定作者，返回空结果
    if not author_names:
        return defaultdict(int), None, None
    
    # 构建Git命令
    cmd = [
        'git',
        '--git-dir', git_dir,
        'log',
        '--all',
        '--pretty=format:%cd',
        '--date=iso',
        '--reverse'
    ]
    
    # 添加多个作者条件
    for name in author_names:
        cmd.extend(['--author', name])
    
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"Git命令执行失败: {result.stderr}")
    
    dates = []
    for line in result.stdout.splitlines():
        if line.strip():
            # 提取日期部分（去掉时区）
            date_str = line.split()[0]
            dates.append(date_str)
    
    if not dates:
        return defaultdict(int), None, None
    
    # 统计每日提交次数
    counts = defaultdict(int)
    for date_str in dates:
        counts[date_str] += 1
    
    min_date = min(dates)
    max_date = max(dates)
    
    return counts, min_date, max_date

def generate_contribution_matrix(commit_counts, min_date, max_date):
    """
    生成贡献热力图矩阵
    :param commit_counts: 每日提交计数字典
    :param min_date: 最早提交日期 (YYYY-MM-DD)
    :param max_date: 最晚提交日期 (YYYY-MM-DD)
    :return: 二维矩阵 (周数 × 星期)
    """
    if not min_date or not max_date:
        return []
    
    start_date = datetime.datetime.strptime(min_date, "%Y-%m-%d").date()
    end_date = datetime.datetime.strptime(max_date, "%Y-%m-%d").date()
    
    # 计算总天数
    total_days = (end_date - start_date).days + 1
    
    # 计算需要的周数
    weeks = (total_days + 6) // 7
    
    # 初始化矩阵
    matrix = [[0] * 7 for _ in range(weeks)]
    
    # 填充矩阵
    for day_offset in range(total_days):
        current_date = start_date + datetime.timedelta(days=day_offset)
        date_str = current_date.strftime("%Y-%m-%d")
        
        # 获取星期几（0=周日，6=周六）
        weekday = current_date.weekday() + 1
        if weekday == 7:  # 将周日从7改为0
            weekday = 0
        
        # 计算周数
        week_num = day_offset // 7
        
        # 确保周数在范围内
        if week_num < weeks:
            matrix[week_num][weekday] = commit_counts.get(date_str, 0)
    
    return matrix, start_date, end_date

def print_contribution_graph(author, matrix, global_min_date, global_max_date):
    """
    打印贡献图（终端可视化）
    """
    print(f"\n{'=' * 60}")
    print(f"贡献者: {author}")
    print(f"时间段: {global_min_date} 至 {global_max_date}")
    print("    Sun Mon Tue Wed Thu Fri Sat")
    
    # 计算总周数和列数
    weeks = len(matrix)
    cols = min(20, weeks)  # 每行最多显示20周
    
    # 分块打印
    for start_week in range(0, weeks, cols):
        end_week = min(start_week + cols, weeks)
        
        print(f"\n周数 {start_week}-{end_week-1}:")
        
        for week_num in range(start_week, end_week):
            if week_num >= len(matrix):
                break
                
            week = matrix[week_num]
            print(f"{week_num:3d} |", end="")
            
            for day in week:
                if day == 0:
                    print("  ", end=" ")
                elif day < 3:
                    print(f"\033[48;5;22m{day}\033[0m", end=" ")  # 浅色背景
                elif day < 6:
                    print(f"\033[48;5;58m{day}\033[0m", end=" ")  # 中等背景
                else:
                    print(f"\033[48;5;94m{day}\033[0m", end=" ")  # 深色背景
            print(f"| {week_num}")
    
    # 打印图例
    print("\n图例:")
    print("  0: 无提交  \033[48;5;22m1-2\033[0m: 1-2次  \033[48;5;58m3-5\033[0m: 3-5次  \033[48;5;94m6+\033[0m: 6+次")

def save_to_json(author, matrix, global_min_date, global_max_date, filename="contributions.json"):
    """保存贡献数据到JSON文件"""
    data = {
        "author": author,
        "min_date": global_min_date,
        "max_date": global_max_date,
        "matrix": matrix
    }
    
    # 读取现有数据或创建新文件
    if os.path.exists(filename):
        with open(filename, "r") as f:
            all_data = json.load(f)
    else:
        all_data = {"contributors": []}
    
    # 添加新数据
    all_data["contributors"].append(data)
    
    # 保存
    with open(filename, "w") as f:
        json.dump(all_data, f, indent=2)
    
    print(f"\n已保存 {author} 的贡献数据到 {filename}")

def save_contribution_heatmap(author, matrix, global_min_date, global_max_date, output_dir):
    """
    保存贡献热力图到文件
    :param author: 贡献者名称
    :param matrix: 贡献矩阵
    :param global_min_date: 全局最早提交日期
    :param global_max_date: 全局最晚提交日期
    :param output_dir: 输出目录
    """
    if not matrix or not global_min_date or not global_max_date:
        return
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 转置矩阵以便绘图（行: 周日-周六, 列: 周数）
    data = np.array(matrix).T
    
    # 定义GitHub风格的颜色映射
    colors = ['#ebedf0', '#9be9a8', '#40c463', '#30a14e', '#216e39']
    # 定义边界：0, 1, 2, 4, 6+
    bounds = [0, 1, 2, 4, 6, 1000]
    cmap = ListedColormap(colors)
    norm = BoundaryNorm(bounds, cmap.N)
    
    # 计算图片尺寸（限制最大宽度）
    weeks = len(matrix)
    cell_size = 0.25  # 每个格子的英寸大小
    max_width = 25    # 最大宽度（英寸）
    
    # 动态调整宽度
    width_in_inches = min(weeks * cell_size, max_width)
    height_in_inches = 3.5  # 增加高度以容纳年份标签
    
    # 创建图形
    fig, ax = plt.subplots(figsize=(width_in_inches, height_in_inches))
    
    # 绘制热力图
    im = ax.imshow(data, cmap=cmap, norm=norm, aspect='auto')
    
    # 设置y轴（星期）
    ax.set_yticks(range(7))
    ax.set_yticklabels(['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat'], fontsize=9)
    ax.invert_yaxis()  # 反转y轴使周日在上
    
    # 添加网格线增强可读性
    ax.xaxis.set_major_locator(MultipleLocator(1))
    ax.yaxis.set_major_locator(MultipleLocator(1))
    ax.xaxis.set_minor_locator(MultipleLocator(0.5))
    ax.yaxis.set_minor_locator(MultipleLocator(0.5))
    # 关闭主刻度网格线
    ax.grid(which='major', axis='both', visible=False)
    # 开启次刻度网格线
    ax.grid(which='minor', axis='both', linestyle='-', linewidth=1, alpha=0.7, color='black')
    ax.set_xticklabels([])  # 隐藏X轴标签

    # 添加年份标记
    start_date = datetime.datetime.strptime(global_min_date, "%Y-%m-%d")
    end_date = datetime.datetime.strptime(global_max_date, "%Y-%m-%d")
    
    # 找出每个年份开始的位置（周数）
    year_positions = {}
    current_year = start_date.year
    while current_year <= end_date.year:
        first_day_of_year = datetime.datetime(current_year, 1, 1)
        if first_day_of_year < start_date:
            first_day_of_year = start_date
        
        # 计算该日期在总时间中的位置（周数）
        days_diff = (first_day_of_year - start_date).days
        week_num = days_diff // 7
        year_positions[current_year] = week_num
        current_year += 1
    
    # 在顶部添加年份标签
    for year, pos in year_positions.items():
        if 0 <= pos < weeks:
            ax.text(pos, -1.2, str(year), fontsize=9, ha='center', va='center')
    
    # 设置标题
    plt.title(f"{author}: {global_min_date} to {global_max_date}", 
              fontsize=12, pad=15)
    
    # 添加颜色条
    cbar = fig.colorbar(im, ax=ax, orientation='horizontal', 
                        pad=0.12, aspect=50)
    cbar.set_ticks([0.5, 1.5, 3.0, 5.0, 8.0])  # 5 个刻度位置
    cbar.set_ticklabels(['0', '1', '2-3', '4-5', '6+'])
    cbar.ax.tick_params(labelsize=7)
    cbar.set_label('Commit Count', fontsize=9)
    
    # 调整布局
    plt.tight_layout()
    plt.subplots_adjust(top=0.85, bottom=0.2)
    
    # 保存图片
    safe_author = "".join(c if c.isalnum() else "_" for c in author)
    output_path = os.path.join(output_dir, f"{safe_author}.png")
    plt.savefig(output_path, bbox_inches='tight', dpi=150)  # 提高DPI获得更高清图片
    plt.close(fig)
    
    print(f"  已保存热力图到: {output_path}")

def generate_contributor_graphs(git_dir, output_dir="contribution_heatmaps", save_json=False):
    """为每个贡献者生成贡献图"""
    contributors, name_mapping = get_all_contributors(git_dir)
    print(f"发现 {len(contributors)} 位贡献者 (合并后)")
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 第一步：收集所有有效贡献者的提交日期范围
    all_min_dates = []
    all_max_dates = []
    valid_contributors = []
    
    # 创建反向映射：标准名称 -> 所有原始名称
    reverse_mapping = defaultdict(list)
    for orig, canonical in name_mapping.items():
        reverse_mapping[canonical].append(orig)
    
    # 收集所有贡献者的日期范围
    for author in contributors:
        # 获取该标准名称对应的所有原始名称
        author_aliases = reverse_mapping[author]
        commit_counts, min_date, max_date = get_commit_counts(git_dir, author_aliases)
        
        if not min_date or not max_date:
            print(f"  - {author} 没有提交记录 (别名: {', '.join(author_aliases)})")
            continue
            
        valid_contributors.append(author)
        all_min_dates.append(min_date)
        all_max_dates.append(max_date)

    if not valid_contributors:
        print("没有有效的贡献者")
        return

    global_min_date = min(all_min_dates)
    global_max_date = max(all_max_dates)

    # 第二步：使用统一时间范围生成每个贡献者的贡献图
    for i, author in enumerate(contributors):
        print(f"\n处理贡献者 {i+1}/{len(contributors)}: {author}")
        author_aliases = reverse_mapping[author]
        commit_counts, min_date, max_date = get_commit_counts(git_dir, author_aliases)
        
        if not min_date or not max_date:
            print(f"  - {author} 没有提交记录 (别名: {', '.join(author_aliases)})")
            matrix = generate_contribution_matrix({}, global_min_date, global_max_date)[0]
        else:
            matrix = generate_contribution_matrix(commit_counts, global_min_date, global_max_date)[0]
        
        print_contribution_graph(author, matrix, global_min_date, global_max_date)
        save_contribution_heatmap(author, matrix, global_min_date, global_max_date, output_dir)
        
        if save_json:
            save_to_json(author, matrix, global_min_date, global_max_date)

def main():
    # 自动检测当前目录的.git
    git_dir = os.path.abspath(".git")
    if not os.path.exists(git_dir):
        raise FileNotFoundError("在当前目录未找到.git文件夹")
    
    # 输出目录
    output_dir = "contribution_heatmaps"
    
    # 为所有贡献者生成图
    generate_contributor_graphs(git_dir, output_dir, save_json=False)
    
    print("\n所有贡献者分析完成!")
    print(f"热力图已保存到: {os.path.abspath(output_dir)}")

if __name__ == "__main__":
    main()