#include <string>
#include <vector>
#include <stdio.h>
using namespace std;
struct fsHeader
{
    int numFiles;
};
struct fsEntry
{
    char fileName[512];
    unsigned int filePtr;
};
struct fsFile
{
    char fileName[512];
    char *data;
    unsigned int size;
};

class SimpleArchive
{
public:
    SimpleArchive(string fileName, bool readFile)
    {
        this->fileName = fileName;

        if(readFile) readFromFile();
    }

    void addFile(string fileName, char *data, int size)
    {
        fsFile file = {0};

        // Convert string to char[512] safely
        strncpy(file.fileName, fileName.c_str(), 512);
        file.fileName[512-1] = '\0';

        // Make a copy of data.
        file.data = (char *)malloc(size);
        memcpy(file.data, data, size);

        // Set the size.
        file.size = size;

        files.push_back(file);
    }

    bool removeFile(string fileName)
    {
        int fileID = getFileID(fileName);
        if(fileID == -1) return false;

        files.erase(files.begin()+fileID);
        return true;
    }

    bool writeToFile()
    {
        int filePtr = 0;
        fsHeader header;
        header.numFiles = files.size();

        fsEntry *entries = (fsEntry *)malloc(sizeof(fsEntry) * header.numFiles);

        // Offset so filePtr is the first byte after the header and all entries.
        filePtr += sizeof(fsHeader) + sizeof(fsEntry) * header.numFiles;

        // Calculate entries
        for(int i=0; i<header.numFiles; ++i)
        {
            // Copy filename
            memcpy(entries[i].fileName, files[i].fileName, 512);

            // Set fileptr, then adjust fileptr.
            entries[i].filePtr = filePtr;
            filePtr += files[i].size;
        }

        // Open file
        FILE *file = fopen(fileName.c_str(), "wb");
        if(file == NULL) return false;

        // Write header+entries
        fwrite(&header, sizeof(fsHeader), 1, file);
        fwrite(entries, sizeof(fsEntry), header.numFiles, file);

        // Write files
        for(int i=0; i<header.numFiles; ++i)
        {
            fwrite(files[i].data, files[i].size, 1, file);
        }

        free(entries);

        return true;
    }

    bool readFromFile()
    {
        int fileSize = -1;
        FILE *file = fopen(fileName.c_str(), "rb"); // Read file as binary.

        if(file == NULL) return false;

        // Trick: Get filesize.
        fseek(file, 0L, SEEK_END);
        fileSize = ftell(file);
        fseek(file, 0L, SEEK_SET);

        fsHeader header;
        fread(&header, sizeof(fsHeader), 1, file);

        fsEntry *entries = (fsEntry *)malloc(sizeof(fsEntry) * header.numFiles);

        // Get all of the file entries.
        for(int i=0; i<header.numFiles; ++i)
        {
            fread(&entries[i], sizeof(fsEntry), 1, file);
        }

        // Get all of the actual files.
        for(int i=0; i<header.numFiles-1; ++i)
        {
            int size = -1;
            // Special case for last file: use container filesize to calculate filesize.
            if(i != header.numFiles-1)
            {
                // Calculate size.
                size = entries[i+1].filePtr - entries[i].filePtr;
            }
            else
            {
                // Calculate size.
                size = fileSize - entries[header.numFiles-1].filePtr;
            }

            // Allocate data.
            char *data = (char *)malloc(size);

            // Read actual data.
            fread(data, size, 1, file);

            // Create file entry for memory storage.
            fsFile f = {0};
            memcpy(f.fileName, entries[i].fileName, 512);
            f.data = data;
            f.size = size;

            // Push file.
            files.push_back(f);
        }

        fclose(file);

        free(entries);

        return true;
    }

    fsFile *getFile(string fileName)
    {
        int fileID = getFileID(fileName);
        if(fileID == -1) return NULL;

        return &files[fileID];
    }

private:
    int getFileID(string fileName)
    {
        for(int i=files.size()-1; i>=0; --i)
        {
            if(fileName == files[i].fileName)
            {
                return i;
            }
        }
        return -1;
    }

    vector<fsFile> files;
    string fileName;
};