#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <stdio.h>
#include <stdint.h>
#include <vector>
#include "crypt.h"

#pragma pack(push, 1)
struct PackTail {
    uint32_t encSize;
    uint32_t magic;
};
#pragma pack(pop)

static const uint32_t kMagic = 0x32505353u; /* SSP2 */

static bool ReadAll(const wchar_t* path, std::vector<uint8_t>& out) {
    HANDLE f = CreateFileW(path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
    if (f == INVALID_HANDLE_VALUE) return false;
    DWORD sz = GetFileSize(f, NULL);
    if (!sz) { CloseHandle(f); return false; }
    out.resize(sz);
    DWORD rd = 0;
    BOOL ok = ReadFile(f, out.data(), sz, &rd, NULL);
    CloseHandle(f);
    return ok && rd == sz;
}

int wmain(int argc, wchar_t** argv) {
    if (argc != 4) {
        fwprintf(stderr, L"usage: pack.exe stub.exe core.exe output.exe\n");
        return 1;
    }

    std::vector<uint8_t> stub, core;
    if (!ReadAll(argv[1], stub)) { fwprintf(stderr, L"stub read failed\n"); return 2; }
    if (!ReadAll(argv[2], core)) { fwprintf(stderr, L"core read failed\n"); return 3; }

    uint8_t nonce[12]{};
    LARGE_INTEGER pc, t;
    QueryPerformanceCounter(&pc);
    GetSystemTimeAsFileTime((FILETIME*)&t);
    uint64_t mix = (uint64_t)pc.QuadPart ^ (uint64_t)t.QuadPart ^ GetTickCount64();
    for (int i = 0; i < 12; ++i)
        nonce[i] = (uint8_t)((mix >> ((i * 7) % 61)) ^ SS_KEY[i] ^ SS_KEY[20 + (i % 12)] ^ (uint8_t)(i * 37));

    uint64_t sum = ss_fnv64(core.data(), core.size());
    std::vector<uint8_t> enc = core;
    ss_crypt(enc.data(), enc.size(), nonce);

    HANDLE out = CreateFileW(argv[3], GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (out == INVALID_HANDLE_VALUE) { fwprintf(stderr, L"output create failed\n"); return 4; }

    auto write = [&](const void* p, DWORD n) -> bool {
        DWORD w = 0;
        return WriteFile(out, p, n, &w, NULL) && w == n;
    };

    if (!write(stub.data(), (DWORD)stub.size())) { CloseHandle(out); return 5; }
    if (!write(nonce, 12)) { CloseHandle(out); return 6; }
    if (!write(&sum, 8)) { CloseHandle(out); return 7; }
    if (!write(enc.data(), (DWORD)enc.size())) { CloseHandle(out); return 8; }

    PackTail tail{};
    tail.encSize = (uint32_t)enc.size();
    tail.magic = kMagic;
    if (!write(&tail, sizeof(tail))) { CloseHandle(out); return 9; }

    CloseHandle(out);
    wprintf(L"packed %u bytes -> %s\n", (unsigned)enc.size(), argv[3]);
    return 0;
}