// The MIT License (MIT)
//
// Copyright (C) 2026 Scott Duensing
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to
// deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.

// 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>
#include <sys/stat.h>

// Permission bits (setuid/setgid/sticky plus rwxrwxrwx) copied from the
// original file onto its replacement when a resource add rewrites through a
// temp-and-rename swap.  Without this the fresh temp file keeps only the
// umask default and strips the executable bit off host tools.
#define DVX_RES_PERM_MASK 07777

// ============================================================
// 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;
    }

    // Bound the directory size read from the footer before allocating: an
    // entryCount that can't fit in the file would wrap the 32-bit multiply
    // below and overrun the allocation.  Treat a corrupt tail as 'no
    // existing resources' (return 0) rather than aborting the append.
    if (footer.entryCount > (uint32_t)((fileSize - (long)sizeof(DvxResFooterT)) / (long)sizeof(DvxResDirEntryT))) {
        fclose(f);
        return 0;
    }

    long maxDirOffset = fileSize - (long)sizeof(DvxResFooterT);

    if (maxDirOffset < 0 || footer.dirOffset > (uint32_t)maxDirOffset) {
        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) {
    // Write the rebuilt file to a sibling temp file and atomically swap it
    // in only after every write succeeds.  Rewriting the DXE in place would
    // truncate the original; a disk-full or power loss mid-write would then
    // leave a corrupt, unrecoverable .app/.wgt.
    char tmpPath[DVX_MAX_PATH];

    int32_t tmpLen = snprintf(tmpPath, sizeof(tmpPath), "%s.tmp", path);

    // A truncated temp path can collapse to the original path itself, and
    // fopen(tmpPath, "wb") would then destroy the file the temp-file scheme
    // exists to protect.  Abort instead of writing anywhere unexpected.
    if (tmpLen < 0 || (size_t)tmpLen >= sizeof(tmpPath)) {
        return -1;
    }

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

    if (!f) {
        return -1;
    }

    // Capture the original file's permission bits before the swap.  A freshly
    // created temp file is born with only the umask default, so without this
    // the swap silently strips the executable bit off host tools like bascomp
    // and leaves them unrunnable (exit 127) after a resource add.
    struct stat srcStat;

    bool haveSrcMode = (stat(path, &srcStat) == 0);

    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(tmpPath, "wb");

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

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

    free(dxeBuf);

    for (uint32_t i = 0; i < count; i++) {
        entries[i].offset = (uint32_t)ftell(f);

        if (fwrite(data[i], 1, entries[i].size, f) != entries[i].size) {
            fclose(f);
            remove(tmpPath);
            return -1;
        }
    }

    uint32_t dirOffset = (uint32_t)ftell(f);

    if (fwrite(entries, sizeof(DvxResDirEntryT), count, f) != count) {
        fclose(f);
        remove(tmpPath);
        return -1;
    }

    DvxResFooterT footer;

    footer.magic      = DVX_RES_MAGIC;
    footer.dirOffset  = dirOffset;
    footer.entryCount = count;
    footer.reserved   = 0;

    if (fwrite(&footer, sizeof(footer), 1, f) != 1) {
        fclose(f);
        remove(tmpPath);
        return -1;
    }

    if (fclose(f) != 0) {
        remove(tmpPath);
        return -1;
    }

    // Carry the original file's permission bits onto the replacement so the
    // atomic swap does not downgrade an executable to a plain data file.
    if (haveSrcMode) {
        chmod(tmpPath, srcStat.st_mode & DVX_RES_PERM_MASK);
    }

    // DOS rename cannot overwrite an existing file, so the original is renamed
    // aside to a backup rather than deleted outright.  Only after the
    // replacement is renamed into place is the backup removed; if the swap
    // fails the original is restored, so it survives every error path.
    char bakPath[DVX_MAX_PATH];

    int32_t bakLen = snprintf(bakPath, sizeof(bakPath), "%s.bak", path);

    if (bakLen < 0 || (size_t)bakLen >= sizeof(bakPath)) {
        remove(tmpPath);
        return -1;
    }

    // A stale backup from an earlier failed swap would block renaming the
    // original aside.
    remove(bakPath);

    bool hadOriginal = (rename(path, bakPath) == 0);

    if (rename(tmpPath, path) != 0) {
        // Swap failed: restore the original and drop the written replacement.
        if (hadOriginal) {
            rename(bakPath, path);
        }
        remove(tmpPath);
        return -1;
    }

    // The replacement is in place; the saved original is no longer needed.
    if (hadOriginal) {
        remove(bakPath);
    }

    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;
    int32_t          readResult = dvxResReadExisting(path, dxeSize, &entries, &count, &data);

    // A read error has already freed any partial allocations and reset the
    // outputs; bail rather than truncate the file to only the new entry.
    if (readResult < 0) {
        return -1;
    }

    // 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.  Use tmp pointers so a realloc failure doesn't
    // strand the existing blocks as unreachable leaks.  Abort the
    // write on OOM -- nothing useful to recover to here.
    DvxResDirEntryT *newEntries = (DvxResDirEntryT *)realloc(entries, (count + 1) * sizeof(DvxResDirEntryT));

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

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

    entries = newEntries;

    uint8_t **newData = (uint8_t **)realloc(data, (count + 1) * sizeof(uint8_t *));

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

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

    data = newData;

    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);

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

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

    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
