import os
import sys
import argparse
import json
from datetime import datetime
import matplotlib.pyplot as plt
from collections import defaultdict

class DiskSpaceAnalyzer:
    def __init__(self, path="/"):
        self.path = path
        self.total_size = 0
        self.file_count = 0
        self.dir_count = 0
        self.largest_files = []
        self.dir_sizes = {}
        self.file_types = defaultdict(int)
        self.errors = []
        
    def get_disk_usage(self):
        """獲取磁盤使用情況"""
        try:
            if os.name == 'nt':  # Windows
                import shutil
                total, used, free = shutil.disk_usage(self.path)
            else:  # Unix/Linux/macOS
                statvfs = os.statvfs(self.path)
                total = statvfs.f_frsize * statvfs.f_blocks
                free = statvfs.f_frsize * statvfs.f_bavail
                used = total - free
                
            return {
                'total': total,
                'used': used,
                'free': free,
                'percentage': (used / total) * 100 if total > 0 else 0
            }
        except Exception as e:
            self.errors.append(f"獲取磁盤使用情況失敗: {e}")
            return None
            
    def format_bytes(self, bytes_value):
        """格式化字節(jié)單位"""
        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
            if bytes_value < 1024.0:
                return f"{bytes_value:.2f} {unit}"
            bytes_value /= 1024.0
        return f"{bytes_value:.2f} PB"
        
    def get_file_extension(self, filename):
        """獲取文件擴(kuò)展名"""
        _, ext = os.path.splitext(filename)
        return ext.lower() if ext else 'no_extension'
        
    def analyze_directory(self, max_depth=3, min_size_mb=1):
        """分析目錄大小"""
        print(f"開始分析目錄: {self.path}")
        print(f"最大深度: {max_depth}, 最小文件大小: {min_size_mb}MB")
        
        try:
            self._walk_directory(self.path, 0, max_depth, min_size_mb * 1024 * 1024)
            self._sort_results()
            return True
        except Exception as e:
            self.errors.append(f"分析目錄時(shí)出錯(cuò): {e}")
            return False
            
    def _walk_directory(self, current_path, current_depth, max_depth, min_size):
        """遞歸遍歷目錄"""
        try:
            # 檢查是否超過(guò)最大深度
            if current_depth > max_depth:
                return
                
            # 獲取目錄內(nèi)容
            items = os.listdir(current_path)
            dir_size = 0
            
            for item in items:
                item_path = os.path.join(current_path, item)
                
                try:
                    # 跳過(guò)隱藏文件和目錄
                    if item.startswith('.'):
                        continue
                        
                    # 獲取文件/目錄信息
                    stat_info = os.stat(item_path)
                    
                    if os.path.isfile(item_path):
                        # 處理文件
                        file_size = stat_info.st_size
                        dir_size += file_size
                        self.file_count += 1
                        
                        # 記錄大于最小大小的文件
                        if file_size >= min_size:
                            self.largest_files.append({
                                'path': item_path,
                                'size': file_size,
                                'modified': datetime.fromtimestamp(stat_info.st_mtime)
                            })
                            
                        # 統(tǒng)計(jì)文件類型
                        ext = self.get_file_extension(item)
                        self.file_types[ext] += file_size
                        
                    elif os.path.isdir(item_path):
                        # 處理目錄
                        self.dir_count += 1
                        subdir_size = self._walk_directory(item_path, current_depth + 1, max_depth, min_size)
                        dir_size += subdir_size
                        
                except PermissionError:
                    self.errors.append(f"權(quán)限不足: {item_path}")
                except Exception as e:
                    self.errors.append(f"處理 {item_path} 時(shí)出錯(cuò): {e}")
                    
            # 記錄當(dāng)前目錄大小
            self.dir_sizes[current_path] = dir_size
            return dir_size
            
        except PermissionError:
            self.errors.append(f"權(quán)限不足，無(wú)法訪問(wèn)目錄: {current_path}")
            return 0
        except Exception as e:
            self.errors.append(f"遍歷目錄 {current_path} 時(shí)出錯(cuò): {e}")
            return 0
            
    def _sort_results(self):
        """對(duì)結(jié)果進(jìn)行排序"""
        # 按大小排序目錄
        self.dir_sizes = dict(sorted(self.dir_sizes.items(), key=lambda x: x[1], reverse=True))
        
        # 按大小排序文件
        self.largest_files.sort(key=lambda x: x['size'], reverse=True)
        
        # 按大小排序文件類型
        self.file_types = dict(sorted(self.file_types.items(), key=lambda x: x[1], reverse=True))
        
    def get_top_directories(self, count=10):
        """獲取最大的幾個(gè)目錄"""
        return dict(list(self.dir_sizes.items())[:count])
        
    def get_top_files(self, count=10):
        """獲取最大的幾個(gè)文件"""
        return self.largest_files[:count]
        
    def generate_text_report(self):
        """生成文本報(bào)告"""
        report = []
        report.append("=" * 80)
        report.append("磁盤空間使用分析報(bào)告")
        report.append("=" * 80)
        report.append(f"分析時(shí)間: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        report.append(f"分析路徑: {self.path}")
        report.append("")
        
        # 磁盤使用情況
        disk_usage = self.get_disk_usage()
        if disk_usage:
            report.append("磁盤使用情況:")
            report.append(f"  總空間: {self.format_bytes(disk_usage['total'])}")
            report.append(f"  已使用: {self.format_bytes(disk_usage['used'])}")
            report.append(f"  可用空間: {self.format_bytes(disk_usage['free'])}")
            report.append(f"  使用率: {disk_usage['percentage']:.2f}%")
            report.append("")
            
        # 統(tǒng)計(jì)信息
        report.append("統(tǒng)計(jì)信息:")
        report.append(f"  文件總數(shù): {self.file_count:,}")
        report.append(f"  目錄總數(shù): {self.dir_count:,}")
        report.append(f"  總大小: {self.format_bytes(sum(self.dir_sizes.values()))}")
        report.append("")
        
        # 最大的目錄
        report.append("最大的目錄 (Top 10):")
        report.append("-" * 50)
        top_dirs = self.get_top_directories(10)
        for i, (dir_path, size) in enumerate(top_dirs.items(), 1):
            # 只顯示相對(duì)路徑
            relative_path = os.path.relpath(dir_path, self.path) if dir_path != self.path else "."
            report.append(f"{i:2d}. {relative_path:<50} {self.format_bytes(size)}")
        report.append("")
        
        # 最大的文件
        report.append("最大的文件 (Top 10):")
        report.append("-" * 50)
        top_files = self.get_top_files(10)
        for i, file_info in enumerate(top_files, 1):
            # 只顯示相對(duì)路徑
            relative_path = os.path.relpath(file_info['path'], self.path)
            modified_time = file_info['modified'].strftime('%Y-%m-%d %H:%M:%S')
            report.append(f"{i:2d}. {relative_path:<50} {self.format_bytes(file_info['size'])} ({modified_time})")
        report.append("")
        
        # 文件類型統(tǒng)計(jì)
        report.append("文件類型統(tǒng)計(jì) (Top 10):")
        report.append("-" * 50)
        top_types = dict(list(self.file_types.items())[:10])
        for i, (ext, size) in enumerate(top_types.items(), 1):
            ext_display = ext if ext else '(無(wú)擴(kuò)展名)'
            report.append(f"{i:2d}. {ext_display:<20} {self.format_bytes(size)}")
        report.append("")
        
        # 錯(cuò)誤信息
        if self.errors:
            report.append("錯(cuò)誤信息:")
            report.append("-" * 50)
            for error in self.errors[:10]:  # 只顯示前10個(gè)錯(cuò)誤
                report.append(f"  {error}")
            if len(self.errors) > 10:
                report.append(f"  ... 還有 {len(self.errors) - 10} 個(gè)錯(cuò)誤")
            report.append("")
            
        report.append("=" * 80)
        return "\n".join(report)
        
    def generate_json_report(self, filename=None):
        """生成JSON報(bào)告"""
        report = {
            'analysis_time': datetime.now().isoformat(),
            'path': self.path,
            'disk_usage': self.get_disk_usage(),
            'statistics': {
                'total_files': self.file_count,
                'total_directories': self.dir_count,
                'total_size': sum(self.dir_sizes.values())
            },
            'top_directories': self.get_top_directories(20),
            'top_files': [
                {
                    'path': f['path'],
                    'size': f['size'],
                    'modified': f['modified'].isoformat()
                }
                for f in self.get_top_files(20)
            ],
            'file_types': self.file_types,
            'errors': self.errors
        }
        
        if filename:
            try:
                with open(filename, 'w', encoding='utf-8') as f:
                    json.dump(report, f, indent=2, ensure_ascii=False)
                print(f"JSON報(bào)告已保存到: {filename}")
                return True
            except Exception as e:
                print(f"保存JSON報(bào)告時(shí)出錯(cuò): {e}")
                return False
        else:
            return json.dumps(report, indent=2, ensure_ascii=False)
            
    def generate_charts(self, prefix="disk_analysis"):
        """生成圖表"""
        try:
            # 創(chuàng)建圖表目錄
            chart_dir = f"{prefix}_charts"
            if not os.path.exists(chart_dir):
                os.makedirs(chart_dir)
                
            # 1. 目錄大小餅圖
            top_dirs = self.get_top_directories(8)
            other_size = sum(self.dir_sizes.values()) - sum(top_dirs.values())
            
            if top_dirs:
                labels = [os.path.basename(path) if path != self.path else '根目錄' for path in top_dirs.keys()]
                sizes = list(top_dirs.values())
                
                if other_size > 0:
                    labels.append('其他')
                    sizes.append(other_size)
                    
                plt.figure(figsize=(10, 8))
                plt.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=90)
                plt.title('目錄大小分布')
                plt.axis('equal')
                plt.savefig(os.path.join(chart_dir, f'{prefix}_directories.png'), 
                           dpi=300, bbox_inches='tight')
                plt.close()
                
            # 2. 文件類型柱狀圖
            top_types = dict(list(self.file_types.items())[:10])
            if top_types:
                labels = [ext if ext else '(無(wú)擴(kuò)展名)' for ext in top_types.keys()]
                sizes = [size / (1024*1024) for size in top_types.values()]  # 轉(zhuǎn)換為MB
                
                plt.figure(figsize=(12, 6))
                plt.bar(labels, sizes)
                plt.title('文件類型大小分布 (MB)')
                plt.xlabel('文件類型')
                plt.ylabel('大小 (MB)')
                plt.xticks(rotation=45)
                plt.tight_layout()
                plt.savefig(os.path.join(chart_dir, f'{prefix}_filetypes.png'), 
                           dpi=300, bbox_inches='tight')
                plt.close()
                
            # 3. 最大文件柱狀圖
            top_files = self.get_top_files(10)
            if top_files:
                labels = [os.path.basename(f['path']) for f in top_files]
                sizes = [f['size'] / (1024*1024) for f in top_files]  # 轉(zhuǎn)換為MB
                
                plt.figure(figsize=(12, 6))
                plt.bar(labels, sizes)
                plt.title('最大文件大小 (MB)')
                plt.xlabel('文件名')
                plt.ylabel('大小 (MB)')
                plt.xticks(rotation=45)
                plt.tight_layout()
                plt.savefig(os.path.join(chart_dir, f'{prefix}_largest_files.png'), 
                           dpi=300, bbox_inches='tight')
                plt.close()
                
            print(f"圖表已保存到目錄: {chart_dir}")
            return True
            
        except Exception as e:
            print(f"生成圖表時(shí)出錯(cuò): {e}")
            return False

def main():
    parser = argparse.ArgumentParser(description="磁盤空間使用報(bào)告工具")
    parser.add_argument("path", nargs='?', default=".", help="要分析的目錄路徑 (默認(rèn): 當(dāng)前目錄)")
    parser.add_argument("-d", "--depth", type=int, default=3, help="分析深度 (默認(rèn): 3)")
    parser.add_argument("-s", "--min-size", type=float, default=1.0, 
                       help="最小文件大小(MB) (默認(rèn): 1.0)")
    parser.add_argument("-o", "--output", help="保存文本報(bào)告到文件")
    parser.add_argument("-j", "--json", help="保存JSON報(bào)告到文件")
    parser.add_argument("-c", "--charts", action="store_true", help="生成圖表")
    parser.add_argument("--chart-prefix", default="disk_analysis", help="圖表文件前綴")
    
    args = parser.parse_args()
    
    # 檢查路徑是否存在
    if not os.path.exists(args.path):
        print(f"錯(cuò)誤: 路徑 '{args.path}' 不存在")
        sys.exit(1)
        
    if not os.path.isdir(args.path):
        print(f"錯(cuò)誤: '{args.path}' 不是一個(gè)目錄")
        sys.exit(1)
        
    try:
        analyzer = DiskSpaceAnalyzer(args.path)
        
        # 分析目錄
        if not analyzer.analyze_directory(args.depth, args.min_size):
            print("目錄分析失敗")
            sys.exit(1)
            
        # 生成文本報(bào)告
        report = analyzer.generate_text_report()
        
        # 輸出報(bào)告
        if args.output:
            try:
                with open(args.output, 'w', encoding='utf-8') as f:
                    f.write(report)
                print(f"文本報(bào)告已保存到: {args.output}")
            except Exception as e:
                print(f"保存文本報(bào)告時(shí)出錯(cuò): {e}")
        else:
            print(report)
            
        # 生成JSON報(bào)告
        if args.json:
            analyzer.generate_json_report(args.json)
            
        # 生成圖表
        if args.charts:
            analyzer.generate_charts(args.chart_prefix)
            
    except KeyboardInterrupt:
        print("\n\n用戶中斷操作")
    except Exception as e:
        print(f"程序執(zhí)行出錯(cuò): {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()

# 分析當(dāng)前目錄
python disk_space_report.py

# 分析指定目錄
python disk_space_report.py /path/to/directory

# 設(shè)置分析深度為5層
python disk_space_report.py /path/to/directory -d 5

# 只分析大于10MB的文件
python disk_space_report.py /path/to/directory -s 10

# 保存報(bào)告到文件
python disk_space_report.py /path/to/directory -o report.txt

# 生成JSON報(bào)告
python disk_space_report.py /path/to/directory -j report.json

# 生成圖表
python disk_space_report.py /path/to/directory -c --chart-prefix my_analysis

python disk_space_report.py /home/user

python disk_space_report.py /var/log -d 5 -s 0.5 -o log_analysis.txt -j log_analysis.json -c

python disk_space_report.py / -d 2 -s 50 -o system_report.txt