// test_ane_advanced.m — Probe advanced ANE interfaces
// SharedEvents, weightsBuffer, procedureIndex, VirtualClient, ChainingRequest
#import <Foundation/Foundation.h>
#import <objc/runtime.h>
#import <objc/message.h>
#import <dlfcn.h>
#import <IOSurface/IOSurface.h>
#import <mach/mach_time.h>

static mach_timebase_info_data_t g_tb;
static double tb_ms(uint64_t t) { return (double)t * g_tb.numer / g_tb.denom / 1e6; }

static void dump_class(const char *name) {
    Class cls = NSClassFromString([NSString stringWithUTF8String:name]);
    if (!cls) { printf("  %s: NOT FOUND\n", name); return; }
    printf("\n=== %s ===\n", name);
    unsigned int count;
    Method *methods = class_copyMethodList(object_getClass(cls), &count);
    if (count) printf("  Class methods:\n");
    for (unsigned int i = 0; i < count; i++) {
        SEL s = method_getName(methods[i]);
        const char *enc = method_getTypeEncoding(methods[i]);
        printf("    + %s  [%s]\n", sel_getName(s), enc ? enc : "?");
    }
    free(methods);
    methods = class_copyMethodList(cls, &count);
    if (count) printf("  Instance methods:\n");
    for (unsigned int i = 0; i < count; i++) {
        SEL s = method_getName(methods[i]);
        const char *enc = method_getTypeEncoding(methods[i]);
        printf("    - %s  [%s]\n", sel_getName(s), enc ? enc : "?");
    }
    free(methods);
    unsigned int pcount;
    objc_property_t *props = class_copyPropertyList(cls, &pcount);
    if (pcount) printf("  Properties:\n");
    for (unsigned int i = 0; i < pcount; i++) {
        const char *pname = property_getName(props[i]);
        const char *pattr = property_getAttributes(props[i]);
        printf("    @property %s  [%s]\n", pname, pattr ? pattr : "?");
    }
    free(props);
}

static IOSurfaceRef make_surface(size_t bytes) {
    return IOSurfaceCreate((__bridge CFDictionaryRef)@{
        (id)kIOSurfaceWidth:@(bytes), (id)kIOSurfaceHeight:@1,
        (id)kIOSurfaceBytesPerElement:@1, (id)kIOSurfaceBytesPerRow:@(bytes),
        (id)kIOSurfaceAllocSize:@(bytes), (id)kIOSurfacePixelFormat:@0});
}

int main() {
    @autoreleasepool {
        setbuf(stdout, NULL);
        mach_timebase_info(&g_tb);
        dlopen("/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/AppleNeuralEngine", RTLD_NOW);

        printf("=== ANE Advanced Interface Probe ===\n");

        // === Part 1: Event/Sync classes ===
        printf("\n--- Part 1: Event/Sync Classes ---\n");
        dump_class("_ANESharedEvents");
        dump_class("_ANESharedSignalEvent");
        dump_class("_ANESharedWaitEvent");
        dump_class("_ANEEvent");
        dump_class("_ANEFenceEvent");

        // Try instantiate
        const char *event_classes[] = {
            "_ANESharedEvents", "_ANESharedSignalEvent", "_ANESharedWaitEvent",
            "_ANEEvent", "_ANEFenceEvent", NULL
        };
        for (int i = 0; event_classes[i]; i++) {
            Class cls = NSClassFromString([NSString stringWithUTF8String:event_classes[i]]);
            if (!cls) continue;
            @try {
                id obj = [[cls alloc] init];
                printf("  %s alloc/init: %s\n", event_classes[i],
                       obj ? [[obj description] UTF8String] : "nil");
            } @catch (NSException *ex) {
                printf("  %s alloc/init: EXCEPTION: %s\n", event_classes[i], [[ex reason] UTF8String]);
            }
        }

        // === Part 2: VirtualClient and ChainingRequest ===
        printf("\n--- Part 2: VirtualClient / ChainingRequest ---\n");
        dump_class("_ANEVirtualClient");
        dump_class("_ANEChainingRequest");
        dump_class("_ANEMultiRequest");
        dump_class("_ANEBatchRequest");

        // === Part 3: weightsBuffer parameter test ===
        printf("\n--- Part 3: weightsBuffer IOSurface test ---\n");
        Class g_D  = NSClassFromString(@"_ANEInMemoryModelDescriptor");
        Class g_I  = NSClassFromString(@"_ANEInMemoryModel");
        Class g_AR = NSClassFromString(@"_ANERequest");
        Class g_AIO= NSClassFromString(@"_ANEIOSurfaceObject");

        int IC = 4, OC = 4, SP = 4;
        _Float16 weights[16];
        for (int i = 0; i < 16; i++) weights[i] = (i/4 == i%4) ? (_Float16)1.0f : (_Float16)0.0f;

        int ws = 16*2, tot = 128+ws;
        uint8_t *blob = (uint8_t*)calloc(tot,1);
        blob[0]=1; blob[4]=2; blob[64]=0xEF; blob[65]=0xBE; blob[66]=0xAD; blob[67]=0xDE; blob[68]=1;
        *(uint32_t*)(blob+72)=ws; *(uint32_t*)(blob+80)=128;
        memcpy(blob+128, weights, ws);
        NSData *wdata = [NSData dataWithBytesNoCopy:blob length:tot freeWhenDone:YES];

        NSString *mil = [NSString stringWithFormat:
            @"program(1.3)\n"
            "[buildInfo = dict<string, string>({{\"coremlc-component-MIL\", \"3510.2.1\"}, "
            "{\"coremlc-version\", \"3505.4.1\"}, {\"coremltools-component-milinternal\", \"\"}, "
            "{\"coremltools-version\", \"9.0\"}})]\n"
            "{\n"
            "    func main<ios18>(tensor<fp16, [1, %d, 1, %d]> x) {\n"
            "        string pt = const()[name=string(\"pt\"), val=string(\"valid\")];\n"
            "        tensor<int32, [2]> st = const()[name=string(\"st\"), val=tensor<int32, [2]>([1,1])];\n"
            "        tensor<int32, [4]> pd = const()[name=string(\"pd\"), val=tensor<int32, [4]>([0,0,0,0])];\n"
            "        tensor<int32, [2]> dl = const()[name=string(\"dl\"), val=tensor<int32, [2]>([1,1])];\n"
            "        int32 gr = const()[name=string(\"gr\"), val=int32(1)];\n"
            "        tensor<fp16, [%d,%d,1,1]> W = const()[name=string(\"W\"), "
            "val=tensor<fp16, [%d,%d,1,1]>(BLOBFILE(path=string(\"@model_path/weights/weight.bin\"), offset=uint64(64)))];\n"
            "        tensor<fp16, [1,%d,1,%d]> y = conv(dilations=dl,groups=gr,pad=pd,pad_type=pt,strides=st,weight=W,x=x)"
            "[name=string(\"conv\")];\n"
            "    } -> (y);\n"
            "}\n", IC, SP, OC, IC, OC, IC, OC, SP];

        NSData *md = [mil dataUsingEncoding:NSUTF8StringEncoding];
        id desc = ((id(*)(Class,SEL,id,id,id))objc_msgSend)(g_D, @selector(modelWithMILText:weights:optionsPlist:),
            md, @{@"@model_path/weights/weight.bin": @{@"offset":@0, @"data":wdata}}, nil);
        id mdl = ((id(*)(Class,SEL,id))objc_msgSend)(g_I, @selector(inMemoryModelWithDescriptor:), desc);
        id hx = ((id(*)(id,SEL))objc_msgSend)(mdl, @selector(hexStringIdentifier));
        NSString *td = [NSTemporaryDirectory() stringByAppendingPathComponent:hx];
        NSFileManager *fm = [NSFileManager defaultManager];
        [fm createDirectoryAtPath:[td stringByAppendingPathComponent:@"weights"]
            withIntermediateDirectories:YES attributes:nil error:nil];
        [md writeToFile:[td stringByAppendingPathComponent:@"model.mil"] atomically:YES];
        [wdata writeToFile:[td stringByAppendingPathComponent:@"weights/weight.bin"] atomically:YES];

        NSError *e = nil;
        ((BOOL(*)(id,SEL,unsigned int,id,NSError**))objc_msgSend)(mdl, @selector(compileWithQoS:options:error:), 21, @{}, &e);
        ((BOOL(*)(id,SEL,unsigned int,id,NSError**))objc_msgSend)(mdl, @selector(loadWithQoS:options:error:), 21, @{}, &e);

        IOSurfaceRef ioIn = make_surface(IC*SP*2);
        IOSurfaceRef ioOut = make_surface(OC*SP*2);

        // Write input
        IOSurfaceLock(ioIn, 0, NULL);
        _Float16 *inp = (_Float16*)IOSurfaceGetBaseAddress(ioIn);
        for (int c = 0; c < IC; c++) for (int s = 0; s < SP; s++) inp[c*SP+s] = (_Float16)(s+1.0f);
        IOSurfaceUnlock(ioIn, 0, NULL);

        // Normal eval first (baseline)
        id wI = ((id(*)(Class,SEL,IOSurfaceRef))objc_msgSend)(g_AIO, @selector(objectWithIOSurface:), ioIn);
        id wO = ((id(*)(Class,SEL,IOSurfaceRef))objc_msgSend)(g_AIO, @selector(objectWithIOSurface:), ioOut);
        id req0 = ((id(*)(Class,SEL,id,id,id,id,id,id,id))objc_msgSend)(g_AR,
            @selector(requestWithInputs:inputIndices:outputs:outputIndices:weightsBuffer:perfStats:procedureIndex:),
            @[wI], @[@0], @[wO], @[@0], nil, nil, @0);
        BOOL ok = ((BOOL(*)(id,SEL,unsigned int,id,id,NSError**))objc_msgSend)(
            mdl, @selector(evaluateWithQoS:options:request:error:), 21, @{}, req0, &e);
        printf("  Baseline eval (weightsBuffer=nil, procIdx=0): %s\n", ok ? "OK" : "FAIL");

        IOSurfaceLock(ioOut, kIOSurfaceLockReadOnly, NULL);
        _Float16 *out0 = (_Float16*)IOSurfaceGetBaseAddress(ioOut);
        printf("  Output: [%.1f, %.1f, %.1f, %.1f, ...]\n",
               (float)out0[0], (float)out0[1], (float)out0[2], (float)out0[3]);
        IOSurfaceUnlock(ioOut, kIOSurfaceLockReadOnly, NULL);

        // Test weightsBuffer: create IOSurface with weight data
        printf("\n  Testing weightsBuffer IOSurface...\n");
        _Float16 weights2[16];
        for (int i = 0; i < 16; i++) weights2[i] = (i/4 == i%4) ? (_Float16)3.0f : (_Float16)0.0f;

        IOSurfaceRef ioW = make_surface(ws);
        IOSurfaceLock(ioW, 0, NULL);
        memcpy(IOSurfaceGetBaseAddress(ioW), weights2, ws);
        IOSurfaceUnlock(ioW, 0, NULL);
        id wW = ((id(*)(Class,SEL,IOSurfaceRef))objc_msgSend)(g_AIO, @selector(objectWithIOSurface:), ioW);

        // Try with weightsBuffer
        wI = ((id(*)(Class,SEL,IOSurfaceRef))objc_msgSend)(g_AIO, @selector(objectWithIOSurface:), ioIn);
        wO = ((id(*)(Class,SEL,IOSurfaceRef))objc_msgSend)(g_AIO, @selector(objectWithIOSurface:), ioOut);
        id req_wb = ((id(*)(Class,SEL,id,id,id,id,id,id,id))objc_msgSend)(g_AR,
            @selector(requestWithInputs:inputIndices:outputs:outputIndices:weightsBuffer:perfStats:procedureIndex:),
            @[wI], @[@0], @[wO], @[@0], wW, nil, @0);
        printf("  Request with weightsBuffer: %s\n", req_wb ? "created" : "nil");

        if (req_wb) {
            ok = ((BOOL(*)(id,SEL,unsigned int,id,id,NSError**))objc_msgSend)(
                mdl, @selector(evaluateWithQoS:options:request:error:), 21, @{}, req_wb, &e);
            printf("  Eval with weightsBuffer: %s\n", ok ? "OK" : [[e description] UTF8String]);
            if (ok) {
                IOSurfaceLock(ioOut, kIOSurfaceLockReadOnly, NULL);
                _Float16 *outW = (_Float16*)IOSurfaceGetBaseAddress(ioOut);
                printf("  Output (3x identity via weightsBuffer): [%.1f, %.1f, %.1f, %.1f, ...]\n",
                       (float)outW[0], (float)outW[1], (float)outW[2], (float)outW[3]);
                bool is_3x = fabsf((float)outW[0] - 3.0f) < 0.1f;
                printf("  weightsBuffer override %s\n", is_3x ? "WORKS!" : "does NOT work (output unchanged)");
                IOSurfaceUnlock(ioOut, kIOSurfaceLockReadOnly, NULL);
            }
        }
        CFRelease(ioW);

        // === Part 4: procedureIndex sweep ===
        printf("\n--- Part 4: procedureIndex sweep (0-15) ---\n");
        for (int pi = 0; pi < 16; pi++) {
            wI = ((id(*)(Class,SEL,IOSurfaceRef))objc_msgSend)(g_AIO, @selector(objectWithIOSurface:), ioIn);
            wO = ((id(*)(Class,SEL,IOSurfaceRef))objc_msgSend)(g_AIO, @selector(objectWithIOSurface:), ioOut);
            id req_p = ((id(*)(Class,SEL,id,id,id,id,id,id,id))objc_msgSend)(g_AR,
                @selector(requestWithInputs:inputIndices:outputs:outputIndices:weightsBuffer:perfStats:procedureIndex:),
                @[wI], @[@0], @[wO], @[@0], nil, nil, @(pi));
            if (!req_p) { printf("  procIdx %2d: request=nil\n", pi); continue; }
            ok = ((BOOL(*)(id,SEL,unsigned int,id,id,NSError**))objc_msgSend)(
                mdl, @selector(evaluateWithQoS:options:request:error:), 21, @{}, req_p, &e);
            printf("  procIdx %2d: %s%s\n", pi, ok ? "OK" : "FAIL",
                   !ok && e ? [NSString stringWithFormat:@" (%@)", [e localizedDescription]].UTF8String : "");
        }

        // === Part 5: Scan all ANE classes ===
        printf("\n--- Part 5: All ANE-prefixed classes ---\n");
        unsigned int classCount;
        Class *allClasses = objc_copyClassList(&classCount);
        for (unsigned int i = 0; i < classCount; i++) {
            const char *name = class_getName(allClasses[i]);
            if (strstr(name, "ANE") || strstr(name, "ane")) {
                printf("  %s\n", name);
            }
        }
        free(allClasses);

        // Cleanup
        ((BOOL(*)(id,SEL,unsigned int,NSError**))objc_msgSend)(mdl, @selector(unloadWithQoS:error:), 21, &e);
        [fm removeItemAtPath:td error:nil];
        CFRelease(ioIn); CFRelease(ioOut);

        printf("\nDone.\n");
    }
    return 0;
}