// mtprobe: mimics a multicore DAW engine. The IO thread wakes N worker threads each cycle, each
// burns a slice of the budget, and the IO thread waits for all of them before the deadline.
// mode 0: workers are ordinary threads (not in the audio workgroup)
// mode 1: workers join the device IO workgroup (what a well-behaved DAW does)
// Output is SILENCE. Counts cycles where workers did not finish inside the budget.
// usage: mtprobe "<device>" <frames> <seconds> <workers> <load per worker 0..1> <mode>
#include <CoreAudio/CoreAudio.h>
#include <mach/mach_time.h>
#include <os/workgroup.h>
#include <pthread.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <math.h>
#include <mach/mach.h>
#include <mach/semaphore.h>
static semaphore_t g_sem;

#define MAXW 16
static int g_nw = 4, g_mode = 0;
static double g_load = 0.5, g_tickToNs = 1, g_budgetNs = 0;
static UInt32 g_frames = 64;
static _Atomic long g_gen = 0, g_done = 0;
static volatile int g_run = 1;
static long g_cycles = 0, g_miss = 0, g_overloads = 0;
static double g_worstMs = 0;
static os_workgroup_t g_wg = NULL;

static void burn(double ns) {
    uint64_t until = mach_absolute_time() + (uint64_t)(ns / g_tickToNs);
    volatile double x = 1.0;
    while (mach_absolute_time() < until) x = sin(x) + 1.0;
}

static void *worker(void *arg) {
    os_workgroup_join_token_s tok;
    int joined = 0;
    if (g_mode == 1 && g_wg) joined = (os_workgroup_join(g_wg, &tok) == 0);
    while (g_run) {
        semaphore_wait(g_sem);
        if (!g_run) break;
        burn(g_budgetNs * g_load);
        atomic_fetch_add(&g_done, 1);
    }
    if (joined) os_workgroup_leave(g_wg, &tok);
    return NULL;
}

static OSStatus ol(AudioObjectID o, UInt32 n, const AudioObjectPropertyAddress *a, void *c) { g_overloads++; return 0; }

static OSStatus ioproc(AudioObjectID dev, const AudioTimeStamp *now, const AudioBufferList *in,
                       const AudioTimeStamp *inTime, AudioBufferList *out, const AudioTimeStamp *outTime, void *ctx) {
    for (UInt32 i = 0; i < out->mNumberBuffers; i++) memset(out->mBuffers[i].mData, 0, out->mBuffers[i].mDataByteSize);
    uint64_t t0 = mach_absolute_time();
    atomic_store(&g_done, 0);
    for (int i = 0; i < g_nw; i++) semaphore_signal(g_sem);
    uint64_t deadline = t0 + (uint64_t)(g_budgetNs * 0.95 / g_tickToNs);
    while (atomic_load(&g_done) < g_nw && mach_absolute_time() < deadline) { }
    double ms = (mach_absolute_time() - t0) * g_tickToNs / 1e6;
    if (ms > g_worstMs) g_worstMs = ms;
    if (atomic_load(&g_done) < g_nw) g_miss++;
    g_cycles++;
    return 0;
}

int main(int argc, char **argv) {
    if (argc < 7) { fprintf(stderr, "usage\n"); return 2; }
    mach_timebase_info_data_t tb; mach_timebase_info(&tb); g_tickToNs = (double)tb.numer / tb.denom;
    g_frames = atoi(argv[2]); int secs = atoi(argv[3]); g_nw = atoi(argv[4]); g_load = atof(argv[5]); g_mode = atoi(argv[6]);
    if (g_nw > MAXW) g_nw = MAXW;
    AudioObjectPropertyAddress pa = {kAudioHardwarePropertyDevices, kAudioObjectPropertyScopeGlobal, kAudioObjectPropertyElementMain};
    UInt32 sz = 0; AudioObjectGetPropertyDataSize(kAudioObjectSystemObject, &pa, 0, NULL, &sz);
    int n = sz / sizeof(AudioObjectID); AudioObjectID *ids = malloc(sz);
    AudioObjectGetPropertyData(kAudioObjectSystemObject, &pa, 0, NULL, &sz, ids);
    AudioObjectID dev = 0;
    for (int i = 0; i < n; i++) {
        CFStringRef name = NULL; UInt32 s2 = sizeof(name);
        AudioObjectPropertyAddress na = {kAudioObjectPropertyName, kAudioObjectPropertyScopeGlobal, kAudioObjectPropertyElementMain};
        if (AudioObjectGetPropertyData(ids[i], &na, 0, NULL, &s2, &name) == 0 && name) {
            char buf[256]; CFStringGetCString(name, buf, sizeof buf, kCFStringEncodingUTF8); CFRelease(name);
            if (strstr(buf, argv[1])) { dev = ids[i]; break; }
        }
    }
    if (!dev) { fprintf(stderr, "device not found\n"); return 1; }
    AudioObjectPropertyAddress bf = {kAudioDevicePropertyBufferFrameSize, kAudioObjectPropertyScopeGlobal, kAudioObjectPropertyElementMain};
    UInt32 fr = g_frames; AudioObjectSetPropertyData(dev, &bf, 0, NULL, sizeof fr, &fr);
    sz = sizeof g_frames; AudioObjectGetPropertyData(dev, &bf, 0, NULL, &sz, &g_frames);
    g_budgetNs = g_frames / 48000.0 * 1e9;
    AudioObjectPropertyAddress wa = {kAudioDevicePropertyIOThreadOSWorkgroup, kAudioObjectPropertyScopeGlobal, kAudioObjectPropertyElementMain};
    sz = sizeof g_wg;
    OSStatus we = AudioObjectGetPropertyData(dev, &wa, 0, NULL, &sz, &g_wg);
    AudioObjectPropertyAddress oa = {kAudioDeviceProcessorOverload, kAudioObjectPropertyScopeGlobal, kAudioObjectPropertyElementMain};
    AudioObjectAddPropertyListener(dev, &oa, ol, NULL);
    semaphore_create(mach_task_self(), &g_sem, SYNC_POLICY_FIFO, 0);
    pthread_t th[MAXW];
    for (int i = 0; i < g_nw; i++) pthread_create(&th[i], NULL, worker, NULL);
    AudioDeviceIOProcID pid = NULL;
    AudioDeviceCreateIOProcID(dev, ioproc, NULL, &pid);
    AudioDeviceStart(dev, pid);
    for (int i = 0; i < secs; i++) sleep(1);
    AudioDeviceStop(dev, pid);
    g_run = 0;
    for (int i = 0; i < g_nw; i++) semaphore_signal(g_sem);
    for (int i = 0; i < g_nw; i++) pthread_join(th[i], NULL);
    printf("RESULT frames=%u secs=%d workers=%d load=%.2f mode=%s workgroup=%s cycles=%ld missed=%ld (%.3f%%) overloads=%ld worst_ms=%.3f budget_ms=%.3f\n",
           g_frames, secs, g_nw, g_load, g_mode ? "workgroup" : "plain", we == 0 && g_wg ? "yes" : "no",
           g_cycles, g_miss, g_cycles ? 100.0 * g_miss / g_cycles : 0, g_overloads, g_worstMs, g_budgetNs / 1e6);
    return 0;
}
