import os
import argparse
import re
from collections import defaultdict
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment, Border, Side

# 配置默认参数
DEFAULT_EXTENSIONS = {'py', 'pyw', 'api', 'html', 'js', 'css'}
AUTO_EXCLUDE_DIRS = {'.env', 'layui', 'node_modules'}  # 不区分大小写
AUTO_EXCLUDE_KEYWORDS = {
    'font-awesome',
    'bootstrap',
    re.compile(r'.*\.lib\.js$', re.IGNORECASE),
    re.compile(r'.*\.lib\.css$', re.IGNORECASE)
}

class CodeCounter:
    def __init__(self, target_dir, exclude_paths=None, extensions=None):
        self.target_dir = os.path.abspath(target_dir)
        self.exclude_paths = [os.path.abspath(p) for p in exclude_paths] if exclude_paths else []
        self.extensions = extensions or DEFAULT_EXTENSIONS
        self.file_stats = []  # 存储每个文件的统计信息
        self.total_lines = 0
        self.total_chars = 0
        self.total_files = 0

    def _should_exclude_file(self, file_path):
        """判断文件是否需要排除"""
        # 检查自动排除的目录（任何深度）
        dir_parts = file_path.lower().split(os.sep)
        if any(ed.lower() in dir_parts for ed in AUTO_EXCLUDE_DIRS):
            return True

        # 检查手动排除的路径
        for ep in self.exclude_paths:
            if file_path.startswith(ep):
                return True

        # 检查自动排除的关键字/正则
        file_name = os.path.basename(file_path).lower()
        for keyword in AUTO_EXCLUDE_KEYWORDS:
            if isinstance(keyword, re.Pattern):
                if keyword.match(file_name):
                    return True
            else:
                if keyword.lower() in file_path.lower():
                    return True

        return False

    def _get_file_ext(self, file_path):
        """获取文件扩展名（不含点）"""
        ext = os.path.splitext(file_path)[1].lstrip('.').lower()
        return ext

    def _count_file(self, file_path):
        """统计单个文件的行数和字符数"""
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                lines = len(content.splitlines())
                chars = len(content)
                return lines, chars
        except Exception as e:
            print(f"警告：无法读取文件 {file_path} - {e}")
            return 0, 0

    def scan(self):
        """扫描目录并统计代码量"""
        print(f"开始扫描目录：{self.target_dir}")
        print(f"包含扩展名：{', '.join(self.extensions)}")
        print(f"排除路径：{', '.join(self.exclude_paths) if self.exclude_paths else '无'}")
        print("-" * 80)

        for root, dirs, files in os.walk(self.target_dir):
            # 跳过自动排除的目录（优化遍历性能）
            dirs[:] = [d for d in dirs if d.lower() not in AUTO_EXCLUDE_DIRS]

            for file in files:
                file_path = os.path.join(root, file)
                ext = self._get_file_ext(file_path)

                # 检查扩展名是否符合要求
                if ext not in self.extensions:
                    continue

                # 检查是否需要排除
                if self._should_exclude_file(file_path):
                    continue

                # 统计文件
                lines, chars = self._count_file(file_path)
                self.file_stats.append({
                    'path': file_path,
                    'lines': lines,
                    'chars': chars
                })
                self.total_lines += lines
                self.total_chars += chars
                self.total_files += 1

                # 实时输出进度
                print(f"已统计：{file_path} | 行数：{lines} | 字符数：{chars}")

        print("-" * 80)
        print(f"统计完成！")
        print(f"总文件数：{self.total_files}")
        print(f"总行数：{self.total_lines}")
        print(f"总字符数：{self.total_chars}")

    def export_xlsx(self, output_path):
        """导出统计结果到XLSX文件"""
        wb = Workbook()
        ws = wb.active
        ws.title = "代码量统计"

        # 设置表头
        headers = ['文件路径', '行数', '字符数']
        for col, header in enumerate(headers, 1):
            cell = ws.cell(row=1, column=col, value=header)
            cell.font = Font(bold=True, size=12)
            cell.alignment = Alignment(horizontal='center')
            cell.border = Border(
                left=Side(style='thin'),
                right=Side(style='thin'),
                top=Side(style='thin'),
                bottom=Side(style='thin')
            )

        # 填充数据
        for row, stat in enumerate(self.file_stats, 2):
            ws.cell(row=row, column=1, value=stat['path'])
            ws.cell(row=row, column=2, value=stat['lines'])
            ws.cell(row=row, column=3, value=stat['chars'])

            # 设置单元格格式
            for col in range(1, 4):
                cell = ws.cell(row=row, column=col)
                cell.border = Border(
                    left=Side(style='thin'),
                    right=Side(style='thin'),
                    top=Side(style='thin'),
                    bottom=Side(style='thin')
                )
                if col > 1:
                    cell.alignment = Alignment(horizontal='center')

        # 添加总计行
        total_row = len(self.file_stats) + 3
        ws.cell(row=total_row, column=1, value='总计')
        ws.cell(row=total_row, column=2, value=self.total_lines)
        ws.cell(row=total_row, column=3, value=self.total_chars)
        for col in range(1, 4):
            cell = ws.cell(row=total_row, column=col)
            cell.font = Font(bold=True)
            cell.alignment = Alignment(horizontal='center')
            cell.border = Border(
                left=Side(style='thin'),
                right=Side(style='thin'),
                top=Side(style='thin'),
                bottom=Side(style='thin')
            )

        # 调整列宽
        ws.column_dimensions['A'].width = 80
        ws.column_dimensions['B'].width = 15
        ws.column_dimensions['C'].width = 15

        # 保存文件
        try:
            wb.save(output_path)
            print(f"\nXLSX文件已导出：{os.path.abspath(output_path)}")
        except Exception as e:
            print(f"\n导出XLSX失败：{e}")

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Python代码量统计工具')
    parser.add_argument('dir', help='要统计的目标目录')
    parser.add_argument('--exclude', nargs='+', default=[], 
                        help='需要排除的路径（如：TempLate/cdn TempLate/AKS）')
    parser.add_argument('--extensions', nargs='+', default=list(DEFAULT_EXTENSIONS),
                        help='要统计的文件扩展名（默认：py pyw api html js css）')
    parser.add_argument('--export', help='导出XLSX文件路径（如：code_stats.xlsx）')
    
    args = parser.parse_args()

    # 处理排除路径（相对路径转绝对路径）
    exclude_paths = []
    for ep in args.exclude:
        if os.path.isabs(ep):
            exclude_paths.append(ep)
        else:
            exclude_paths.append(os.path.join(os.path.abspath(args.dir), ep))

    # 创建统计器并执行统计
    counter = CodeCounter(
        target_dir=args.dir,
        exclude_paths=exclude_paths,
        extensions=set(args.extensions)
    )
    counter.scan()

    # 导出XLSX（如果指定）
    if args.export:
        counter.export_xlsx(args.export)

if __name__ == '__main__':
    main()