#!/usr/bin/env python3
"""Disk usage analyzer - reports filesystem usage for key directories."""
import os
import sys
import shutil
from pathlib import Path

def get_dir_size(path):
    """Calculate total size of directory in bytes."""
    total = 0
    try:
        for entry in os.scandir(path):
            if entry.is_file(follow_symlinks=False):
                total += entry.stat().st_size
            elif entry.is_dir(follow_symlinks=False):
                total += get_dir_size(entry.path)
    except (PermissionError, FileNotFoundError):
        pass
    return total

def format_size(size_bytes):
    """Format bytes to human readable string."""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_bytes < 1024.0:
            return f"{size_bytes:.2f} {unit}"
        size_bytes /= 1024.0
    return f"{size_bytes:.2f} PB"

def main():
    output_dir = Path("/root/.hermes/wizard-outputs/spark")
    output_file = output_dir / "disk_usage_report.txt"
    
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Key directories to check
    dirs_to_check = [
        "/root",
        "/tmp", 
        "/root/.hermes",
    ]
    
    results = []
    results.append("Disk Usage Report")
    results.append("=" * 50)
    results.append("")
    
    total_size = 0
    for dir_path in dirs_to_check:
        if os.path.exists(dir_path):
            size = get_dir_size(dir_path)
            total_size += size
            results.append(f"{dir_path}: {format_size(size)}")
        else:
            results.append(f"{dir_path}: NOT FOUND")
    
    results.append("")
    results.append(f"Total: {format_size(total_size)}")
    
    content = "\n".join(results)
    output_file.write_text(content)
    print(content)
    return 0

if __name__ == "__main__":
    sys.exit(main())