// .nib -> .dsk converter for Apple II 5.25" disks.
//
// The .nib format stores 6656 raw nibbles per track (35 tracks total =
// 232,960 bytes) with the sector framing intact: D5 AA 96 marks the
// address field, D5 AA AD marks the data field, and the data field
// holds 342 6-and-2 encoded nibbles plus a checksum that decodes back
// to 256 bytes per sector.
//
// .dsk is just 256 * 16 * 35 = 143,360 bytes of decoded sector data
// in DOS 3.3 sector-interleave order (this tool emits DOS order,
// matching the .DSK conventions on Asimov; pass `--prodos` to emit
// the ProDOS .po block order instead).

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

#define NIB_BYTES_PER_TRACK   6656
#define DSK_BYTES_PER_TRACK   4096
#define TRACKS                35
#define SECTORS               16


// 6-and-2 decode lookup: maps an Apple disk nibble (0x96..0xFF) to its
// 6-bit value (0x00..0x3F) or 0xFF if the nibble is invalid.
static uint8_t  decTable[256];
static uint8_t  dosToProDos[16] = { 0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15 };
static uint8_t  dosToPhysical[16] = { 0, 7, 14, 6, 13, 5, 12, 4, 11, 3, 10, 2, 9, 1, 8, 15 };


static void buildDecTable(void);
static int  decodeTrack(const uint8_t *nibTrack, uint8_t outSectors[SECTORS][256]);
static int  findMarker(const uint8_t *buf, int *pos, int len, uint8_t a, uint8_t b, uint8_t c);


static void buildDecTable(void) {
        static const uint8_t physicalCodes[64] = {
                0x96, 0x97, 0x9A, 0x9B, 0x9D, 0x9E, 0x9F, 0xA6,
                0xA7, 0xAB, 0xAC, 0xAD, 0xAE, 0xAF, 0xB2, 0xB3,
                0xB4, 0xB5, 0xB6, 0xB7, 0xB9, 0xBA, 0xBB, 0xBC,
                0xBD, 0xBE, 0xBF, 0xCB, 0xCD, 0xCE, 0xCF, 0xD3,
                0xD6, 0xD7, 0xD9, 0xDA, 0xDB, 0xDC, 0xDD, 0xDE,
                0xDF, 0xE5, 0xE6, 0xE7, 0xE9, 0xEA, 0xEB, 0xEC,
                0xED, 0xEE, 0xEF, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6,
                0xF7, 0xF9, 0xFA, 0xFB, 0xFC, 0xFD, 0xFE, 0xFF
        };
        for (int i = 0; i < 256; i++) {
                decTable[i] = 0xFF;
        }
        for (int i = 0; i < 64; i++) {
                decTable[physicalCodes[i]] = (uint8_t)i;
        }
}


static int findMarker(const uint8_t *buf, int *pos, int len, uint8_t a, uint8_t b, uint8_t c) {
        for (int p = *pos; p + 2 < len; p++) {
                if (buf[p] == a && buf[p + 1] == b && buf[p + 2] == c) {
                        *pos = p + 3;
                        return 1;
                }
        }
        return 0;
}


// Read 8 4-and-4 encoded nibbles -> 4 bytes (volume, track, sector,
// checksum). Returns 1 on success, 0 if there isn't enough buffer.
static int read44(const uint8_t *buf, int *pos, int len, uint8_t out[4]) {
        if (*pos + 8 > len) {
                return 0;
        }
        for (int i = 0; i < 4; i++) {
                uint8_t hi = buf[*pos + i * 2];
                uint8_t lo = buf[*pos + i * 2 + 1];
                out[i] = (uint8_t)(((hi << 1) | 1) & lo);
        }
        *pos += 8;
        return 1;
}


// Decode a 6-and-2 data field at `*pos`. Outputs 256 bytes, advances
// `*pos` past the consumed nibbles. Returns 1 on success, 0 on
// checksum mismatch or truncation.
static int decodeDataField(const uint8_t *buf, int *pos, int len, uint8_t out[256]) {
        if (*pos + 343 > len) {
                return 0;
        }
        uint8_t  twos[86];
        uint8_t  sixes[256];

        // 86 lower-2-bit nibbles followed by 256 upper-6-bit nibbles
        // followed by 1 checksum nibble. Each nibble is XORed with
        // the running checksum during decode.
        uint8_t running = 0;
        for (int i = 85; i >= 0; i--) {
                uint8_t nib = decTable[buf[(*pos)++]];
                if (nib == 0xFF) {
                        return 0;
                }
                running ^= nib;
                twos[i] = running;
        }
        for (int i = 0; i < 256; i++) {
                uint8_t nib = decTable[buf[(*pos)++]];
                if (nib == 0xFF) {
                        return 0;
                }
                running ^= nib;
                sixes[i] = running;
        }
        uint8_t cksum = decTable[buf[(*pos)++]];
        if (cksum == 0xFF || running != cksum) {
                return 0;
        }

        // Combine the 6-bit "sixes" with the 2-bit "twos" (3 sectors of
        // bit-pair stream; bit-pair n of a sixes byte goes to byte n,
        // n+86, or n+172).
        for (int i = 0; i < 256; i++) {
                uint8_t two_idx = (uint8_t)(i % 86);
                uint8_t two_pos = (uint8_t)(i / 86);
                uint8_t low2    = (uint8_t)((twos[two_idx] >> (two_pos * 2)) & 0x03);
                // Bit-reverse the 2-bit nibble (Apple writes the high
                // bit of each pair first).
                low2 = (uint8_t)(((low2 & 1) << 1) | ((low2 >> 1) & 1));
                out[i] = (uint8_t)((sixes[i] << 2) | low2);
        }
        return 1;
}


static int decodeTrack(const uint8_t *nibTrack, uint8_t outSectors[SECTORS][256]) {
        int pos     = 0;
        int decoded = 0;
        uint8_t haveSector[SECTORS] = { 0 };
        uint8_t addr[4];

        while (pos < NIB_BYTES_PER_TRACK && decoded < SECTORS) {
                if (!findMarker(nibTrack, &pos, NIB_BYTES_PER_TRACK, 0xD5, 0xAA, 0x96)) {
                        break;
                }
                if (!read44(nibTrack, &pos, NIB_BYTES_PER_TRACK, addr)) {
                        break;
                }
                if (addr[2] >= SECTORS) {
                        continue;
                }
                if (haveSector[addr[2]]) {
                        continue;
                }
                if (!findMarker(nibTrack, &pos, NIB_BYTES_PER_TRACK, 0xD5, 0xAA, 0xAD)) {
                        break;
                }
                if (!decodeDataField(nibTrack, &pos, NIB_BYTES_PER_TRACK, outSectors[addr[2]])) {
                        continue;
                }
                haveSector[addr[2]] = 1;
                decoded++;
        }
        return decoded;
}


int main(int argc, char **argv) {
        const uint8_t *interleave = dosToPhysical;
        int            argIndex   = 1;

        if (argIndex < argc && strcmp(argv[argIndex], "--prodos") == 0) {
                interleave = dosToProDos;
                argIndex++;
        }
        if (argc - argIndex != 2) {
                fprintf(stderr, "usage: %s [--prodos] in.nib out.dsk\n", argv[0]);
                return 1;
        }

        const char *inPath  = argv[argIndex];
        const char *outPath = argv[argIndex + 1];

        FILE *f = fopen(inPath, "rb");
        if (f == NULL) {
                fprintf(stderr, "cannot open %s\n", inPath);
                return 1;
        }
        uint8_t *nib = malloc(TRACKS * NIB_BYTES_PER_TRACK);
        if (nib == NULL) {
                fclose(f);
                return 1;
        }
        size_t nRead = fread(nib, 1, TRACKS * NIB_BYTES_PER_TRACK, f);
        fclose(f);
        if (nRead != TRACKS * NIB_BYTES_PER_TRACK) {
                fprintf(stderr, "short read: got %zu bytes (expected %d)\n",
                        nRead, TRACKS * NIB_BYTES_PER_TRACK);
                free(nib);
                return 1;
        }

        buildDecTable();

        uint8_t (*track)[256] = malloc(SECTORS * 256);
        if (track == NULL) {
                free(nib);
                return 1;
        }
        FILE *out = fopen(outPath, "wb");
        if (out == NULL) {
                fprintf(stderr, "cannot open %s\n", outPath);
                free(nib);
                free(track);
                return 1;
        }

        int totalDecoded = 0;
        for (int t = 0; t < TRACKS; t++) {
                memset(track, 0, SECTORS * 256);
                int got = decodeTrack(&nib[t * NIB_BYTES_PER_TRACK], (uint8_t (*)[256])track);
                totalDecoded += got;
                for (int dosSector = 0; dosSector < SECTORS; dosSector++) {
                        int physSector = interleave[dosSector];
                        fwrite(track[physSector], 1, 256, out);
                }
                if (got < SECTORS) {
                        fprintf(stderr, "track %2d: %d/%d sectors decoded\n", t, got, SECTORS);
                }
        }
        fclose(out);
        free(nib);
        free(track);

        printf("decoded %d / %d sectors -> %s\n", totalDecoded, TRACKS * SECTORS, outPath);
        return totalDecoded == TRACKS * SECTORS ? 0 : 1;
}