// dvxResWrite.h -- DVX resource writing library interface
//
// Shared implementation for appending resources to DXE files.
// Used by both the dvxres command-line tool and the runtime
// dvxResAppend() function.
//
// Include this header in exactly one .c file. The tool includes
// it directly; the runtime includes it via dvxResource.c.

#ifndef DVX_RES_WRITE_H
#define DVX_RES_WRITE_H

#include "../libs/kpunch/libdvx/dvxRes.h"

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

// ============================================================
// dvxResDxeContentSize
// ============================================================
//
// Returns the size of the DXE3 content (before any appended
// resources). If the file has a resource footer, returns the
// start of the resource data. Otherwise returns the full file size.

static long dvxResDxeContentSize(const char *path) {
    FILE *f = fopen(path, "rb");

    if (!f) {
        return -1;
    }

    fseek(f, 0, SEEK_END);
    long fileSize = ftell(f);

    if (fileSize < (long)sizeof(DvxResFooterT)) {
        fclose(f);
        return fileSize;
    }

    fseek(f, -(long)sizeof(DvxResFooterT), SEEK_END);
    DvxResFooterT footer;

    if (fread(&footer, sizeof(footer), 1, f) == 1 && footer.magic == DVX_RES_MAGIC) {
        fseek(f, footer.dirOffset, SEEK_SET);
        DvxResDirEntryT entry;
        long earliest = footer.dirOffset;

        for (uint32_t i = 0; i < footer.entryCount; i++) {
            if (fread(&entry, sizeof(entry), 1, f) == 1) {
                if ((long)entry.offset < earliest) {
                    earliest = (long)entry.offset;
                }
            }
        }

        fclose(f);
        return earliest;
    }

    fclose(f);
    return fileSize;
}


// ============================================================
// dvxResReadExisting
// ============================================================
//
// Reads any existing resource block from the file. Returns arrays
// of directory entries and data blobs. Caller frees everything.

static int dvxResReadExisting(const char *path, long dxeSize, DvxResDirEntryT **outEntries, uint32_t *outCount, uint8_t ***outData) {
    *outEntries = NULL;
    *outCount   = 0;
    *outData    = NULL;

    FILE *f = fopen(path, "rb");

    if (!f) {
        return -1;
    }

    fseek(f, 0, SEEK_END);
    long fileSize = ftell(f);

    if (fileSize <= dxeSize) {
        fclose(f);
        return 0;
    }

    fseek(f, -(long)sizeof(DvxResFooterT), SEEK_END);
    DvxResFooterT footer;

    if (fread(&footer, sizeof(footer), 1, f) != 1 || footer.magic != DVX_RES_MAGIC) {
        fclose(f);
        return 0;
    }

    fseek(f, footer.dirOffset, SEEK_SET);
    DvxResDirEntryT *entries = (DvxResDirEntryT *)malloc(footer.entryCount * sizeof(DvxResDirEntryT));

    if (!entries) {
        fclose(f);
        return -1;
    }

    if (fread(entries, sizeof(DvxResDirEntryT), footer.entryCount, f) != footer.entryCount) {
        free(entries);
        fclose(f);
        return -1;
    }

    uint8_t **data = (uint8_t **)malloc(footer.entryCount * sizeof(uint8_t *));

    if (!data) {
        free(entries);
        fclose(f);
        return -1;
    }

    for (uint32_t i = 0; i < footer.entryCount; i++) {
        data[i] = (uint8_t *)malloc(entries[i].size);

        if (!data[i]) {
            for (uint32_t j = 0; j < i; j++) {
                free(data[j]);
            }

            free(data);
            free(entries);
            fclose(f);
            return -1;
        }

        fseek(f, entries[i].offset, SEEK_SET);

        if (fread(data[i], 1, entries[i].size, f) != entries[i].size) {
            for (uint32_t j = 0; j <= i; j++) {
                free(data[j]);
            }

            free(data);
            free(entries);
            fclose(f);
            return -1;
        }
    }

    fclose(f);
    *outEntries = entries;
    *outCount   = footer.entryCount;
    *outData    = data;
    return 0;
}


// ============================================================
// dvxResWriteBlock
// ============================================================
//
// Truncates the file to dxeSize and writes the resource block.

static int dvxResWriteBlock(const char *path, long dxeSize, DvxResDirEntryT *entries, uint32_t count, uint8_t **data) {
    FILE *f = fopen(path, "rb");

    if (!f) {
        return -1;
    }

    uint8_t *dxeBuf = (uint8_t *)malloc(dxeSize);

    if (!dxeBuf) {
        fclose(f);
        return -1;
    }

    if (fread(dxeBuf, 1, dxeSize, f) != (size_t)dxeSize) {
        free(dxeBuf);
        fclose(f);
        return -1;
    }

    fclose(f);

    f = fopen(path, "wb");

    if (!f) {
        free(dxeBuf);
        return -1;
    }

    fwrite(dxeBuf, 1, dxeSize, f);
    free(dxeBuf);

    for (uint32_t i = 0; i < count; i++) {
        entries[i].offset = (uint32_t)ftell(f);
        fwrite(data[i], 1, entries[i].size, f);
    }

    uint32_t dirOffset = (uint32_t)ftell(f);
    fwrite(entries, sizeof(DvxResDirEntryT), count, f);

    DvxResFooterT footer;
    footer.magic      = DVX_RES_MAGIC;
    footer.dirOffset  = dirOffset;
    footer.entryCount = count;
    footer.reserved   = 0;
    fwrite(&footer, sizeof(footer), 1, f);

    fclose(f);
    return 0;
}


// ============================================================
// dvxResAppendEntry
// ============================================================
//
// High-level: append a single resource to a DXE file.
// Reads existing resources, adds the new one (replacing if
// same name exists), and rewrites the resource block.
// Returns 0 on success.

static int dvxResAppendEntry(const char *path, const char *name, uint32_t type, const void *resData, uint32_t resSize) {
    if (!path || !name || !resData || resSize == 0) {
        return -1;
    }

    long dxeSize = dvxResDxeContentSize(path);

    if (dxeSize < 0) {
        return -1;
    }

    DvxResDirEntryT *entries = NULL;
    uint32_t         count   = 0;
    uint8_t        **data    = NULL;
    dvxResReadExisting(path, dxeSize, &entries, &count, &data);

    // Remove existing entry with same name
    for (uint32_t i = 0; i < count; i++) {
        if (strcmp(entries[i].name, name) == 0) {
            free(data[i]);

            for (uint32_t j = i; j < count - 1; j++) {
                entries[j] = entries[j + 1];
                data[j]    = data[j + 1];
            }

            count--;
            break;
        }
    }

    // Add new entry
    entries = (DvxResDirEntryT *)realloc(entries, (count + 1) * sizeof(DvxResDirEntryT));
    data    = (uint8_t **)realloc(data, (count + 1) * sizeof(uint8_t *));

    memset(&entries[count], 0, sizeof(DvxResDirEntryT));
    strncpy(entries[count].name, name, DVX_RES_NAME_MAX - 1);
    entries[count].type = type;
    entries[count].size = resSize;
    data[count] = (uint8_t *)malloc(resSize);
    memcpy(data[count], resData, resSize);
    count++;

    int result = dvxResWriteBlock(path, dxeSize, entries, count, data);

    for (uint32_t i = 0; i < count; i++) {
        free(data[i]);
    }

    free(data);
    free(entries);
    return result;
}

#endif // DVX_RES_WRITE_H
