#!/usr/bin/env python3
import ctypes
import re
import ctypes.util
import subprocess
import json

# Shows Safari tabs with their memory usage, sorted largest descending.
#
# Enable Safari's debug menu with
#   defaults write com.apple.Safari IncludeInternalDebugMenu 1
# Enable showing pid in tab titles by choosing
#   Debug...Miscellaneous Flags...Show Web Process IDs in Page Titles
#
# License: MIT

def main():
    data = [(usage(pid), pid, title, url) for title, url, pid in tabs_and_pids()]
    data.sort(reverse=True)
    for mem, pid, title, url in data:
        print(human_readable_size(mem), pid, title, url)


TAB_TITLES_AND_URLS_SCRIPT = """
	function run() {
			const ret = [];
			const Safari = Application('Safari');
			const windows = Safari.windows();
			for (let windowIndex = 0; windowIndex < windows.length; windowIndex++) {
					const tabs = windows[windowIndex].tabs();
									if (tabs === null) {
										continue;
									}
					for (let tabIndex = 0; tabIndex < tabs.length; tabIndex++) {
							const currentTab = tabs[tabIndex];
							const tabName = currentTab.name();
							const tabURL = currentTab.url();
							ret.push([tabName, tabURL]);
					}
			}
			return JSON.stringify(ret);
	}
"""


def tabs_and_pids():
    result = subprocess.check_output(
        ["osascript", "-l", "JavaScript", "-e", TAB_TITLES_AND_URLS_SCRIPT],
        shell=False,
        text=True,
    )
    for title, url in json.loads(result):
        pid = re.search(r"\[WP\s*(\d+)\]", title).group(1)
        yield title, url, int(pid)


def human_readable_size(size_bytes):
    units = ["B", "KB", "MB", "GB"]
    size = float(size_bytes)
    unit_index = 0

    while size >= 1024 and unit_index < len(units) - 1:
        size /= 1024
        unit_index += 1

    return f"{size:.1f}{units[unit_index]}"


# See /Library/Developer/CommandLineTools/SDKs/MacOSX13.3.sdk/usr/include/libproc.h
# and System/Library/Frameworks/Kernel.framework/Versions/A/Headers/sys/resource.h
RUSAGE_INFO_V6 = 6


class rusage_info_v6(ctypes.Structure):
    _fields_ = [
        ("ri_uuid", ctypes.c_uint8 * 16),
        ("ri_user_time", ctypes.c_uint64),
        ("ri_system_time", ctypes.c_uint64),
        ("ri_pkg_idle_wkups", ctypes.c_uint64),
        ("ri_interrupt_wkups", ctypes.c_uint64),
        ("ri_pageins", ctypes.c_uint64),
        ("ri_wired_size", ctypes.c_uint64),
        ("ri_resident_size", ctypes.c_uint64),
        ("ri_phys_footprint", ctypes.c_uint64),
        ("ri_proc_start_abstime", ctypes.c_uint64),
        ("ri_proc_exit_abstime", ctypes.c_uint64),
        ("ri_child_user_time", ctypes.c_uint64),
        ("ri_child_system_time", ctypes.c_uint64),
        ("ri_child_pkg_idle_wkups", ctypes.c_uint64),
        ("ri_child_interrupt_wkups", ctypes.c_uint64),
        ("ri_child_pageins", ctypes.c_uint64),
        ("ri_child_elapsed_abstime", ctypes.c_uint64),
        ("ri_diskio_bytesread", ctypes.c_uint64),
        ("ri_diskio_byteswritten", ctypes.c_uint64),
        ("ri_cpu_time_qos_default", ctypes.c_uint64),
        ("ri_cpu_time_qos_maintenance", ctypes.c_uint64),
        ("ri_cpu_time_qos_background", ctypes.c_uint64),
        ("ri_cpu_time_qos_utility", ctypes.c_uint64),
        ("ri_cpu_time_qos_legacy", ctypes.c_uint64),
        ("ri_cpu_time_qos_user_initiated", ctypes.c_uint64),
        ("ri_cpu_time_qos_user_interactive", ctypes.c_uint64),
        ("ri_billed_system_time", ctypes.c_uint64),
        ("ri_serviced_system_time", ctypes.c_uint64),
        ("ri_logical_writes", ctypes.c_uint64),
        ("ri_lifetime_max_phys_footprint", ctypes.c_uint64),
        ("ri_instructions", ctypes.c_uint64),
        ("ri_cycles", ctypes.c_uint64),
        ("ri_billed_energy", ctypes.c_uint64),
        ("ri_serviced_energy", ctypes.c_uint64),
        ("ri_interval_max_phys_footprint", ctypes.c_uint64),
        ("ri_runnable_time", ctypes.c_uint64),
        ("ri_flags", ctypes.c_uint64),
        ("ri_user_ptime", ctypes.c_uint64),
        ("ri_system_ptime", ctypes.c_uint64),
        ("ri_pinstructions", ctypes.c_uint64),
        ("ri_pcycles", ctypes.c_uint64),
        ("ri_energy_nj", ctypes.c_uint64),
        ("ri_penergy_nj", ctypes.c_uint64),
        ("ri_reserved", ctypes.c_uint64 * 14),
    ]


proc_lib = ctypes.CDLL(ctypes.util.find_library("proc"))
proc_pid_rusage = proc_lib.proc_pid_rusage
proc_pid_rusage.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(rusage_info_v6)]
proc_pid_rusage.restype = ctypes.c_int


def usage(pid):
    us = rusage_info_v6()
    proc_pid_rusage(ctypes.c_int(pid), RUSAGE_INFO_V6, ctypes.byref(us))
    return us.ri_phys_footprint


if __name__ == "__main__":
    main()