from dataclasses import dataclass
import dataclasses
from decimal import Decimal
from functools import reduce
from typing import NamedTuple, Optional
from line_profiler import profile
import psutil

@dataclass
class ProcessSwapInfo:
    pid: int
    name: str
    swap: int

@dataclass
class ProcessSwapGroup:
    name: str
    processes: list[ProcessSwapInfo] = dataclasses.field(default_factory=list)
    
    @property
    def total_swap(self) -> int:
        return sum(map(lambda x: x.swap, self.processes))

type ProcessName = str

def convert_bytes_to_mb(value: int) -> Decimal:
    return Decimal(value) / (1024 * 1024)

@profile
def get_process_swap_usage() -> None:
    print(f"{'PID':>6} {'Name':<25} {'Swap (MB)':>10}")
    print("-" * 45)
    
    total_swap_used = 0

    swap_used: list[ProcessSwapInfo] = []

    for proc in psutil.process_iter(['pid', 'name', 'memory_full_info']):
        try:
            with proc.oneshot():
                memory_info: Optional[NamedTuple] = proc.info['memory_full_info']
                swap: int = getattr(memory_info, "swap", 0)
                if swap > 0:
                    swap_used.append(ProcessSwapInfo(proc.pid, proc.name(), swap))
                    total_swap_used += swap
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue
    
    swap_used.sort(key=lambda x: x.swap)

    def _add_to_group(groups: dict[ProcessName, ProcessSwapGroup], process: ProcessSwapInfo) -> dict[ProcessName, ProcessSwapGroup]:
        if process.name not in groups:
            groups[process.name] = ProcessSwapGroup(process.name)
        groups[process.name].processes.append(process)
        return groups

    swap_used_groups: list[ProcessSwapGroup] = sorted(
        (item for item in reduce(_add_to_group, swap_used, dict()).values()),
        key=lambda x: x.total_swap
    )

    for group in swap_used_groups:
        print(f"{group.name:<32} {convert_bytes_to_mb(group.total_swap):10.2f}")
        for item in group.processes:
            print(f"{item.pid:>32} {convert_bytes_to_mb(item.swap):10.2f}")
        print()

    print("-" * 45)
    print(f"{'Total swap used by processes:':<32} {convert_bytes_to_mb(total_swap_used):10.2f} MB")

def show_swap_summary() -> None:
    swap = psutil.swap_memory()
    print("\n=== Swap Summary ===")
    print(f"Total: {convert_bytes_to_mb(swap.total):.2f} MB")
    print(f"Used:  {convert_bytes_to_mb(swap.used):.2f} MB")
    print(f"Free:  {convert_bytes_to_mb(swap.free):.2f} MB")
    print(f"Percent Used: {swap.percent:.1f}%")

def main() -> None:
    print("=== Process Swap Usage ===")
    get_process_swap_usage()
    show_swap_summary()

if __name__ == "__main__":
    main()