/*
 * mkexe — convert a flat binary or Intel HEX image into a Sprinter ESTEX .EXE.
 *
 * Single-image mode (no banks):
 *   +00   3 bytes  "EXE"
 *   +03   1 byte   version (0)
 *   +04   4 bytes  offset to image (0x00000200)
 *   +08   2 bytes  loader size (0 — whole file at once)
 *   +0A   6 bytes  reserved (0)
 *   +10   2 bytes  load address
 *   +12   2 bytes  start (PC)
 *   +14   2 bytes  initial SP
 *   +16   490 bytes  reserved (0)
 *   +200  ...      contiguous image bytes
 *
 * Multi-bank mode (detected from IHX ELA records):
 *   header.loader = HOME image size (so ESTEX loads only HOME and keeps the
 *                   .EXE file open with the handle in IX-3; crt0_banked.s
 *                   then reads each bank with ESTEX READ)
 *   file layout: header (512) + HOME image + bank1 (16KB) + bank2 (16KB) + ...
 *
 *   IHX convention from SDCC `-Wl-b _BANKn=0xNC000`:
 *     ELA = 0x00       → HOME (low16 = 0x4100..0x7FFF range)
 *     ELA = 0x01       → BANK1 (low16 = 0xC000..0xFFFF)
 *     ELA = 0x02       → BANK2
 *     ...
 *
 * Build:   cc -O2 -Wall -Wextra -std=c99 -o mkexe mkexe.c
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <ctype.h>

#define EXE_HEADER_SIZE   512
#define SEGMENT_SIZE      0x10000u
#define BANK_SIZE         0x4000u   /* 16 KB per bank                    */
#define DEFAULT_BANK_BASE 0xC000u   /* HUGE mode (banks live in window 3)*/
#define BIG_BANK_BASE     0x4000u   /* BIG  mode (banks live in window 1)*/
#define MAX_BANKS         15
#define HOME_SEG          0

#define DEFAULT_STACK     0xBFFEu
#define DEFAULT_LOAD      0x4100u

typedef struct {
    uint8_t  bytes[SEGMENT_SIZE];
    uint8_t  present[SEGMENT_SIZE];
    uint32_t lo;
    uint32_t hi;
    int      any;
} segment_t;

static segment_t segments[MAX_BANKS + 1];

static void segments_init(void) {
    for (int i = 0; i <= MAX_BANKS; i++) {
        memset(segments[i].bytes, 0, sizeof(segments[i].bytes));
        memset(segments[i].present, 0, sizeof(segments[i].present));
        segments[i].lo  = 0xFFFFFFFFu;
        segments[i].hi  = 0;
        segments[i].any = 0;
    }
}

static int segment_set(unsigned seg, uint32_t addr, uint8_t v) {
    if (seg > MAX_BANKS) {
        fprintf(stderr, "mkexe: bank id %u exceeds MAX_BANKS=%d\n", seg, MAX_BANKS);
        return -1;
    }
    if (addr >= SEGMENT_SIZE) {
        fprintf(stderr, "mkexe: address 0x%X outside the 64K Z80 space\n", addr);
        return -1;
    }
    segment_t *s = &segments[seg];
    s->bytes[addr]   = v;
    s->present[addr] = 1;
    if (addr < s->lo) s->lo = addr;
    if (addr > s->hi) s->hi = addr;
    s->any = 1;
    return 0;
}

static int hex_nib(int c) {
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    return -1;
}

static int hex_byte(const char *p) {
    int h = hex_nib((unsigned char)p[0]);
    int l = hex_nib((unsigned char)p[1]);
    if (h < 0 || l < 0) return -1;
    return (h << 4) | l;
}

static int load_ihx(const char *path) {
    FILE *f = fopen(path, "rb");
    if (!f) {
        fprintf(stderr, "mkexe: cannot open '%s': %s\n", path, strerror(errno));
        return -1;
    }
    char line[1024];
    int rc = 0, lineno = 0;
    unsigned seg = 0;       /* upper 8 bits of 24-bit virtual addr */
    uint32_t para_base = 0; /* paragraph base from type-02 records */
    while (fgets(line, sizeof(line), f)) {
        lineno++;
        size_t n = strlen(line);
        while (n > 0 && (line[n-1] == '\n' || line[n-1] == '\r' || line[n-1] == ' ')) {
            line[--n] = 0;
        }
        if (n == 0) continue;
        if (line[0] != ':') {
            fprintf(stderr, "mkexe: %s:%d: expected ':' record start\n", path, lineno);
            rc = -1; break;
        }
        if (n < 11) {
            fprintf(stderr, "mkexe: %s:%d: record too short\n", path, lineno);
            rc = -1; break;
        }
        int count = hex_byte(line + 1);
        int aHi   = hex_byte(line + 3);
        int aLo   = hex_byte(line + 5);
        int type  = hex_byte(line + 7);
        if (count < 0 || aHi < 0 || aLo < 0 || type < 0) {
            fprintf(stderr, "mkexe: %s:%d: bad hex in header\n", path, lineno);
            rc = -1; break;
        }
        if (n < (size_t)(11 + 2 * count)) {
            fprintf(stderr, "mkexe: %s:%d: truncated record\n", path, lineno);
            rc = -1; break;
        }
        unsigned sum = (unsigned)count + (unsigned)aHi + (unsigned)aLo + (unsigned)type;
        uint16_t addr = (uint16_t)((aHi << 8) | aLo);

        if (type == 0x00) {
            for (int i = 0; i < count; i++) {
                int b = hex_byte(line + 9 + 2 * i);
                if (b < 0) { rc = -1; goto done; }
                sum += (unsigned)b;
                uint32_t a = para_base + (uint32_t)addr + (uint32_t)i;
                if (segment_set(seg, a, (uint8_t)b) < 0) { rc = -1; goto done; }
            }
        } else if (type == 0x01) {
            int cc = hex_byte(line + 9);
            if (cc < 0) { rc = -1; goto done; }
            sum += (unsigned)cc;
            if ((sum & 0xFF) != 0) {
                fprintf(stderr, "mkexe: %s:%d: bad checksum on EOF record\n", path, lineno);
                rc = -1; goto done;
            }
            break;
        } else if (type == 0x02) {
            int bHi = hex_byte(line + 9);
            int bLo = hex_byte(line + 11);
            if (bHi < 0 || bLo < 0) { rc = -1; goto done; }
            sum += (unsigned)bHi + (unsigned)bLo;
            para_base = ((uint32_t)((bHi << 8) | bLo)) << 4;
        } else if (type == 0x04) {
            int bHi = hex_byte(line + 9);
            int bLo = hex_byte(line + 11);
            if (bHi < 0 || bLo < 0) { rc = -1; goto done; }
            sum += (unsigned)bHi + (unsigned)bLo;
            uint32_t high = ((uint32_t)bHi << 8) | (uint32_t)bLo;
            if (high > MAX_BANKS) {
                fprintf(stderr, "mkexe: %s:%d: ELA upper16=0x%04X means bank id %u, exceeds MAX_BANKS=%d\n",
                        path, lineno, (unsigned)high, (unsigned)high, MAX_BANKS);
                rc = -1; goto done;
            }
            seg = (unsigned)high;
        } else if (type == 0x03 || type == 0x05) {
            for (int i = 0; i < count; i++) {
                int b = hex_byte(line + 9 + 2 * i);
                if (b < 0) { rc = -1; goto done; }
                sum += (unsigned)b;
            }
        } else {
            fprintf(stderr, "mkexe: %s:%d: unsupported HEX type 0x%02X\n", path, lineno, type);
            rc = -1; goto done;
        }
        int cc = hex_byte(line + 9 + 2 * count);
        if (cc < 0) {
            fprintf(stderr, "mkexe: %s:%d: missing checksum\n", path, lineno);
            rc = -1; goto done;
        }
        sum += (unsigned)cc;
        if ((sum & 0xFF) != 0) {
            fprintf(stderr, "mkexe: %s:%d: checksum mismatch (sum=%02X)\n", path, lineno, sum & 0xFF);
            rc = -1; goto done;
        }
    }
done:
    fclose(f);
    return rc;
}

static int load_bin(const char *path, uint32_t load_addr) {
    FILE *f = fopen(path, "rb");
    if (!f) {
        fprintf(stderr, "mkexe: cannot open '%s': %s\n", path, strerror(errno));
        return -1;
    }
    uint8_t buf[4096];
    uint32_t addr = load_addr;
    size_t n;
    while ((n = fread(buf, 1, sizeof(buf), f)) > 0) {
        for (size_t i = 0; i < n; i++) {
            if (segment_set(HOME_SEG, addr++, buf[i]) < 0) { fclose(f); return -1; }
        }
    }
    fclose(f);
    return 0;
}

static int has_suffix(const char *s, const char *suf) {
    size_t ls = strlen(s), lf = strlen(suf);
    if (lf > ls) return 0;
    return strcasecmp(s + ls - lf, suf) == 0;
}

static void usage(FILE *o) {
    fprintf(o,
        "mkexe — assemble a Sprinter ESTEX .EXE from an IHX or BIN image\n"
        "\n"
        "Usage:\n"
        "  mkexe [options] -o OUT.exe INPUT(.ihx|.bin)\n"
        "\n"
        "Options:\n"
        "  -o FILE         output .exe path (required)\n"
        "  -L ADDR         load address (default: lowest in .ihx, or 0x%04X for .bin)\n"
        "  -E ADDR         entry point / start PC (default: same as load)\n"
        "  -S ADDR         initial SP (default: 0x%04X)\n"
        "  -p PAD          padding byte for gaps in image (default: 0xFF)\n"
        "  -B ADDR         bank low-16 base (default: 0x%04X HUGE; pass 0x%04X BIG)\n"
        "  -v              verbose\n"
        "  -h              this help\n"
        "\n"
        "If the input .ihx contains ELA records (type 04) with high16 = N where\n"
        "1 <= N <= %d, those bytes are packed as bank N (16 KB each, base 0x%04X).\n"
        "The header.loader is then set to the HOME image size so ESTEX keeps the\n"
        "file open for crt0_banked.s to read each bank with RST 10h C=13h.\n"
        "\n"
        "Addresses accept 0x.. / $.. / decimal.\n",
        DEFAULT_LOAD, DEFAULT_STACK,
        DEFAULT_BANK_BASE, BIG_BANK_BASE,
        MAX_BANKS, DEFAULT_BANK_BASE);
}

static int parse_addr(const char *s, uint32_t *out) {
    if (!s || !*s) return -1;
    char *end = NULL;
    int base = 10;
    if (s[0] == '$') { s++; base = 16; }
    else if (s[0] == '0' && (s[1] == 'x' || s[1] == 'X')) { s += 2; base = 16; }
    unsigned long v = strtoul(s, &end, base);
    if (end == s || (end && *end)) return -1;
    *out = (uint32_t)v;
    return 0;
}

static void put16_le(uint8_t *p, uint16_t v) {
    p[0] = (uint8_t)(v & 0xFF);
    p[1] = (uint8_t)(v >> 8);
}

int main(int argc, char **argv) {
    const char *out_path = NULL;
    const char *in_path  = NULL;
    uint32_t load_addr   = 0xFFFFFFFFu;
    uint32_t start_addr  = 0xFFFFFFFFu;
    uint32_t stack_addr  = DEFAULT_STACK;
    uint32_t pad_byte    = 0xFF;
    uint32_t bank_base   = DEFAULT_BANK_BASE;
    int verbose          = 0;

    for (int i = 1; i < argc; i++) {
        const char *a = argv[i];
        if (!strcmp(a, "-h") || !strcmp(a, "--help")) { usage(stdout); return 0; }
        if (!strcmp(a, "-v")) { verbose = 1; continue; }
        if (!strcmp(a, "-o") && i + 1 < argc) { out_path = argv[++i]; continue; }
        if (!strcmp(a, "-L") && i + 1 < argc) {
            if (parse_addr(argv[++i], &load_addr) < 0) { fprintf(stderr, "mkexe: bad -L\n"); return 1; }
            continue;
        }
        if (!strcmp(a, "-E") && i + 1 < argc) {
            if (parse_addr(argv[++i], &start_addr) < 0) { fprintf(stderr, "mkexe: bad -E\n"); return 1; }
            continue;
        }
        if (!strcmp(a, "-S") && i + 1 < argc) {
            if (parse_addr(argv[++i], &stack_addr) < 0) { fprintf(stderr, "mkexe: bad -S\n"); return 1; }
            continue;
        }
        if (!strcmp(a, "-p") && i + 1 < argc) {
            if (parse_addr(argv[++i], &pad_byte) < 0) { fprintf(stderr, "mkexe: bad -p\n"); return 1; }
            continue;
        }
        if (!strcmp(a, "-B") && i + 1 < argc) {
            if (parse_addr(argv[++i], &bank_base) < 0) { fprintf(stderr, "mkexe: bad -B\n"); return 1; }
            if (bank_base + BANK_SIZE > SEGMENT_SIZE) {
                fprintf(stderr, "mkexe: -B 0x%X + 16KB exceeds 0x10000\n", bank_base);
                return 1;
            }
            continue;
        }
        if (a[0] == '-') { fprintf(stderr, "mkexe: unknown option '%s'\n", a); usage(stderr); return 1; }
        if (in_path) { fprintf(stderr, "mkexe: extra positional '%s'\n", a); return 1; }
        in_path = a;
    }

    if (!in_path || !out_path) { usage(stderr); return 1; }

    segments_init();

    if (has_suffix(in_path, ".ihx") || has_suffix(in_path, ".hex")) {
        if (load_ihx(in_path) < 0) return 2;
    } else if (has_suffix(in_path, ".bin")) {
        if (load_addr == 0xFFFFFFFFu) load_addr = DEFAULT_LOAD;
        if (load_bin(in_path, load_addr) < 0) return 2;
    } else {
        fprintf(stderr, "mkexe: input must end with .ihx, .hex, or .bin\n");
        return 1;
    }

    segment_t *home = &segments[HOME_SEG];
    if (!home->any) {
        fprintf(stderr, "mkexe: input has no HOME (segment 0) bytes\n");
        return 2;
    }

    if (load_addr == 0xFFFFFFFFu) load_addr = home->lo;
    if (start_addr == 0xFFFFFFFFu) start_addr = load_addr;

    if (load_addr  > 0xFFFFu) { fprintf(stderr, "mkexe: load address 0x%X > 0xFFFF\n", load_addr);  return 1; }
    if (start_addr > 0xFFFFu) { fprintf(stderr, "mkexe: start address 0x%X > 0xFFFF\n", start_addr); return 1; }
    if (stack_addr > 0xFFFFu) { fprintf(stderr, "mkexe: stack address 0x%X > 0xFFFF\n", stack_addr); return 1; }
    if (pad_byte > 0xFFu) { fprintf(stderr, "mkexe: pad byte 0x%X > 0xFF\n", pad_byte); return 1; }

    if (home->lo < load_addr) {
        fprintf(stderr, "mkexe: HOME has bytes at 0x%04X below requested load 0x%04X\n",
                home->lo, load_addr);
        return 2;
    }

    uint32_t home_size = home->hi + 1 - load_addr;
    if (load_addr + home_size > 0x10000u) {
        fprintf(stderr, "mkexe: HOME extends past 0xFFFF (load=0x%04X, size=0x%X)\n",
                load_addr, home_size);
        return 2;
    }

    /* Inventory banks. */
    int max_bank = 0;
    for (int i = 1; i <= MAX_BANKS; i++) {
        if (segments[i].any) {
            if (segments[i].lo < bank_base) {
                /* Bytes below the bank base inside a BANK segment almost
                   certainly mean the *previous* bank overflowed 16 KB and
                   bled into this 64 KB virtual slot.  Tell the user clearly. */
                int prev = i - 1;
                if (prev >= 1) {
                    fprintf(stderr,
                        "mkexe: BANK%d appears to overflow its 16 KB limit:\n"
                        "       data lands at virtual 0x%X%04X (low16 < 0x%04X)\n"
                        "       Move some code from BANK%d's .c files into a new bank.\n",
                        prev, i, segments[i].lo, bank_base, prev);
                } else {
                    fprintf(stderr,
                        "mkexe: bank %d has bytes at 0x%04X, below the bank base 0x%04X\n",
                        i, segments[i].lo, bank_base);
                }
                return 2;
            }
            if (segments[i].hi >= bank_base + BANK_SIZE) {
                fprintf(stderr,
                    "mkexe: BANK%d overflows its 16 KB limit:\n"
                    "       last byte at 0x%04X (limit 0x%04X)\n",
                    i, segments[i].hi, bank_base + BANK_SIZE - 1);
                return 2;
            }
            if (i > max_bank) max_bank = i;
        }
    }

    /* HOME must not extend past 0xBFFF (W3 = bank territory in HUGE,
       free in BIG but still off-limits to HOME). */
    if (home->hi >= 0xC000u) {
        fprintf(stderr,
            "mkexe: HOME image extends to 0x%04X, past window 2 end (0xBFFF).\n"
            "       Code grew too big for HOME — move some .c files into a bank.\n",
            home->hi);
        return 2;
    }
    /* In BIG mode HOME starts at 0x8100 (W2).  Reject if it accidentally
       overlaps the bank window in W1 (0x4000..0x7FFF) — that means the
       user passed a bad --code-loc and the image would clobber banks. */
    if (bank_base == BIG_BANK_BASE && home->lo < 0x8000u) {
        fprintf(stderr,
            "mkexe: HOME image starts at 0x%04X in BIG mode, but banks own\n"
            "       0x4000..0x7FFF — HOME must live in W2 (0x8000..0xBFFF).\n"
            "       Check --code-loc / -L.\n",
            home->lo);
        return 2;
    }

    /* Banks must be densely numbered from 1 so the loader index matches. */
    for (int i = 1; i <= max_bank; i++) {
        if (!segments[i].any) {
            fprintf(stderr, "mkexe: bank %d is empty but bank %d exists — banks must be consecutive from 1\n",
                    i, max_bank);
            return 2;
        }
    }

    uint32_t loader_size = (max_bank > 0) ? home_size : 0u;

    /* --- Write the EXE. --- */
    uint8_t header[EXE_HEADER_SIZE];
    memset(header, 0, sizeof(header));
    header[0x00] = 'E';
    header[0x01] = 'X';
    header[0x02] = 'E';
    header[0x03] = 0x00;
    header[0x04] = 0x00;
    header[0x05] = 0x02;
    header[0x06] = 0x00;
    header[0x07] = 0x00;
    put16_le(header + 0x08, (uint16_t)loader_size);
    put16_le(header + 0x10, (uint16_t)load_addr);
    put16_le(header + 0x12, (uint16_t)start_addr);
    put16_le(header + 0x14, (uint16_t)stack_addr);

    FILE *out = fopen(out_path, "wb");
    if (!out) {
        fprintf(stderr, "mkexe: cannot open '%s' for writing: %s\n", out_path, strerror(errno));
        return 2;
    }
    if (fwrite(header, 1, EXE_HEADER_SIZE, out) != EXE_HEADER_SIZE) {
        fprintf(stderr, "mkexe: short write on header\n"); fclose(out); return 2;
    }

    /* HOME image */
    for (uint32_t a = load_addr; a < load_addr + home_size; a++) {
        uint8_t b = home->present[a] ? home->bytes[a] : (uint8_t)pad_byte;
        if (fwrite(&b, 1, 1, out) != 1) {
            fprintf(stderr, "mkexe: short write on HOME at 0x%04X\n", a);
            fclose(out); return 2;
        }
    }

    /* Banks — full 16 KB each, padding empty cells. */
    for (int i = 1; i <= max_bank; i++) {
        segment_t *bnk = &segments[i];
        for (uint32_t a = bank_base; a < bank_base + BANK_SIZE; a++) {
            uint8_t b = bnk->present[a] ? bnk->bytes[a] : (uint8_t)pad_byte;
            if (fwrite(&b, 1, 1, out) != 1) {
                fprintf(stderr, "mkexe: short write on bank %d at 0x%04X\n", i, a);
                fclose(out); return 2;
            }
        }
    }

    fclose(out);

    if (verbose) {
        fprintf(stderr,
            "mkexe: wrote %s\n"
            "  load=0x%04X start=0x%04X stack=0x%04X loader=0x%04X\n"
            "  HOME 0x%04X..0x%04X (%u bytes)\n",
            out_path, load_addr, start_addr, stack_addr, loader_size,
            home->lo, home->hi, home_size);
        for (int i = 1; i <= max_bank; i++) {
            fprintf(stderr, "  BANK%-2d 0x%04X..0x%04X (%u live bytes, padded to 16 KB)\n",
                    i, segments[i].lo, segments[i].hi,
                    segments[i].hi + 1 - segments[i].lo);
        }
        uint32_t total = EXE_HEADER_SIZE + home_size + (uint32_t)max_bank * BANK_SIZE;
        fprintf(stderr, "  total .exe = %u bytes\n", total);
    }
    return 0;
}