#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

#define IMAGE_WIDTH  2048 // Nothing above 61 seen in images.dat
#define IMAGE_HEIGHT 2048

uint8_t image[IMAGE_WIDTH*IMAGE_HEIGHT];

#ifdef WIN32
#include <windows.h>
BOOL make_executable(void* data, uint32_t size) {
    DWORD old;
    return VirtualProtect(data, size, PAGE_EXECUTE_READWRITE, &old);
}
#endif

#ifdef _MSC_VER
void call_code(void* code) {
    __asm {
        pushad                      // code seems to expect some register to be preserved?
        mov     edi, offset image   // image pointer
        mov     edx, IMAGE_WIDTH    // stride
        call    [code]
        popad
    }
}
#endif

void* read_file(const char* filename, uint32_t* size) {
    FILE* fp = fopen(filename, "rb");
    if (!fp) {
        perror(filename);
        return NULL;
    }

    fseek(fp, 0L, SEEK_END);
    *size = (uint32_t)ftell(fp);
    fseek(fp, 0L, SEEK_SET);

    void* data = malloc(*size);
    if (!data) {
        fclose(fp);
        return NULL;
    }

    fread(data, *size, 1, fp);

    if (ferror(fp)) {
        free(data);
        data = NULL;
    }
    fclose(fp);
    return data;
}

int write_image(const char* filename, int w, int h) {
    FILE* fp = fopen(filename, "wb");
    if (!fp) {
        perror(filename);
        return 0;
    }
    fprintf(fp, "P6\n%d %d 255\n", w, h);
    for (int y = 0; y < h; ++y) {
        for (int x = 0; x < w; ++x) {
            const uint8_t c = image[x+y*IMAGE_WIDTH];
            fputc(c, fp);
            fputc(c, fp);
            fputc(c, fp);
        }
    }
    fclose(fp);
    return 1;
}

//
// Returns one past next 'ret' (0xC3) instruction
//
// Recognizes only the subset of x86 actually used and takes
// as many shortcuts as possible
uint8_t* find_end(uint8_t* start, uint8_t* end) {
    while (start < end) {
        uint8_t inst = *start++;
        const int op_size = inst == 0x66; // Operand size prefix
        if (op_size) inst = *start++;

        if ((inst&0xf0) == 0x40) {
            // inc/dec r16/r32
        } else if ((inst&0xf8) == 0xb0) {
            // MOV r8, imm8
            start++; // imm8
        } else if ((inst&0xf8) == 0xb8) {
            // MOV  r16/r32, imm16/imm32
            start += op_size ? 2 : 4;
        } else if (inst == 0x03) {
            // ADD  r16/32  r/m16/32
            start++; // MODR/M
        } else if (inst == 0x83) {
            // ALUOP r32, imm8
            start++; // MODR/M
            start++; // imm8
        } else if (inst == 0x8a || inst == 0x8b) {
            // MOV  r16/32  r/m16/32
            const uint8_t modrm = *start++; // MODR/M
            if (modrm&0x40) start++;
        } else if (inst == 0x69) {
            // IMUL r16/32  r/m16/32 imm16/32
            start++; // MODR/M
            start += op_size ? 2 : 4;
        } else if (inst == 0x6b) {
            // IMUL r16/32 r/m16/32 imm8
            start++; // MODR/M
            start++; // imm8
        } else if (inst == 0x86) {
            // XCHG r8, r/m8
            start++; // MODR/M
        } else if (inst == 0xaa) {
            // stosb
        } else if (inst == 0xab) {
            // stosw/stosd
        } else if (inst == 0xc3) {
            // ret!
            break;
        } else {
            printf("Unhandled instruction %02X\n", inst);
            abort();
        }
    }
    return start;
}

int main() {
    uint32_t size;
    uint8_t* data = read_file("images.dat", &size);
    if (!data) {
        return 1;
    }
    if (!make_executable(data, size)) {
        free(data);
        return 2;
    }
    int cnt = 0, maxw = 0, maxh = 0;
    // 4 byte header?
    for (uint8_t* code = data + 4; code < data + size; ++cnt) {
        memset(image, 0, sizeof(image));
        call_code(code);
        int w=0,h=0;
        for (int y = 0; y < IMAGE_HEIGHT; ++y) {
            for (int x = 0; x < IMAGE_WIDTH; ++x) {
                if (image[x+y*IMAGE_WIDTH]) {
                    if (x > w) w = x + 1;
                    if (y > h) h = y + 1;
                }
            }
        }
        printf("%08X: %dx%d\n", code-data-4, w, h);
        if (w > maxw) maxw = w;
        if (h > maxh) maxh = h;

        char filename[32];
        snprintf(filename, sizeof(filename), "img%06X.ppm", code-data);
        write_image(filename, w, h);

        code = find_end(code, data + size);
    }
    printf("Max size: %d x %d\n", maxw, maxh);
    free(data);
    return 0;
}