/* Targeted macOS bench harness. Never selects a system audio output or unmutes a profile. */
#include "nxs_protocol.h"
#include <CoreAudio/CoreAudio.h>
#include <CoreFoundation/CoreFoundation.h>
#include <IOKit/hid/IOHIDManager.h>
#include <IOKit/hid/IOHIDKeys.h>
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define REQUIRE(condition) do { if (!(condition)) { fprintf(stderr, "Failed at line %d: %s\n", __LINE__, #condition); exit(1); } } while (0)

static uint32_t next_request_id;

static uint64_t microseconds(void) {
    struct timespec now;
    REQUIRE(clock_gettime(CLOCK_MONOTONIC, &now) == 0);
    return (uint64_t)now.tv_sec * 1000000u + (uint64_t)now.tv_nsec / 1000u;
}

static void audio_defaults(void) {
    const AudioObjectPropertySelector selectors[] = {
        kAudioHardwarePropertyDefaultOutputDevice, kAudioHardwarePropertyDefaultSystemOutputDevice};
    AudioDeviceID devices[2];
    for (size_t index = 0; index < 2; ++index) {
        AudioObjectPropertyAddress property = {selectors[index], kAudioObjectPropertyScopeGlobal, kAudioObjectPropertyElementMain};
        UInt32 size = sizeof(devices[index]);
        REQUIRE(AudioObjectGetPropertyData(kAudioObjectSystemObject, &property, 0, NULL, &size, &devices[index]) == noErr);
    }
    printf("{\"kind\":\"audio_defaults\",\"output_device_id\":%u,\"system_output_device_id\":%u}\n", devices[0], devices[1]);
    for (size_t index = 0; index < 2; ++index) {
        const AudioObjectPropertySelector properties[] = {kAudioDevicePropertyDeviceUID, kAudioObjectPropertyName};
        char strings[2][512];
        for (size_t field = 0; field < 2; ++field) {
            AudioObjectPropertyAddress property = {properties[field], kAudioObjectPropertyScopeGlobal, kAudioObjectPropertyElementMain};
            CFStringRef value;
            UInt32 size = sizeof(value);
            REQUIRE(AudioObjectGetPropertyData(devices[index], &property, 0, NULL, &size, &value) == noErr);
            REQUIRE(CFStringGetCString(value, strings[field], sizeof(strings[field]), kCFStringEncodingUTF8));
            CFRelease(value);
        }
        printf("{\"kind\":\"audio_default_identity\",\"role\":%zu,\"uid\":\"%s\",\"name\":\"%s\"}\n", index, strings[0], strings[1]);
    }
}

static void match_number(CFMutableDictionaryRef matching, CFStringRef key, int value) {
    CFNumberRef number = CFNumberCreate(kCFAllocatorDefault, kCFNumberIntType, &value);
    REQUIRE(number != NULL);
    CFDictionarySetValue(matching, key, number);
    CFRelease(number);
}

static IOHIDDeviceRef selected_device(IOHIDManagerRef manager, const char *serial, CFSetRef *retained_devices) {
    REQUIRE(strlen(serial) == 16);
    for (size_t index = 0; index < 16; ++index) REQUIRE(isxdigit((unsigned char)serial[index]));
    CFMutableDictionaryRef matching = CFDictionaryCreateMutable(kCFAllocatorDefault, 0,
        &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);
    REQUIRE(matching != NULL);
    match_number(matching, CFSTR(kIOHIDVendorIDKey), NXS_USB_VENDOR_ID);
    match_number(matching, CFSTR(kIOHIDProductIDKey), NXS_USB_PRODUCT_ID);
    match_number(matching, CFSTR(kIOHIDPrimaryUsagePageKey), NXS_HID_USAGE_PAGE);
    match_number(matching, CFSTR(kIOHIDPrimaryUsageKey), NXS_HID_USAGE);
    CFStringRef identifier = CFStringCreateWithCString(kCFAllocatorDefault, serial, kCFStringEncodingUTF8);
    REQUIRE(identifier != NULL);
    CFDictionarySetValue(matching, CFSTR(kIOHIDSerialNumberKey), identifier);
    IOHIDManagerSetDeviceMatching(manager, matching);
    CFRelease(identifier);
    CFRelease(matching);
    REQUIRE(IOHIDManagerOpen(manager, kIOHIDOptionsTypeNone) == kIOReturnSuccess);
    *retained_devices = IOHIDManagerCopyDevices(manager);
    REQUIRE(*retained_devices != NULL && CFSetGetCount(*retained_devices) == 1);
    const void *values[1];
    CFSetGetValues(*retained_devices, values);
    return (IOHIDDeviceRef)values[0];
}

static nxs_message_t prepare(nxs_operation_t operation, uint32_t generation, const nxs_snapshot_t *snapshot) {
    nxs_message_t request = {.operation = operation, .request_id = next_request_id++, .generation = generation};
    REQUIRE(request.request_id != 0);
    if (snapshot) {
        request.payload_length = NXS_SNAPSHOT_SIZE;
        REQUIRE(nxs_snapshot_encode(snapshot, request.payload, request.payload_length) == NXS_RESULT_OK);
    }
    return request;
}

static nxs_message_t exchange(IOHIDDeviceRef device, const nxs_message_t *request, nxs_result_t expected) {
    uint8_t report[NXS_REPORT_SIZE];
    REQUIRE(nxs_message_encode(request, report, sizeof(report)) == NXS_RESULT_OK);
    uint64_t started = microseconds();
    IOReturn sent = IOHIDDeviceSetReport(device, kIOHIDReportTypeFeature, NXS_REPORT_ID, report, sizeof(report));
    if (sent != kIOReturnSuccess) fprintf(stderr, "SET_REPORT: 0x%08x\n", (unsigned int)sent);
    REQUIRE(sent == kIOReturnSuccess);
    IOReturn received = kIOReturnError;
    CFIndex length = 0;
    do {
        length = sizeof(report);
        memset(report, 0, sizeof(report));
        received = IOHIDDeviceGetReport(device, kIOHIDReportTypeFeature, NXS_REPORT_ID, report, &length);
        nxs_message_t reply;
        if (received == kIOReturnSuccess && nxs_message_decode(report, (size_t)length, &reply) == NXS_RESULT_OK &&
            reply.request_id == request->request_id && reply.operation == request->operation && (reply.flags & NXS_FLAG_RESPONSE)) {
            printf("{\"kind\":\"reply\",\"operation\":%u,\"request_id\":%u,\"result\":%u,\"generation\":%u,\"elapsed_us\":%llu,\"report_hex\":\"",
                (unsigned int)reply.operation, reply.request_id, (unsigned int)reply.result, reply.generation,
                (unsigned long long)(microseconds() - started));
            for (size_t index = 0; index < sizeof(report); ++index) printf("%02x", report[index]);
            puts("\"}");
            fflush(stdout);
            REQUIRE(reply.result == expected);
            return reply;
        }
        struct timespec delay = {.tv_nsec = 1000000};
        nanosleep(&delay, NULL);
    } while (microseconds() - started < 3000000u);
    fprintf(stderr, "Timed out waiting for request %u; last GET_REPORT 0x%08x, length %ld\n",
        request->request_id, (unsigned int)received, (long)length);
    exit(1);
}

static nxs_message_t query(IOHIDDeviceRef device, nxs_operation_t operation) {
    nxs_message_t request = prepare(operation, 0, NULL);
    return exchange(device, &request, NXS_RESULT_OK);
}

static nxs_runtime_status_t status(IOHIDDeviceRef device) {
    nxs_message_t reply = query(device, NXS_OPERATION_STATUS);
    nxs_runtime_status_t value;
    REQUIRE(nxs_runtime_status_decode(reply.payload, reply.payload_length, &value) == NXS_RESULT_OK);
    printf("{\"kind\":\"status\",\"persisted_generation\":%u,\"storage_available\":%u,\"streaming\":%u,\"input_peak\":%.9g,\"output_peak\":%.9g,\"underruns\":%u,\"overruns\":%u,\"max_processing_us\":%u,\"deadline_us\":%u}\n",
        value.persisted_generation, value.storage_available, value.streaming, value.input_peak, value.output_peak,
        value.underrun_count, value.overrun_count, value.max_processing_us, value.block_deadline_us);
    return value;
}

static nxs_snapshot_t maximum_muted(nxs_snapshot_t snapshot) {
    static const float frequencies[] = {20, 100, 1000, 10000, 20000};
    snapshot.high_pass_enabled = 1;
    snapshot.high_pass_hz = nxs_parameter_range(NXS_PARAMETER_HIGH_PASS_HZ).maximum;
    snapshot.attenuation_db = nxs_parameter_range(NXS_PARAMETER_ATTENUATION_DB).maximum;
    snapshot.limiter_threshold_db = nxs_parameter_range(NXS_PARAMETER_LIMITER_THRESHOLD_DB).minimum;
    snapshot.limiter_release_ms = nxs_parameter_range(NXS_PARAMETER_LIMITER_RELEASE_MS).maximum;
    snapshot.muted = 1;
    for (size_t band = 0; band < NXS_EQ_BAND_COUNT; ++band) {
        snapshot.bands[band].frequency_hz = frequencies[band];
        snapshot.bands[band].gain_db = nxs_parameter_range(NXS_PARAMETER_EQ_GAIN_DB).maximum;
        snapshot.bands[band].quality = nxs_parameter_range(NXS_PARAMETER_EQ_QUALITY).maximum;
        snapshot.bands[band].enabled = 1;
    }
    REQUIRE(nxs_snapshot_is_valid(&snapshot));
    return snapshot;
}

static void retain_baseline(const char *serial, const nxs_snapshot_t *snapshot) {
    const char *path = getenv("NXS_BENCH_BASELINE");
    REQUIRE(path != NULL && snapshot->muted == 1);
    uint8_t payload[NXS_SNAPSHOT_SIZE];
    REQUIRE(nxs_snapshot_encode(snapshot, payload, sizeof(payload)) == NXS_RESULT_OK);
    FILE *file = fopen(path, "wbx");
    REQUIRE(file != NULL);
    REQUIRE(fwrite("NXSB", 1, 4, file) == 4 && fwrite(serial, 1, 16, file) == 16);
    REQUIRE(fwrite(payload, 1, sizeof(payload), file) == sizeof(payload));
    REQUIRE(fclose(file) == 0);
}

static nxs_snapshot_t restore_baseline(const char *serial, const nxs_snapshot_t *current) {
    const char *path = getenv("NXS_BENCH_BASELINE");
    REQUIRE(path != NULL);
    FILE *file = fopen(path, "rb");
    REQUIRE(file != NULL);
    char header[20];
    uint8_t payload[NXS_SNAPSHOT_SIZE];
    REQUIRE(fread(header, 1, sizeof(header), file) == sizeof(header));
    REQUIRE(!memcmp(header, "NXSB", 4) && !memcmp(header + 4, serial, 16));
    REQUIRE(fread(payload, 1, sizeof(payload), file) == sizeof(payload) && fgetc(file) == EOF);
    REQUIRE(fclose(file) == 0);
    nxs_snapshot_t baseline;
    REQUIRE(nxs_snapshot_decode(payload, sizeof(payload), &baseline) == NXS_RESULT_OK && baseline.muted == 1);
    nxs_snapshot_t expected = maximum_muted(baseline);
    uint8_t expected_bytes[NXS_SNAPSHOT_SIZE], current_bytes[NXS_SNAPSHOT_SIZE];
    REQUIRE(nxs_snapshot_encode(&expected, expected_bytes, sizeof(expected_bytes)) == NXS_RESULT_OK);
    REQUIRE(nxs_snapshot_encode(current, current_bytes, sizeof(current_bytes)) == NXS_RESULT_OK);
    REQUIRE(!memcmp(expected_bytes, current_bytes, sizeof(expected_bytes)));
    return baseline;
}

int main(int count, char **arguments) {
    if (count == 2 && !strcmp(arguments[1], "audio-defaults")) {
        audio_defaults();
        return 0;
    }
    if (count < 3 || count > 4 || (strcmp(arguments[2], "inspect") && strcmp(arguments[2], "save-muted") &&
        strcmp(arguments[2], "stress-muted") && strcmp(arguments[2], "restore-muted"))) {
        fprintf(stderr, "Usage: control-stress SERIAL inspect|stress-muted|restore-muted | SERIAL save-muted ATTENUATION_DB\n");
        return 2;
    }
    bool stress = !strcmp(arguments[2], "stress-muted"), restore = !strcmp(arguments[2], "restore-muted");
    bool attenuation = !strcmp(arguments[2], "save-muted");
    bool save = stress || restore || attenuation;
    REQUIRE((attenuation && count == 4) || (!attenuation && count == 3));
    next_request_id = arc4random_uniform(0x7fffffffu) + 1u;
    IOHIDManagerRef manager = IOHIDManagerCreate(kCFAllocatorDefault, kIOHIDOptionsTypeNone);
    REQUIRE(manager != NULL);
    CFSetRef devices;
    IOHIDDeviceRef device = selected_device(manager, arguments[1], &devices);
    nxs_message_t reply = query(device, NXS_OPERATION_CAPABILITIES);
    nxs_capabilities_t capabilities;
    REQUIRE(nxs_capabilities_decode(reply.payload, reply.payload_length, &capabilities) == NXS_RESULT_OK);
    REQUIRE(!strcmp(capabilities.device_serial, arguments[1]));
    REQUIRE(capabilities.system_clock_hz == NXS_SYSTEM_CLOCK_HZ && capabilities.flash_size_bytes == 4194304u);
    REQUIRE(capabilities.feature_flags == (NXS_FEATURE_PERSISTENCE | NXS_FEATURE_FLASH_VERIFIED));
    printf("{\"kind\":\"capabilities\",\"serial\":\"%s\",\"firmware\":\"%s\",\"clock_hz\":%u,\"flash_bytes\":%u,\"features\":%u}\n",
        capabilities.device_serial, capabilities.firmware_version, capabilities.system_clock_hz,
        capabilities.flash_size_bytes, capabilities.feature_flags);
    reply = query(device, NXS_OPERATION_READ_STATE);
    nxs_snapshot_t active;
    REQUIRE(nxs_snapshot_decode(reply.payload, reply.payload_length, &active) == NXS_RESULT_OK);
    printf("{\"kind\":\"state\",\"generation\":%u,\"muted\":%u,\"attenuation_db\":%.9g}\n", reply.generation, active.muted, active.attenuation_db);
    REQUIRE(active.muted == 1);
    uint32_t generation = reply.generation;
    nxs_runtime_status_t before = status(device);
    REQUIRE(before.storage_available == 1);
    if (save) {
        if (stress) {
            retain_baseline(arguments[1], &active);
            active = maximum_muted(active);
            puts("{\"kind\":\"stress_stage\",\"action\":\"apply_and_persist_maximum_muted\"}");
        } else if (restore) {
            active = restore_baseline(arguments[1], &active);
            puts("{\"kind\":\"stress_stage\",\"action\":\"restore_and_persist_baseline\"}");
        } else {
            char *end;
            active.attenuation_db = strtof(arguments[3], &end);
            REQUIRE(*end == '\0');
        }
        REQUIRE(nxs_snapshot_is_valid(&active) && active.muted == 1 && generation != UINT32_MAX);
        nxs_message_t apply = prepare(NXS_OPERATION_APPLY_STATE, generation, &active);
        reply = exchange(device, &apply, NXS_RESULT_OK);
        REQUIRE(reply.generation == generation + 1u);
        reply = exchange(device, &apply, NXS_RESULT_OK);
        REQUIRE(reply.generation == generation + 1u);
        reply = query(device, NXS_OPERATION_READ_STATE);
        REQUIRE(reply.generation == generation + 1u && reply.payload_length == apply.payload_length);
        REQUIRE(!memcmp(reply.payload, apply.payload, reply.payload_length));
        nxs_message_t stale = prepare(NXS_OPERATION_PERSIST_STATE, generation, NULL);
        exchange(device, &stale, NXS_RESULT_STALE_GENERATION);
        nxs_message_t persist = prepare(NXS_OPERATION_PERSIST_STATE, generation + 1u, NULL);
        reply = exchange(device, &persist, NXS_RESULT_OK);
        REQUIRE(reply.generation == generation + 1u);
        exchange(device, &persist, NXS_RESULT_OK);
        nxs_runtime_status_t after = status(device);
        REQUIRE(after.persisted_generation == generation + 1u && after.storage_available == 1);
        puts("{\"kind\":\"muted_save_verified\"}");
    }
    REQUIRE(IOHIDManagerClose(manager, kIOHIDOptionsTypeNone) == kIOReturnSuccess);
    CFRelease(devices);
    CFRelease(manager);
    return 0;
}