// omfEmit — wrap a flat binary in a minimal Apple IIgs OMF v2.1
// container so GS/OS can load and execute it.
//
// Single-segment output (CODE, kind=0), no INTERSEG opcodes (multi-
// segment output is a follow-on).  Header layout per OMF 2.1 spec:
// 44-byte fixed header + 10-byte LOAD_NAME + 32-byte SEG_NAME, then
// the body (DS opcode for the payload, END opcode terminator).
//
// CLI mirrors the Python tool exactly:
//   omfEmit --input flat.bin --map flat.map --base 0x8000
//           --entry main --output prog.omf [--name SEG]

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <map>
#include <string>
#include <vector>

namespace {

[[noreturn]] static void die(const std::string &msg) {
    std::fprintf(stderr, "omfEmit: %s\n", msg.c_str());
    std::exit(1);
}

static std::vector<uint8_t> readFile(const std::string &path) {
    std::ifstream f(path, std::ios::binary);
    if (!f) die("cannot open '" + path + "' for reading");
    return std::vector<uint8_t>((std::istreambuf_iterator<char>(f)),
                                 std::istreambuf_iterator<char>());
}

static std::map<std::string, uint32_t> readMap(const std::string &path) {
    std::map<std::string, uint32_t> syms;
    std::ifstream f(path);
    if (!f) die("cannot open '" + path + "' for reading");
    std::string line;
    while (std::getline(f, line)) {
        auto eq = line.find(" = ");
        if (eq == std::string::npos) continue;
        std::string name = line.substr(0, eq);
        std::string addr = line.substr(eq + 3);
        // Trim trailing whitespace.
        while (!name.empty() && std::isspace((unsigned char)name.back()))
            name.pop_back();
        while (!addr.empty() && std::isspace((unsigned char)addr.back()))
            addr.pop_back();
        try {
            syms[name] = std::stoul(addr, nullptr, 16);
        } catch (...) { /* skip non-hex entries */ }
    }
    return syms;
}

// Emit little-endian.
static void put32(std::vector<uint8_t> &v, uint32_t x) {
    v.push_back(x & 0xFF);
    v.push_back((x >> 8) & 0xFF);
    v.push_back((x >> 16) & 0xFF);
    v.push_back((x >> 24) & 0xFF);
}
static void put16(std::vector<uint8_t> &v, uint16_t x) {
    v.push_back(x & 0xFF);
    v.push_back((x >> 8) & 0xFF);
}

static std::vector<uint8_t> emitOMF(const std::vector<uint8_t> &image,
                                     uint32_t entryOffset,
                                     const std::string &name) {
    // Body: DS (literal data) + END.
    std::vector<uint8_t> body;
    if (!image.empty()) {
        body.push_back(0xF1);                       // DS opcode
        put32(body, static_cast<uint32_t>(image.size()));
        body.insert(body.end(), image.begin(), image.end());
    }
    body.push_back(0x00);                           // END opcode

    // LOAD_NAME: 10 bytes, space-padded.
    std::string loadName = name.substr(0, 10);
    while (loadName.size() < 10) loadName += ' ';

    // SEG_NAME: 1-byte length prefix + 31 bytes (truncated, padded with NUL).
    std::string segNameTxt = name.substr(0, 31);
    std::vector<uint8_t> segName;
    segName.push_back(static_cast<uint8_t>(segNameTxt.size()));
    for (char c : segNameTxt) segName.push_back((uint8_t)c);
    while (segName.size() < 32) segName.push_back(0);

    constexpr uint16_t DISPNAME = 44;
    const uint16_t DISPDATA = DISPNAME + 10 + 32;
    const uint32_t LENGTH   = static_cast<uint32_t>(image.size());
    const uint32_t BYTECNT  = DISPDATA + static_cast<uint32_t>(body.size());
    const uint32_t RESSPC   = 0;
    const uint32_t BANKSIZE = 0x10000;
    const uint16_t KIND     = 0x0000;  // CODE
    const uint32_t ORG      = 0;
    const uint32_t ALIGN    = 0;
    const uint8_t  NUMSEX   = 0;
    const uint16_t SEGNUM   = 1;
    const uint32_t ENTRY    = entryOffset;

    std::vector<uint8_t> hdr;
    put32(hdr, BYTECNT);
    put32(hdr, RESSPC);
    put32(hdr, LENGTH);
    hdr.push_back(0x00);                            // undefined
    hdr.push_back(10);                              // LABLEN
    hdr.push_back(4);                               // NUMLEN
    hdr.push_back(0x21);                            // VERSION 2.1
    put32(hdr, BANKSIZE);
    put16(hdr, KIND);
    hdr.push_back(0x00); hdr.push_back(0x00);       // undefined (2 bytes)
    put32(hdr, ORG);
    put32(hdr, ALIGN);
    hdr.push_back(NUMSEX);
    hdr.push_back(0x00);                            // undefined
    put16(hdr, SEGNUM);
    put32(hdr, ENTRY);
    put16(hdr, DISPNAME);
    put16(hdr, DISPDATA);

    if (hdr.size() != 44) die("internal: header size != 44");

    std::vector<uint8_t> out;
    out.insert(out.end(), hdr.begin(), hdr.end());
    out.insert(out.end(), loadName.begin(), loadName.end());
    out.insert(out.end(), segName.begin(), segName.end());
    out.insert(out.end(), body.begin(), body.end());
    return out;
}

static uint32_t parseInt(const std::string &s) {
    return static_cast<uint32_t>(std::stoul(s, nullptr, 0));
}

static void usage(const char *argv0) {
    std::fprintf(stderr,
        "usage: %s --input FLAT --map FILE --base ADDR --entry SYM\n"
        "           --output OMF [--name NAME]\n",
        argv0);
    std::exit(2);
}

} // namespace

int main(int argc, char **argv) {
    std::string input, mapFile, output, entry = "main", name;
    uint32_t base = 0;
    bool baseSet = false;

    int i = 1;
    while (i < argc) {
        std::string a = argv[i];
        if (a == "--input") { if (++i >= argc) usage(argv[0]); input = argv[i++]; }
        else if (a == "--map") { if (++i >= argc) usage(argv[0]); mapFile = argv[i++]; }
        else if (a == "--base") { if (++i >= argc) usage(argv[0]); base = parseInt(argv[i++]); baseSet = true; }
        else if (a == "--entry") { if (++i >= argc) usage(argv[0]); entry = argv[i++]; }
        else if (a == "--name") { if (++i >= argc) usage(argv[0]); name = argv[i++]; }
        else if (a == "--output" || a == "-o") { if (++i >= argc) usage(argv[0]); output = argv[i++]; }
        else if (a == "-h" || a == "--help") usage(argv[0]);
        else die("unknown option '" + a + "'");
    }
    if (input.empty() || mapFile.empty() || !baseSet || output.empty())
        usage(argv[0]);

    auto image = readFile(input);
    auto syms  = readMap(mapFile);

    auto it = syms.find(entry);
    if (it == syms.end())
        die("entry symbol '" + entry + "' not in map");
    uint32_t entryAddr = it->second;
    if (entryAddr < base || entryAddr >= base + image.size())
        die("entry symbol outside linked image");
    uint32_t entryOff = entryAddr - base;

    if (name.empty()) {
        // Default name: output basename without extension.
        size_t slash = output.find_last_of('/');
        std::string base_n = (slash == std::string::npos) ? output
                                                          : output.substr(slash + 1);
        size_t dot = base_n.find_last_of('.');
        name = (dot == std::string::npos) ? base_n : base_n.substr(0, dot);
    }

    auto blob = emitOMF(image, entryOff, name);
    std::ofstream f(output, std::ios::binary);
    if (!f) die("cannot open '" + output + "' for writing");
    f.write(reinterpret_cast<const char *>(blob.data()), blob.size());

    std::fprintf(stderr,
        "OMF: 1 segment, %zu bytes payload, entry='%s' at +0x%x -> %s "
        "(%zu bytes total)\n",
        image.size(), entry.c_str(), entryOff,
        output.c_str(), blob.size());
    return 0;
}