/* =========================================================================
 * MZFDTool — Sharp MZ-700 Floppy Disk image tool
 *
 * Creates and manages raw floppy disk images for the MZ-700 FDC system.
 * Supports format, directory listing, adding/extracting MZF files, and
 * setting a boot program.
 *
 * Original V1.01 (c) 2002 by BKK (DOS BIOS int13h version)
 * Updated  V2.00 (c) 2026 Philip Smart — rewritten for Linux with raw
 *          image file I/O, proper directory management, MZF conversion,
 *          and robust error handling.
 *
 * MZ-700 Floppy Disk format (327,680 bytes raw image):
 *
 *   Geometry: 40 cylinders, 2 heads, 16 sectors/track, 256 bytes/sector.
 *   Data on disk is INVERTED (XOR 0xFF) due to MZ-700 hardware.
 *
 *   Logical sector = (cylinder * 2 + head) * 16 + (phys_sector - 1)
 *   Total sectors: 40 * 2 * 16 = 1280
 *
 *   Boot sector:    Logical sector 0 (Cyl 0, Head 0, Sector 1)
 *   Directory:      Logical sectors 16-31 (Cyl 0, Head 1, Sectors 1-16)
 *                   16 sectors * 8 entries/sector = 128 entries
 *                   Entry 0 of sector 16 = volume/FAT header
 *   Data area:      Logical sector 32+ (Cyl 1+)
 *
 *   Boot sector (32 bytes used):
 *     +00: Type (03 = MZ-700 bootable)
 *     +01: Signature "IPLPRO" (6 bytes)
 *     +07: Program name (11 bytes, CR-terminated)
 *     +12: Load address (2 bytes LE)
 *     +14: File size (2 bytes LE)
 *     +16: Reserved (8 bytes)
 *     +1E: Start data sector (2 bytes LE)
 *
 *   Directory entry (32 bytes):
 *     +00: FileType (01=OBJ, 02=BTX, 00=deleted, >=80h=header)
 *     +01: FileName (17 bytes, CR-terminated, space-padded)
 *     +12: LockFlag (1 byte)
 *     +13: DummyFlag (1 byte)
 *     +14: FileSize (2 bytes LE, in bytes)
 *     +16: LoadAddress (2 bytes LE)
 *     +18: ExecAddress (2 bytes LE)
 *     +1A: Reserved (4 bytes)
 *     +1E: SectorAddress (2 bytes LE, logical)
 *
 * Usage:
 *   MZFDTool format  [-o disk.img]            Format empty disk image
 *   MZFDTool dir     [-o disk.img]            List directory
 *   MZFDTool add     <file.mzf> [-o disk.img] Add MZF file to directory
 *   MZFDTool extract <name> [-o disk.img]     Extract file to MZF
 *   MZFDTool boot    <file.mzf> [-o disk.img] Set boot program
 * ========================================================================= */

#define VERSION "2.00"

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

#define CYLS            40
#define HEADS           2
#define SECTORS         16
#define SECSIZE         256
#define TOTAL_SECTORS   (CYLS * HEADS * SECTORS)
#define IMGSIZE         (TOTAL_SECTORS * SECSIZE)
#define DIR_FIRST_SEC   16       /* First directory logical sector */
#define DIR_SECTORS     16       /* Number of directory sectors */
#define DIR_ENTRIES     8        /* Entries per sector */
#define DATA_FIRST_SEC  32       /* First data logical sector */
#define MAXFILETYPES    12
#define DEFAULTIMAGE    "MZ700.img"
#define CMTHDRSIZE      128      /* MZF/CMT file header size */
#define BOOT_SIG_TYPE   3        /* MZ-700 boot type */
#define BOOT_SIG        "IPLPRO"
#define BOOT_SIG_LEN    6

#pragma pack(push, 1)

/* Boot sector structure (first 32 bytes of logical sector 0) */
struct BootSector {
    uint8_t  type;               /* +00: 03 = MZ-700 bootable */
    char     signature[6];       /* +01: "IPLPRO" */
    char     name[11];           /* +07: Program name (CR-terminated) */
    uint16_t loadAddress;        /* +12: Load address */
    uint16_t fileSize;           /* +14: File size in bytes */
    uint8_t  reserved[8];        /* +16: Reserved */
    uint16_t sectorAddress;      /* +1E: Start data sector */
};

/* Directory entry (32 bytes) */
struct DirEntry {
    uint8_t  fileType;           /* +00: 01=OBJ, 02=BTX, 00=deleted */
    char     fileName[17];       /* +01: CR-terminated, space-padded */
    uint8_t  lockFlag;           /* +12: Lock flag */
    uint8_t  dummyFlag;          /* +13: Reserved */
    uint16_t fileSize;           /* +14: File size in bytes */
    uint16_t loadAddress;        /* +16: Load address */
    uint16_t execAddress;        /* +18: Exec address */
    uint8_t  dummy[4];           /* +1A: Reserved */
    uint16_t sectorAddress;      /* +1E: Start sector (logical) */
};

/* MZF/CMT tape file header (128 bytes) */
struct CMTHeader {
    uint8_t  attribute;          /* File type: 01=OBJ, 02=BTX, etc. */
    char     name[17];           /* Filename (CR-terminated) */
    uint16_t size;               /* Data size */
    uint16_t loadAddress;        /* Load address */
    uint16_t execAddress;        /* Execution address */
    char     comment[104];       /* Comment area */
};

#pragma pack(pop)

static const char *FileTypes[MAXFILETYPES] = {
    "???", "OBJ", "BTX", "BSD", "BRD", "RB ",
    "???", "LIB", "???", "???", "SYS", "GR "
};

static uint8_t  diskImage[IMGSIZE];
static char     imgFileName[256] = DEFAULTIMAGE;

/* ------- Sector I/O helpers ------- */

/* Convert logical sector to byte offset in raw image */
static int sec_offset(int logical_sec)
{
    /* Logical sector mapping:
     *   phys_sector = logical % 16 + 1  (but stored 0-based in image)
     *   head = (logical / 16) % 2
     *   cylinder = logical / 16 / 2
     * Raw image is sequential: cyl0/head0/sec1..16, cyl0/head1/sec1..16, ... */
    return logical_sec * SECSIZE;
}

/* Invert a buffer (XOR 0xFF) — MZ-700 data bus inversion */
static void invert_buf(uint8_t *buf, int len)
{
    for (int i = 0; i < len; i++)
        buf[i] ^= 0xFF;
}

/* Read a logical sector from image into buffer (un-inverted) */
static void read_sector(int logical_sec, uint8_t *buf)
{
    memcpy(buf, &diskImage[sec_offset(logical_sec)], SECSIZE);
    invert_buf(buf, SECSIZE);
}

/* Write a buffer to a logical sector in image (inverts before writing) */
static void write_sector(int logical_sec, const uint8_t *buf)
{
    memcpy(&diskImage[sec_offset(logical_sec)], buf, SECSIZE);
    invert_buf(&diskImage[sec_offset(logical_sec)], SECSIZE);
}

/* Load disk image from file */
static int load_image(void)
{
    FILE *f = fopen(imgFileName, "rb");
    if (!f) {
        fprintf(stderr, "ERROR: Cannot open '%s'\n", imgFileName);
        return 0;
    }
    size_t n = fread(diskImage, 1, IMGSIZE, f);
    fclose(f);
    if (n != IMGSIZE) {
        fprintf(stderr, "ERROR: '%s' is not a valid disk image (%zu bytes, expected %d)\n",
                imgFileName, n, IMGSIZE);
        return 0;
    }
    return 1;
}

/* Save disk image to file */
static int save_image(void)
{
    FILE *f = fopen(imgFileName, "wb");
    if (!f) {
        fprintf(stderr, "ERROR: Cannot write '%s'\n", imgFileName);
        return 0;
    }
    fwrite(diskImage, IMGSIZE, 1, f);
    fclose(f);
    return 1;
}

/* Format a filename for display: stop at CR, null-terminate */
static void format_name(char *dst, const char *src, int maxlen)
{
    memcpy(dst, src, maxlen);
    dst[maxlen] = '\0';
    for (int i = 0; i < maxlen; i++) {
        if (dst[i] == '\r' || dst[i] == '\0') { dst[i] = '\0'; break; }
    }
}

/* ------- Commands ------- */

/* Format: create an empty formatted disk image */
static void cmd_format(void)
{
    uint8_t sec[SECSIZE];

    printf("Formatting '%s' (%d bytes, %d sectors)\n", imgFileName, IMGSIZE, TOTAL_SECTORS);

    /* Fill entire image with 0xFF (which inverts to 0x00 = empty) */
    memset(diskImage, 0xFF, IMGSIZE);

    /* Write volume header in directory sector 16, entry 0 */
    memset(sec, 0, SECSIZE);
    sec[0] = 0x81;              /* Volume header marker */
    sec[1] = 'O';               /* Volume type */
    /* +14: capacity hint = 100 (not used by RFS) */
    sec[0x14] = 100;
    /* +1E: next free sector (first data sector) */
    sec[0x1E] = DATA_FIRST_SEC & 0xFF;
    sec[0x1F] = DATA_FIRST_SEC >> 8;
    write_sector(DIR_FIRST_SEC, sec);

    if (!save_image()) exit(1);
    puts("Done.");
}

/* Dir: list directory contents */
static void cmd_dir(void)
{
    uint8_t sec[SECSIZE];
    char name[18];
    int files = 0;

    if (!load_image()) exit(1);

    /* Check boot sector */
    read_sector(0, sec);
    struct BootSector *boot = (struct BootSector *)sec;
    if (boot->type == BOOT_SIG_TYPE && memcmp(boot->signature, BOOT_SIG, BOOT_SIG_LEN) == 0) {
        format_name(name, boot->name, 11);
        printf("\nBoot program: %-17s  Size=%u  Load=0x%04X  Sector=%u\n",
               name, boot->fileSize, boot->loadAddress, boot->sectorAddress);
    } else {
        printf("\nDisk is not bootable.\n");
    }

    printf("\nDirectory of '%s':\n\n", imgFileName);
    printf(" %-17s %-3s %5s %s %4s %4s %6s  (C H  S)\n",
           "Name", "Typ", "Size", "L", "Load", "Exec", "Sector");
    printf(" %-17s %-3s %5s %s %4s %4s %6s  %s\n",
           "-----------------", "---", "-----", "-", "----", "----", "------", "-------");

    for (int ds = 0; ds < DIR_SECTORS; ds++) {
        read_sector(DIR_FIRST_SEC + ds, sec);

        for (int e = 0; e < DIR_ENTRIES; e++) {
            /* Skip entry 0 of first directory sector (volume header) */
            if (ds == 0 && e == 0) continue;

            struct DirEntry *ent = (struct DirEntry *)&sec[e * sizeof(struct DirEntry)];
            if (ent->fileType == 0x00 || ent->fileType >= 0x80)
                continue;

            uint8_t ft = ent->fileType;
            if (ft >= MAXFILETYPES) ft = 0;
            format_name(name, ent->fileName, 17);

            int ls = ent->sectorAddress;
            int ps = ls % 16 + 1;
            int hd = (ls / 16) % 2;
            int cy = ls / 16 / 2;

            printf(" %-17s %s %5u %c %04X %04X  %5u  (%2d %d %2d)\n",
                   name, FileTypes[ft], ent->fileSize,
                   ent->lockFlag ? 'X' : ' ',
                   ent->loadAddress, ent->execAddress,
                   ent->sectorAddress, cy, hd, ps);
            files++;
        }
    }

    /* Calculate free space */
    int highSec = DATA_FIRST_SEC;
    for (int ds = 0; ds < DIR_SECTORS; ds++) {
        read_sector(DIR_FIRST_SEC + ds, sec);
        for (int e = 0; e < DIR_ENTRIES; e++) {
            struct DirEntry *ent = (struct DirEntry *)&sec[e * sizeof(struct DirEntry)];
            if (ent->fileType > 0 && ent->fileType < 0x80 && ent->sectorAddress > 0) {
                int endSec = ent->sectorAddress + (ent->fileSize + SECSIZE - 1) / SECSIZE;
                if (endSec > highSec) highSec = endSec;
            }
        }
    }
    printf("\n%d file(s), %d sectors used, %d sectors free\n",
           files, highSec, TOTAL_SECTORS - highSec);
}

/* Add: add an MZF file to the disk directory */
static void cmd_add(const char *mzfFileName)
{
    FILE *mzfFile;
    uint8_t sec[SECSIZE];
    struct CMTHeader cmtHdr;
    char name[18];

    if (!load_image()) exit(1);

    /* Open and read MZF file */
    mzfFile = fopen(mzfFileName, "rb");
    if (!mzfFile) {
        fprintf(stderr, "ERROR: Cannot open '%s'\n", mzfFileName);
        exit(1);
    }
    fseek(mzfFile, 0, SEEK_END);
    long mzfSize = ftell(mzfFile);
    fseek(mzfFile, 0, SEEK_SET);

    if (mzfSize < CMTHDRSIZE) {
        fprintf(stderr, "ERROR: '%s' too small for MZF format (%ld bytes)\n", mzfFileName, mzfSize);
        fclose(mzfFile);
        exit(1);
    }
    fread(&cmtHdr, CMTHDRSIZE, 1, mzfFile);
    uint16_t dataSize = cmtHdr.size;
    if (dataSize == 0) dataSize = (uint16_t)(mzfSize - CMTHDRSIZE);

    /* Find free directory entry and highest used sector */
    int freeDirSec = -1, freeDirIdx = -1;
    int highSec = DATA_FIRST_SEC;

    for (int ds = 0; ds < DIR_SECTORS; ds++) {
        read_sector(DIR_FIRST_SEC + ds, sec);
        for (int e = 0; e < DIR_ENTRIES; e++) {
            if (ds == 0 && e == 0) continue; /* Skip volume header */
            struct DirEntry *ent = (struct DirEntry *)&sec[e * sizeof(struct DirEntry)];

            if (ent->fileType == 0x00 && freeDirSec < 0) {
                freeDirSec = ds;
                freeDirIdx = e;
            }
            if (ent->fileType > 0 && ent->fileType < 0x80 && ent->sectorAddress > 0) {
                int endSec = ent->sectorAddress + (ent->fileSize + SECSIZE - 1) / SECSIZE;
                if (endSec > highSec) highSec = endSec;
            }
        }
    }

    if (freeDirSec < 0) {
        fprintf(stderr, "ERROR: Directory full\n");
        fclose(mzfFile);
        exit(1);
    }

    int dataSectors = (dataSize + SECSIZE - 1) / SECSIZE;
    if (highSec + dataSectors > TOTAL_SECTORS) {
        fprintf(stderr, "ERROR: Not enough space (need %d sectors, have %d)\n",
                dataSectors, TOTAL_SECTORS - highSec);
        fclose(mzfFile);
        exit(1);
    }

    /* Write file data sectors */
    for (int s = 0; s < dataSectors; s++) {
        memset(sec, 0, SECSIZE);
        fread(sec, 1, SECSIZE, mzfFile);
        write_sector(highSec + s, sec);
    }
    fclose(mzfFile);

    /* Update directory entry */
    read_sector(DIR_FIRST_SEC + freeDirSec, sec);
    struct DirEntry *ent = (struct DirEntry *)&sec[freeDirIdx * sizeof(struct DirEntry)];
    memset(ent, 0, sizeof(struct DirEntry));

    ent->fileType = cmtHdr.attribute;
    if (ent->fileType == 0x05) ent->fileType = 0x02;  /* CMT type 05 → FD type 02 (BTX) */
    memcpy(ent->fileName, cmtHdr.name, 17);
    ent->fileSize = dataSize;
    ent->loadAddress = cmtHdr.loadAddress;
    ent->execAddress = cmtHdr.execAddress;
    ent->sectorAddress = highSec;
    write_sector(DIR_FIRST_SEC + freeDirSec, sec);

    format_name(name, cmtHdr.name, 17);
    printf("  Added: \"%s\"  type=%s  size=%u  load=0x%04X  exec=0x%04X  sector=%d\n",
           name, FileTypes[ent->fileType < MAXFILETYPES ? ent->fileType : 0],
           dataSize, cmtHdr.loadAddress, cmtHdr.execAddress, highSec);
    printf("  %d sectors free\n", TOTAL_SECTORS - highSec - dataSectors);

    if (!save_image()) exit(1);
}

/* Extract: extract a file from disk to MZF format */
static void cmd_extract(const char *fileName)
{
    uint8_t sec[SECSIZE];
    char entName[18], outName[256];

    if (!load_image()) exit(1);

    /* Search directory for matching filename */
    static struct DirEntry foundEntry;
    int foundIt = 0;

    for (int ds = 0; ds < DIR_SECTORS && !foundIt; ds++) {
        read_sector(DIR_FIRST_SEC + ds, sec);
        for (int e = 0; e < DIR_ENTRIES && !foundIt; e++) {
            if (ds == 0 && e == 0) continue;
            struct DirEntry *ent = (struct DirEntry *)&sec[e * sizeof(struct DirEntry)];
            if (ent->fileType > 0 && ent->fileType < 0x80) {
                format_name(entName, ent->fileName, 17);
                if (strcmp(fileName, entName) == 0) {
                    memcpy(&foundEntry, ent, sizeof(foundEntry));
                    foundIt = 1;
                }
            }
        }
    }

    if (!foundIt) {
        /* Also check boot sector */
        read_sector(0, sec);
        struct BootSector *boot = (struct BootSector *)sec;
        if (boot->type == BOOT_SIG_TYPE && memcmp(boot->signature, BOOT_SIG, BOOT_SIG_LEN) == 0) {
            format_name(entName, boot->name, 11);
            if (strcmp(fileName, entName) == 0) {
                memset(&foundEntry, 0, sizeof(foundEntry));
                foundEntry.fileType = 0x01;
                memcpy(foundEntry.fileName, boot->name, 11);
                foundEntry.fileSize = boot->fileSize;
                foundEntry.loadAddress = boot->loadAddress;
                foundEntry.sectorAddress = boot->sectorAddress;
                foundIt = 1;
            }
        }
    }

    if (!foundIt) {
        fprintf(stderr, "ERROR: '%s' not found on disk\n", fileName);
        exit(1);
    }

    struct DirEntry *found = &foundEntry;

    /* Build output filename */
    snprintf(outName, sizeof(outName), "%s.mzf", fileName);
    for (char *p = outName; *p; p++) {
        if (*p == ' ') *p = '_';
    }

    /* Write MZF file */
    FILE *out = fopen(outName, "wb");
    if (!out) {
        fprintf(stderr, "ERROR: Cannot create '%s'\n", outName);
        exit(1);
    }

    struct CMTHeader hdr;
    memset(&hdr, 0, sizeof(hdr));
    hdr.attribute = (found->fileType == 0x02) ? 0x05 : found->fileType;
    memcpy(hdr.name, found->fileName, 17);
    hdr.size = found->fileSize;
    hdr.loadAddress = found->loadAddress;
    hdr.execAddress = found->execAddress;
    snprintf(hdr.comment, sizeof(hdr.comment), "Extracted by MZFDTool V%s", VERSION);
    fwrite(&hdr, sizeof(hdr), 1, out);

    /* Read and write data sectors */
    uint16_t remaining = found->fileSize;
    int curSec = found->sectorAddress;
    while (remaining > 0) {
        read_sector(curSec, sec);
        uint16_t chunk = (remaining > SECSIZE) ? SECSIZE : remaining;
        fwrite(sec, 1, chunk, out);
        remaining -= chunk;
        curSec++;
    }

    fclose(out);
    format_name(entName, found->fileName, 17);
    printf("  Extracted: \"%s\" → %s (%u bytes)\n", entName, outName, found->fileSize);
}

/* Boot: set boot program from MZF file */
static void cmd_boot(const char *mzfFileName)
{
    FILE *mzfFile;
    uint8_t sec[SECSIZE];
    struct CMTHeader cmtHdr;
    char name[18];

    if (!load_image()) exit(1);

    mzfFile = fopen(mzfFileName, "rb");
    if (!mzfFile) {
        fprintf(stderr, "ERROR: Cannot open '%s'\n", mzfFileName);
        exit(1);
    }
    fseek(mzfFile, 0, SEEK_END);
    long mzfSize = ftell(mzfFile);
    fseek(mzfFile, 0, SEEK_SET);

    if (mzfSize < CMTHDRSIZE) {
        fprintf(stderr, "ERROR: '%s' too small for MZF format\n", mzfFileName);
        fclose(mzfFile);
        exit(1);
    }
    fread(&cmtHdr, CMTHDRSIZE, 1, mzfFile);
    uint16_t dataSize = cmtHdr.size;
    if (dataSize == 0) dataSize = (uint16_t)(mzfSize - CMTHDRSIZE);

    /* Find space for boot data — use sectors 1-15 (logical), before directory */
    int dataSectors = (dataSize + SECSIZE - 1) / SECSIZE;
    int bootDataStart = 1; /* Sector 1 (sector 0 is the boot sector itself) */
    if (dataSectors > 15) {
        fprintf(stderr, "ERROR: Boot program too large (%d sectors, max 15)\n", dataSectors);
        fclose(mzfFile);
        exit(1);
    }

    /* Write boot data to sectors 1-15 */
    for (int s = 0; s < dataSectors; s++) {
        memset(sec, 0, SECSIZE);
        fread(sec, 1, SECSIZE, mzfFile);
        write_sector(bootDataStart + s, sec);
    }
    fclose(mzfFile);

    /* Write boot sector (sector 0) */
    memset(sec, 0, SECSIZE);
    struct BootSector *boot = (struct BootSector *)sec;
    boot->type = BOOT_SIG_TYPE;
    memcpy(boot->signature, BOOT_SIG, BOOT_SIG_LEN);
    format_name(name, cmtHdr.name, 11);
    memcpy(boot->name, cmtHdr.name, 11);
    boot->loadAddress = cmtHdr.loadAddress;
    boot->fileSize = dataSize;
    boot->sectorAddress = bootDataStart;
    write_sector(0, sec);

    format_name(name, cmtHdr.name, 11);
    printf("  Boot set: \"%s\"  size=%u  load=0x%04X  exec=0x%04X  sector=%d\n",
           name, dataSize, cmtHdr.loadAddress, cmtHdr.execAddress, bootDataStart);

    if (!save_image()) exit(1);
}

/* ------- Usage and main ------- */

static void usage(void)
{
    printf("Usage:\n");
    printf("  MZFDTool format  [-o disk.img]             Format empty disk image\n");
    printf("  MZFDTool dir     [-o disk.img]             List directory\n");
    printf("  MZFDTool add     <file.mzf> [-o disk.img]  Add MZF file to disk\n");
    printf("  MZFDTool extract <name> [-o disk.img]      Extract file to MZF\n");
    printf("  MZFDTool boot    <file.mzf> [-o disk.img]  Set boot program\n");
    printf("\n");
    printf("Options:\n");
    printf("  -o <filename>   Disk image file (default: %s)\n", DEFAULTIMAGE);
    printf("\n");
    printf("Disk geometry: %d cyls, %d heads, %d sectors, %d bytes/sector = %d bytes\n",
           CYLS, HEADS, SECTORS, SECSIZE, IMGSIZE);
    printf("\n");
}

int main(int argc, char *argv[])
{
    printf("\nMZFDTool V%s  (c) 2002 BKK, 2026 Philip Smart\n\n", VERSION);

    if (argc < 2) {
        usage();
        return 1;
    }

    /* Parse -o option from any position */
    for (int i = 1; i < argc - 1; i++) {
        if (strcmp("-o", argv[i]) == 0) {
            strncpy(imgFileName, argv[i + 1], sizeof(imgFileName) - 1);
            imgFileName[sizeof(imgFileName) - 1] = '\0';
            for (int j = i; j < argc - 2; j++)
                argv[j] = argv[j + 2];
            argc -= 2;
            break;
        }
    }

    if (strcmp("format", argv[1]) == 0) {
        cmd_format();
    } else if (strcmp("dir", argv[1]) == 0) {
        cmd_dir();
    } else if (strcmp("add", argv[1]) == 0) {
        if (argc < 3) {
            fprintf(stderr, "ERROR: No MZF file specified\n");
            return 2;
        }
        cmd_add(argv[2]);
    } else if (strcmp("extract", argv[1]) == 0) {
        if (argc < 3) {
            fprintf(stderr, "ERROR: No filename specified\n");
            return 2;
        }
        cmd_extract(argv[2]);
    } else if (strcmp("boot", argv[1]) == 0) {
        if (argc < 3) {
            fprintf(stderr, "ERROR: No MZF file specified\n");
            return 2;
        }
        cmd_boot(argv[2]);
    } else {
        fprintf(stderr, "ERROR: Unknown command '%s'\n", argv[1]);
        usage();
        return 1;
    }

    return 0;
}