#import <Foundation/Foundation.h>
#import <Metal/Metal.h>
#import <stdio.h>
#import <stdlib.h>
#import <string.h>

enum { kMaxHits = 256, kThreadsPerGroup = 256, kParamsSize = 104 };

typedef struct {
    uint64_t index;
    uint64_t key;
    uint64_t plain;
} Hit;

static void fail(const char *msg) {
    fprintf(stderr, "%s\n", msg);
    exit(2);
}

static uint64_t parse_hex_u64(const char *hex) {
    const char *p = hex;
    if (p[0] == '0' && (p[1] == 'x' || p[1] == 'X')) p += 2;
    return strtoull(p, NULL, 16);
}

static NSDictionary<NSString *, NSString *> *parse_opts(int argc, const char **argv) {
    NSMutableDictionary *opts = [NSMutableDictionary dictionary];
    for (int i = 2; i + 1 < argc; i += 2) {
        if (strncmp(argv[i], "--", 2) != 0) {
            i -= 1;
            continue;
        }
        opts[@(argv[i] + 2)] = @(argv[i + 1]);
    }
    return opts;
}

static void json_escape(const char *src, char *dst, size_t dstsz) {
    size_t o = 0;
    for (size_t i = 0; src[i] && o + 2 < dstsz; i++) {
        unsigned char c = (unsigned char)src[i];
        if (c == '"' || c == '\\') {
            dst[o++] = '\\';
            dst[o++] = (char)c;
        } else if (c < 32) {
            dst[o++] = '?';
        } else {
            dst[o++] = (char)c;
        }
    }
    dst[o] = 0;
}

static void write_params(void *dst, uint64_t start, uint32_t key_len, uint32_t charset_len,
                         uint32_t pad_byte, uint32_t fill_count, uint32_t batch_count,
                         uint64_t target, const uint64_t *fills, int nfills) {
    uint8_t *p = dst;
    memcpy(p, &start, 8); p += 8;
    memcpy(p, &key_len, 4); p += 4;
    memcpy(p, &charset_len, 4); p += 4;
    memcpy(p, &pad_byte, 4); p += 4;
    memcpy(p, &fill_count, 4); p += 4;
    memcpy(p, &batch_count, 4); p += 4;
    uint32_t z = 0;
    memcpy(p, &z, 4); p += 4;
    memcpy(p, &target, 8); p += 8;
    for (int i = 0; i < 8; i++) {
        uint64_t f = i < nfills ? fills[i] : 0;
        memcpy(p, &f, 8);
        p += 8;
    }
}

static NSArray<id<MTLDevice>> *all_devices(void) {
    NSArray<id<MTLDevice>> *devs = MTLCopyAllDevices();
    if (devs && devs.count > 0) {
        return devs;
    }
    id<MTLDevice> def = MTLCreateSystemDefaultDevice();
    if (!def) {
        return @[];
    }
    return @[def];
}

static id<MTLDevice> pick_device(NSDictionary<NSString *, NSString *> *opts) {
    NSArray<id<MTLDevice>> *devs = all_devices();
    if (devs.count == 0) {
        fail("no Metal GPU");
    }
    int idx = opts[@"device"] ? opts[@"device"].intValue : 0;
    if (idx < 0 || idx >= (int)devs.count) {
        idx = 0;
    }
    return devs[(NSUInteger)idx];
}

static void run_devices(void) {
    NSArray<id<MTLDevice>> *devs = all_devices();
    printf("[");
    for (NSUInteger i = 0; i < devs.count; i++) {
        char name[512];
        json_escape(devs[i].name.UTF8String ?: "Metal GPU", name, sizeof(name));
        uint64_t mem = 0;
        if ([devs[i] respondsToSelector:@selector(recommendedMaxWorkingSetSize)]) {
            mem = (uint64_t)devs[i].recommendedMaxWorkingSetSize;
        }
        printf("%s{\"id\":%lu,\"name\":\"%s\",\"type\":\"metal\",\"memory\":%llu}",
               i ? "," : "", (unsigned long)i, name, (unsigned long long)mem);
    }
    printf("]\n");
}

static id<MTLLibrary> compile_library(id<MTLDevice> device, NSString *path) {
    NSError *err = nil;
    NSString *src = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:&err];
    if (!src) {
        fprintf(stderr, "cannot read shader: %s\n", err.localizedDescription.UTF8String);
        exit(2);
    }
    MTLCompileOptions *copts = [[MTLCompileOptions alloc] init];
    id<MTLLibrary> lib = [device newLibraryWithSource:src options:copts error:&err];
    if (!lib) {
        fprintf(stderr, "Metal compile failed:\n%s\n", err.localizedDescription.UTF8String);
        exit(2);
    }
    return lib;
}

static void run_selftest(id<MTLDevice> device, id<MTLLibrary> lib) {
    NSError *err = nil;
    id<MTLFunction> fn = [lib newFunctionWithName:@"des_known_answer"];
    id<MTLComputePipelineState> pipe = [device newComputePipelineStateWithFunction:fn error:&err];
    if (!pipe) fail("failed to build DES known-answer pipeline");
    id<MTLCommandQueue> queue = [device newCommandQueue];
    id<MTLBuffer> out = [device newBufferWithLength:16 options:MTLResourceStorageModeShared];
    id<MTLCommandBuffer> cmd = [queue commandBuffer];
    id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
    [enc setComputePipelineState:pipe];
    [enc setBuffer:out offset:0 atIndex:0];
    [enc dispatchThreads:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
    [enc endEncoding];
    [cmd commit];
    [cmd waitUntilCompleted];
    uint64_t *words = out.contents;
    uint64_t expected = 0x85E813540F0AB405ULL;
    uint64_t plain = 0x0123456789ABCDEFULL;
    if (words[0] != expected || words[1] != plain) {
        fprintf(stderr, "DES self-test failed ct=%016llx expected=%016llx roundtrip=%016llx\n",
                words[0], expected, words[1]);
        exit(2);
    }
    fprintf(stderr, "metal_selftest ok  gpu=%s\n", device.name.UTF8String);
}

static void run_brute(id<MTLDevice> device, id<MTLLibrary> lib, NSDictionary<NSString *, NSString *> *opts) {
    NSError *err = nil;
    id<MTLFunction> fn = [lib newFunctionWithName:@"des_brute"];
    id<MTLComputePipelineState> pipe = [device newComputePipelineStateWithFunction:fn error:&err];
    if (!pipe) fail("failed to build DES brute pipeline");
    id<MTLCommandQueue> queue = [device newCommandQueue];

    NSString *keyLenStr = opts[@"key-len"] ?: @"0";
    NSString *charsetStr = opts[@"charset"] ?: @"";
    NSString *padStr = opts[@"pad"] ?: @"00";
    NSString *targetStr = opts[@"target"] ?: @"0";
    NSString *startStr = opts[@"start"] ?: @"0";
    NSString *countStr = opts[@"count"] ?: @"0";
    NSString *batchStr = opts[@"batch"] ?: @"16777216";
    uint32_t key_len = (uint32_t)keyLenStr.intValue;
    const char *cs = charsetStr.UTF8String;
    uint32_t charset_len = (uint32_t)strlen(cs);
    uint32_t pad_byte = (uint32_t)(parse_hex_u64(padStr.UTF8String) & 0xFF);
    uint64_t target = parse_hex_u64(targetStr.UTF8String);
    uint64_t start = strtoull(startStr.UTF8String, NULL, 10);
    uint64_t count = strtoull(countStr.UTF8String, NULL, 10);
    uint64_t batch = strtoull(batchStr.UTF8String, NULL, 10);
    if (key_len == 0 || charset_len == 0 || count == 0) fail("brute requires --key-len --charset --count");

    uint64_t fills[8] = {0};
    int nfills = 0;
    NSString *fillsStr = opts[@"fills"] ?: @"0000000000000000";
    NSArray *parts = [fillsStr componentsSeparatedByString:@","];
    for (NSString *part in parts) {
        if (nfills >= 8) break;
        fills[nfills++] = parse_hex_u64(part.UTF8String);
    }

    id<MTLBuffer> paramsBuf = [device newBufferWithLength:kParamsSize options:MTLResourceStorageModeShared];
    id<MTLBuffer> charsetBuf = [device newBufferWithBytes:cs length:charset_len options:MTLResourceStorageModeShared];
    id<MTLBuffer> hitsBuf = [device newBufferWithLength:kMaxHits * sizeof(Hit) options:MTLResourceStorageModeShared];
    id<MTLBuffer> countBuf = [device newBufferWithLength:4 options:MTLResourceStorageModeShared];

    uint64_t done = 0;
    CFAbsoluteTime t0 = CFAbsoluteTimeGetCurrent();
    while (done < count) {
        uint64_t n = batch < (count - done) ? batch : (count - done);
        *(uint32_t *)countBuf.contents = 0;
        write_params(paramsBuf.contents, start + done, key_len, charset_len, pad_byte,
                     (uint32_t)nfills, (uint32_t)n, target, fills, nfills);

        id<MTLCommandBuffer> cmd = [queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:pipe];
        [enc setBuffer:paramsBuf offset:0 atIndex:0];
        [enc setBuffer:charsetBuf offset:0 atIndex:1];
        [enc setBuffer:hitsBuf offset:0 atIndex:2];
        [enc setBuffer:countBuf offset:0 atIndex:3];
        NSUInteger groups = (NSUInteger)((n + kThreadsPerGroup - 1) / kThreadsPerGroup);
        [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
            threadsPerThreadgroup:MTLSizeMake(kThreadsPerGroup, 1, 1)];
        [enc endEncoding];
        [cmd commit];
        [cmd waitUntilCompleted];

        uint32_t hitn = *(uint32_t *)countBuf.contents;
        if (hitn > kMaxHits) hitn = kMaxHits;
        Hit *hits = hitsBuf.contents;
        for (uint32_t i = 0; i < hitn; i++) {
            printf("{\"index\":%llu,\"key_hex\":\"%016llx\",\"plain_hex\":\"%016llx\"}\n",
                   hits[i].index, hits[i].key, hits[i].plain);
            fflush(stdout);
        }
        done += n;
        double elapsed = CFAbsoluteTimeGetCurrent() - t0;
        if (elapsed < 1e-6) elapsed = 1e-6;
        fprintf(stderr, "gpu %llu/%llu  %.0f keys/s\n", done, count, done / elapsed);
    }
}

int main(int argc, const char **argv) {
    if (argc < 2) {
        fail("usage: med_gpu devices | selftest --shader PATH | brute --shader PATH [options]");
    }
    @autoreleasepool {
        if (strcmp(argv[1], "devices") == 0) {
            run_devices();
            return 0;
        }
        NSDictionary *opts = parse_opts(argc, argv);
        NSString *shader = opts[@"shader"];
        if (!shader) fail("missing --shader");
        id<MTLDevice> device = pick_device(opts);
        id<MTLLibrary> lib = compile_library(device, shader);
        if (strcmp(argv[1], "selftest") == 0) {
            run_selftest(device, lib);
        } else if (strcmp(argv[1], "brute") == 0) {
            run_selftest(device, lib);
            run_brute(device, lib, opts);
        } else {
            fail("unknown command");
        }
    }
    return 0;
}