#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "elf.h"

int elf_write(const char *path, CodeBuffer *text, CodeBuffer *rodata,
              CodeBuffer *data) {
    FILE *fp = fopen(path, "wb");
    if (!fp) {
        fprintf(stderr, "error: cannot create %s\n", path);
        return -1;
    }

    int txtoff = 0x1000;
    int txtvaddr = 0x401000;
    int text_rodata_size = text->size + rodata->size;

    int dataoff = 0x10000;
    int datavaddr = 0x410000;
    int data_size = data->size > 0 ? data->size : 2048;

    unsigned char ehdr[64] = {0};
    ehdr[0] = 0x7F; ehdr[1] = 'E'; ehdr[2] = 'L'; ehdr[3] = 'F';
    ehdr[4] = 2;
    ehdr[5] = 1;
    ehdr[6] = 1;
    ehdr[16] = 2;
    ehdr[18] = 0xB7;
    ehdr[20] = 1;
    unsigned int entry = 0x401000;
    memcpy(ehdr + 24, &entry, 4);
    ehdr[32] = 0x40;
    ehdr[52] = 0x40;
    ehdr[54] = 0x38;
    ehdr[56] = 2;
    if (fwrite(ehdr, 1, 64, fp) != 64) { fclose(fp); return -1; }

    unsigned char phdr1[56] = {0};
    phdr1[0] = 1;
    phdr1[4] = 5;
    unsigned int off1 = txtoff;
    unsigned int vaddr1 = txtvaddr;
    unsigned int fsz1 = text_rodata_size;
    unsigned int msz1 = text_rodata_size;
    unsigned long long align1 = 0x10000;
    memcpy(phdr1 + 8, &off1, 4);
    memcpy(phdr1 + 16, &vaddr1, 4);
    memcpy(phdr1 + 24, &vaddr1, 4);
    memcpy(phdr1 + 32, &fsz1, 4);
    memcpy(phdr1 + 40, &msz1, 4);
    memcpy(phdr1 + 48, &align1, 8);
    if (fwrite(phdr1, 1, 56, fp) != 56) { fclose(fp); return -1; }

    unsigned char phdr2[56] = {0};
    phdr2[0] = 1;
    phdr2[4] = 6;
    unsigned int off2 = dataoff;
    unsigned int vaddr2 = datavaddr;
    unsigned int fsz2 = data_size;
    unsigned int msz2 = data_size;
    memcpy(phdr2 + 8, &off2, 4);
    memcpy(phdr2 + 16, &vaddr2, 4);
    memcpy(phdr2 + 24, &vaddr2, 4);
    memcpy(phdr2 + 32, &fsz2, 4);
    memcpy(phdr2 + 40, &msz2, 4);
    memcpy(phdr2 + 48, &align1, 8);
    if (fwrite(phdr2, 1, 56, fp) != 56) { fclose(fp); return -1; }

    unsigned char pad[4096];
    memset(pad, 0, sizeof(pad));
    long pos = ftell(fp);
    if (pos < 0 || pos > txtoff) { fclose(fp); return -1; }
    while (pos < txtoff) {
        size_t chunk = (size_t)(txtoff - pos);
        if (chunk > sizeof(pad)) chunk = sizeof(pad);
        if (fwrite(pad, 1, chunk, fp) != chunk) { fclose(fp); return -1; }
        pos += (long)chunk;
    }

    if (text->size > 0 && fwrite(text->data, 1, (size_t)text->size, fp) != (size_t)text->size) { fclose(fp); return -1; }
    if (rodata->size > 0 && fwrite(rodata->data, 1, (size_t)rodata->size, fp) != (size_t)rodata->size) { fclose(fp); return -1; }

    pos = ftell(fp);
    while (pos < dataoff) {
        size_t chunk = (size_t)(dataoff - pos);
        if (chunk > sizeof(pad)) chunk = sizeof(pad);
        if (fwrite(pad, 1, chunk, fp) != chunk) { fclose(fp); return -1; }
        pos += (long)chunk;
    }

    unsigned char *dbuf = calloc(1, (size_t)data_size);
    if (data->size > 0) memcpy(dbuf, data->data, (size_t)data->size);
    if (fwrite(dbuf, 1, (size_t)data_size, fp) != (size_t)data_size) { free(dbuf); fclose(fp); return -1; }
    free(dbuf);

    fclose(fp);
    return 0;
}
