#!/usr/bin/env python3
"""
File Extension Disk Usage Analyzer

Analyzes a directory and reports disk usage by file extension.
Useful for identifying large file types in project directories.
"""

import os
import sys
from collections import defaultdict
from pathlib import Path


def analyze_directory(directory: str, extension: str = None) -> dict:
    """
    Analyze a directory and return file sizes grouped by extension.

    Args:
        directory: Path to directory to analyze
        extension: Optional specific extension to filter (e.g., '.py')

    Returns:
        Dictionary mapping extensions to total size in bytes and file count
    """
    stats = defaultdict(lambda: {"size": 0, "count": 0})
    dir_path = Path(directory)

    if not dir_path.exists():
        raise FileNotFoundError(f"Directory not found: {directory}")

    if not dir_path.is_dir():
        raise NotADirectoryError(f"Not a directory: {directory}")

    for root, dirs, files in os.walk(dir_path):
        # Skip hidden directories and common ignore patterns
        dirs[:] = [d for d in dirs if not d.startswith('.') and d not in ('__pycache__', 'node_modules', '.git')]

        for filename in files:
            if filename.startswith('.'):
                continue

            filepath = Path(root) / filename
            try:
                size = filepath.stat().st_size
            except OSError:
                continue

            file_ext = filepath.suffix.lower() or '(no extension)'

            if extension and file_ext != extension.lower():
                continue

            stats[file_ext]["size"] += size
            stats[file_ext]["count"] += 1

    # Sort by total size descending
    return dict(sorted(stats.items(), key=lambda x: x[1]["size"], reverse=True))


def format_size(size_bytes: int) -> str:
    """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():
    if len(sys.argv) < 2:
        print("Usage: file_ext_analyzer.py <directory> [extension]")
        print("Example: file_ext_analyzer.py /path/to/project .py")
        sys.exit(1)

    directory = sys.argv[1]
    extension = sys.argv[2] if len(sys.argv) > 2 else None

    try:
        stats = analyze_directory(directory, extension)

        if not stats:
            print("No files found.")
            sys.exit(0)

        print(f"\nDisk Usage Analysis: {directory}")
        print("-" * 50)

        for ext, data in stats.items():
            print(f"{ext:15} {format_size(data['size']):>12} ({data['count']:>5} files)")

        total_size = sum(d["size"] for d in stats.values())
        total_files = sum(d["count"] for d in stats.values())
        print("-" * 50)
        print(f"{'TOTAL':15} {format_size(total_size):>12} ({total_files:>5} files)")

    except (FileNotFoundError, NotADirectoryError) as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()