"""Enumerate IOReport Energy Model rails on Apple Silicon, unprivileged.

Decisive question: are the GPU and ANE energy rails readable without sudo on
macOS 27, and do they advance under a Metal workload?
"""
from ctypes import *

CF = CDLL("/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation")
IOR = CDLL("/usr/lib/libIOReport.dylib")

CFStringRef = c_void_p; CFDictionaryRef = c_void_p; CFArrayRef = c_void_p
CF.CFStringCreateWithCString.restype = CFStringRef
CF.CFStringCreateWithCString.argtypes = [c_void_p, c_char_p, c_uint32]
CF.CFStringGetCStringPtr.restype = c_char_p
CF.CFStringGetCStringPtr.argtypes = [CFStringRef, c_uint32]
CF.CFStringGetCString.restype = c_bool
CF.CFStringGetCString.argtypes = [CFStringRef, c_char_p, c_long, c_uint32]
CF.CFDictionaryGetValue.restype = c_void_p
CF.CFDictionaryGetValue.argtypes = [CFDictionaryRef, c_void_p]
CF.CFArrayGetCount.restype = c_long
CF.CFArrayGetCount.argtypes = [CFArrayRef]
CF.CFArrayGetValueAtIndex.restype = c_void_p
CF.CFArrayGetValueAtIndex.argtypes = [CFArrayRef, c_long]

def cfstr(s): return CF.CFStringCreateWithCString(None, s.encode(), 0x08000100)  # kCFStringEncodingUTF8
def pystr(ref):
    if not ref: return None
    p = CF.CFStringGetCStringPtr(ref, 0x08000100)
    if p: return p.decode(errors="replace")
    buf = create_string_buffer(512)
    return buf.value.decode(errors="replace") if CF.CFStringGetCString(ref, buf, 512, 0x08000100) else None

IOR.IOReportCopyChannelsInGroup.restype = CFDictionaryRef
IOR.IOReportCopyChannelsInGroup.argtypes = [CFStringRef, CFStringRef, c_uint64, c_uint64, c_uint64]
IOR.IOReportCreateSubscription.restype = c_void_p
IOR.IOReportCreateSubscription.argtypes = [c_void_p, CFDictionaryRef, POINTER(CFDictionaryRef), c_uint64, c_void_p]
IOR.IOReportCreateSamples.restype = CFDictionaryRef
IOR.IOReportCreateSamples.argtypes = [c_void_p, CFDictionaryRef, c_void_p]
IOR.IOReportCreateSamplesDelta.restype = CFDictionaryRef
IOR.IOReportCreateSamplesDelta.argtypes = [CFDictionaryRef, CFDictionaryRef, c_void_p]
for fn, rt in (("IOReportChannelGetChannelName", CFStringRef), ("IOReportChannelGetGroup", CFStringRef),
               ("IOReportChannelGetSubGroup", CFStringRef), ("IOReportChannelGetUnitLabel", CFStringRef)):
    getattr(IOR, fn).restype = rt; getattr(IOR, fn).argtypes = [CFDictionaryRef]
IOR.IOReportSimpleGetIntegerValue.restype = c_int64
IOR.IOReportSimpleGetIntegerValue.argtypes = [CFDictionaryRef, c_int32]

chans = IOR.IOReportCopyChannelsInGroup(cfstr("Energy Model"), None, 0, 0, 0)
if not chans: raise SystemExit("no Energy Model group")
subbed = CFDictionaryRef()
sub = IOR.IOReportCreateSubscription(None, chans, byref(subbed), 0, None)
if not sub: raise SystemExit("subscription refused (would mean privilege is required)")

def sample(): return IOR.IOReportCreateSamples(sub, subbed, None)
def rails(delta):
    out = []
    arr = CF.CFDictionaryGetValue(delta, cfstr("IOReportChannels"))
    for i in range(CF.CFArrayGetCount(arr)):
        ch = CF.CFArrayGetValueAtIndex(arr, i)
        out.append((pystr(IOR.IOReportChannelGetGroup(ch)), pystr(IOR.IOReportChannelGetSubGroup(ch)),
                    pystr(IOR.IOReportChannelGetChannelName(ch)), pystr(IOR.IOReportChannelGetUnitLabel(ch)),
                    IOR.IOReportSimpleGetIntegerValue(ch, 0)))
    return out

import time, subprocess, sys
s0 = sample()
# Load the GPU (and ANE if asked) while the rails are watched.
load = subprocess.Popen([sys.argv[1], sys.argv[2]]) if len(sys.argv) > 2 else None
time.sleep(2.0)
s1 = sample()
if load: load.wait()
d = IOR.IOReportCreateSamplesDelta(s0, s1, None)
rs = rails(d)
print(f"subscription: OK (unprivileged, uid={__import__('os').getuid()})   rails in Energy Model: {len(rs)}")
agg = {}
for g, sg, name, unit, val in rs:
    key = (sg or "-", name or "-", unit or "-")
    a = agg.setdefault(key, [0, 0]); a[0] += val; a[1] += 1
print(f"{'subgroup':<22}{'channel':<34}{'unit':<8}{'delta over 2s':>16}{'n':>5}")
for (sg, name, unit), (val, n) in sorted(agg.items(), key=lambda kv: -abs(kv[1][0]))[:28]:
    print(f"{sg:<22}{name:<34}{unit:<8}{val:>16}{n:>5}")