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

#define MAX_FILES 128
#define MAX_FILENAME 64
#define INITRD_MAGIC 0xBF

struct initrd_header
{
    unsigned char magic; // The magic number is there to check for consistency.
    char name[64];
    unsigned int offset; // Offset in the initrd the file starts.
    unsigned int length; // Length of the file.
};

struct mkinitrd_file {
    char  file_src[MAX_FILENAME];
    char  initrd_path[MAX_FILENAME];
};

struct mkinitrd_args {
    char                  outpath[256];
    int                   nfiles;
    struct mkinitrd_file  files[MAX_FILES];
};

void parse_args(int argc, char **argv, struct mkinitrd_args *out)
{
    int count = argc;
    int file_i = 0;
    int i = 0;

    strcpy(out->outpath, "./initrd.img");
    out->nfiles = 0;
    for (i=1; i<argc; ++i) {
        if (strcmp(argv[i], "-o") == 0 && argc > (i+1)) {
            strcpy(out->outpath, argv[i+1]);
            i++; // skip outpath arg.
        } else {
            strcpy(out->files[file_i].file_src, argv[i]);
            strcpy(out->files[file_i].initrd_path, argv[i+1]);
            i++;
            file_i++;
            out->nfiles++;
        }
    }
}

int main(int argc, char **argv)
{
    struct initrd_header headers[64];
    unsigned int off = sizeof(struct initrd_header) * MAX_FILES + sizeof(int);
    struct mkinitrd_args args;
    int i;

    parse_args(argc, argv, &args);

    for(i = 0; i < args.nfiles; i++) {
        strcpy(headers[i].name, args.files[i].initrd_path);
        headers[i].offset = off;
        FILE *stream = fopen(args.files[i].file_src, "r");
        if(stream == 0) {
            printf("Error: file not found: %s\n", args.files[i].file_src);
            return;
        }
        fseek(stream, 0, SEEK_END);
        headers[i].length = ftell(stream);
        off += headers[i].length;
        fclose(stream);
        headers[i].magic = INITRD_MAGIC;
    }

    FILE *wstream = fopen(args.outpath, "w");
    unsigned char *data = (unsigned char *)malloc(off);
    fwrite(&(args.nfiles), sizeof(int), 1, wstream);
    fwrite(headers, sizeof(struct initrd_header), MAX_FILES, wstream);

    for(i = 0; i < args.nfiles; i++) {
        FILE *stream = fopen(args.files[i].file_src, "r");
        unsigned char *buf = (unsigned char *)malloc(headers[i].length);
        fread(buf, 1, headers[i].length, stream);
        fwrite(buf, 1, headers[i].length, wstream);
        fclose(stream);
        free(buf);
    }

    fclose(wstream);
    free(data);

    return 0;
}