#!/usr/bin/env python3
"""
folder_size_report.py

Walks a directory tree and prints the top N largest subfolders and
top N largest individual files, with human-readable sizes. Useful
for quickly finding what's eating disk space, similar to WinDirStat
but scriptable and dependency-free (standard library only).

Usage:
    python folder_size_report.py "C:\\Users\\me" --top 15
"""

import argparse
import sys
from pathlib import Path


def human_size(num_bytes: float) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if num_bytes < 1024:
            return f"{num_bytes:.1f} {unit}"
        num_bytes /= 1024
    return f"{num_bytes:.1f} PB"


def scan(root: Path):
    folder_sizes: dict[Path, int] = {}
    file_sizes: list[tuple[Path, int]] = []

    for path in root.rglob("*"):
        try:
            if path.is_file():
                size = path.stat().st_size
                file_sizes.append((path, size))
                parent = path.parent
                folder_sizes[parent] = folder_sizes.get(parent, 0) + size
        except (PermissionError, OSError):
            continue

    # Roll sizes up into ancestor folders too
    rolled_up: dict[Path, int] = dict(folder_sizes)
    for folder, size in folder_sizes.items():
        for ancestor in folder.parents:
            if root not in ancestor.parents and ancestor != root:
                continue
            if ancestor == root or root in ancestor.parents:
                rolled_up[ancestor] = rolled_up.get(ancestor, 0) + size
            if ancestor == root:
                break

    return rolled_up, file_sizes


def main() -> None:
    parser = argparse.ArgumentParser(description="Report largest folders and files under a path.")
    parser.add_argument("path", type=Path, help="Root folder to scan")
    parser.add_argument("--top", type=int, default=10, help="How many entries to show in each list (default 10)")
    args = parser.parse_args()

    root = args.path
    if not root.is_dir():
        sys.exit(f"Error: '{root}' is not a valid folder.")

    print(f"Scanning '{root}' ... (this may take a while for large trees)\n")
    folder_sizes, file_sizes = scan(root)

    print(f"Top {args.top} largest folders:")
    for folder, size in sorted(folder_sizes.items(), key=lambda kv: kv[1], reverse=True)[: args.top]:
        print(f"  {human_size(size):>10}  {folder}")

    print(f"\nTop {args.top} largest files:")
    for file, size in sorted(file_sizes, key=lambda kv: kv[1], reverse=True)[: args.top]:
        print(f"  {human_size(size):>10}  {file}")


if __name__ == "__main__":
    main()