#define _POSIX_C_SOURCE 2
#include "lzw.h"
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define MAGIC_DSK_OLD 0x58aa
#define MAGIC_DSK_UNCOMPRESSED 0x59aa
#define MAGIC_DSK_COMPRESSED 0x5aaa

#if defined(__DOS__) || defined(__WINDOWS__) || defined(__NT__) || defined(__OS2__)
#define DOSLIKE
#define DIRSEP '\\'
#else
#undef DOSLIKE
#define DIRSEP '/'
#endif

#pragma pack(1)
struct dskheader {
    uint16_t magic; /* magic identifier (0x58aa, 0x59aa or 0x5aaa) */
    uint16_t mediatype; /* disk media type as defined in the first byte of the FAT */
    uint16_t sectorsize; /* sector size in bytes */
    uint8_t clustermask; /* number of sectors per cluster minus one */
    uint8_t clustershift; /* log2(cluster size / sector size) */
    uint16_t reservedsectors; /* number of reserved sectors at the start of the disk */
    uint8_t fatcopies; /* number of copies of the FAT */
    uint16_t rootentries; /* max number of entries in the root directory */
    uint16_t firstclustersector; /* count of sectors for boot sector + FATs + root dir */
    uint16_t imageclusters; /* number of clusters in the image (empty clusters at the end are not saved) */
    uint8_t sectorsperfat; /* size of each FAT copy in sectors */
    uint16_t rootdirsector; /* count of sectors used by boot sector + FATs */
    uint32_t checksum; /* sum of all words in the image */
    uint16_t cylinders; /* number of disk cylinders (40 or 80) */
    uint16_t heads; /* number of disk heads (sides, 1 or 2) */
    uint16_t sectors; /* number of sectors per track */
    uint32_t unused;
    uint16_t imagesectors; /* number of sectors in the image */
    uint16_t commentoffset; /* offset of the image comment string */
    uint16_t firstoffset; /* offset of the start of image contents */
};
#pragma pack()

const struct media_type {
    int sectorsize;
    int cylinders;
    int heads;
    int sectors;
    char* name;
} media_types[] = {
    { 512, 80, 2, 36, "3.5\" DSED 2.88MB" },
    { 512, 80, 2, 18, "3.5\" DSHD 1.44MB" },
    { 512, 80, 2, 9, "3.5\" DSDD 720KB" },
    { 512, 80, 2, 15, "5.25\" DSHD 1.2MB" },
    { 512, 40, 1, 8, "5.25\" SSDD 160KB" },
    { 512, 40, 1, 9, "5.25\" SSDD 180KB" },
    { 512, 40, 2, 8, "5.25\" DSDD 320KB" },
    { 512, 40, 2, 9, "5.25\" DSDD 360KB" },
    { 512, 80, 1, 8, "5.25\" SSQD 320KB" },
    { 512, 80, 2, 8, "5.25\" DSQD 640KB" },
    { 512, 80, 1, 8, "3.5\" SSDD 320KB" },
    { 512, 80, 1, 9, "3.5\" SSDD 360KB" },
    { 512, 80, 2, 8, "3.5\" DSDD 640KB" },
    { 512, 80, 2, 21, "3.5\" DSHD 1.68MB DMF" },
    { 512, 80, 2, 21, "3.5\" DSHD 1.72MB DMF" },
    { -1, -1, -1, -1, "Unknown" }
};

char* guess_output_filename(const char* input_filename, int is_compressed)
{
    char* outname;
    unsigned int ifnlen, ofnlen;
    int i;

    /* NOLINTBEGIN(clang-analyzer-security.insecureAPI.DeprecatedOrUnsafeBufferHandling) */
    ifnlen = strlen(input_filename);
    ofnlen = ifnlen + 9;
    outname = (char*)malloc(ofnlen);
    strncpy(outname, input_filename, ofnlen);
    /* strip extension */
    for (i = ifnlen; i > 0; i--) {
        if (outname[i - 1] == '.') {
            outname[i - 1] = '\0';
            break;
        } else if (outname[i - 1] == DIRSEP)
            break;
    }
#ifdef DOSLIKE
    if (is_compressed) /* use filename.lzw instead of filename.img.lzw on dos and win */
        strcat(outname, ".lzw");
    else
        strcat(outname, ".img");
#else
    strncat(outname, ".img", ofnlen);
    if (is_compressed)
        strncat(outname, ".lzw", ofnlen);
#endif
    /* NOLINTEND(clang-analyzer-security.insecureAPI.DeprecatedOrUnsafeBufferHandling) */

    return outname;
}

void pad_file_to_size(FILE* f, unsigned long size)
{
    long r;
    unsigned long current_pos;

    r = ftell(f);
    if (r < 0) {
        perror("ftell()");
    }
    current_pos = r;
    if (current_pos > size)
        printf("ERROR: output file is larger than expected! (%lu vs %lu)\n", current_pos, size);
    if (current_pos < size - 1) {
        long pos;

        pos = ((long)size) - 1;
        fseek(f, pos, SEEK_SET);
    }
    if (current_pos < size)
        fputc(0, f);
}

int valid_header(const struct dskheader* h)
{
    unsigned int final_sectors;
    unsigned long int final_bytes;

    if (!(h->magic == MAGIC_DSK_COMPRESSED || h->magic == MAGIC_DSK_UNCOMPRESSED || h->magic == MAGIC_DSK_OLD))
        return 0;

    if (h->cylinders < 40 || h->cylinders > 80)
        return 0;
    if (h->heads == 0 || h->heads > 2)
        return 0;
    if (h->sectors < 8 || h->sectors > 36)
        return 0;
    /* reasonable floppy sector sizes are powers of two between 128 and 1024 */
    if (h->sectorsize < 128 || h->sectorsize > 1024 || (h->sectorsize & (h->sectorsize - 1)) != 0)
        return 0;

    final_sectors = h->cylinders * h->heads * h->sectors;
    final_bytes = final_sectors * h->sectorsize;

    if (final_bytes > 4000000UL)
        return 0;

    if ((h->clustershift > 7) || (h->clustermask != (1UL << h->clustershift) - 1))
        return 0;
    if (h->reservedsectors + h->fatcopies * h->sectorsperfat != h->rootdirsector)
        return 0;
    if (h->firstclustersector - 1 + (h->imageclusters << h->clustershift) != final_sectors)
        return 0;
    if (h->firstclustersector != h->reservedsectors + h->fatcopies * h->sectorsperfat + h->rootentries * 32 / h->sectorsize)
        return 0;
    if (h->imagesectors > final_sectors)
        return 0;
    if (h->commentoffset < sizeof(struct dskheader) || h->firstoffset <= h->commentoffset)
        return 0;

    return 1;
}

#if 0
void dump_dsk_header(struct dskheader* h)
{
    printf("=== header =====================================\n");
    printf("magic: 0x%04x\n", h->magic);
    printf("media type: 0x%04x\n", h->mediatype);
    printf("sectorsize: %d\n", h->sectorsize);
    printf("clustermask: %d\n", h->clustermask);
    printf("clustershift: %d\n", h->clustershift);
    printf("reservedsectors: %d\n", h->reservedsectors);
    printf("fatcopies: %d\n", h->fatcopies);
    printf("rootentries: %d\n", h->rootentries);
    printf("firstclustersector: %d\n", h->firstclustersector);
    printf("imageclusters: %d\n", h->imageclusters);
    printf("sectorsperfat: %d\n", h->sectorsperfat);
    printf("rootdirsector: %d\n", h->rootdirsector);
    printf("checksum: 0x%08x\n", h->checksum);
    printf("cylinders: %d\n", h->cylinders);
    printf("heads: %d\n", h->heads);
    printf("sectors: %d\n", h->sectors);
    printf("unused: 0x%08x\n", h->unused);
    printf("imagesectors: %d\n", h->imagesectors);
    printf("commentoffset: %d\n", h->commentoffset);
    printf("firstoffset: %d\n", h->firstoffset);
    printf("================================================\n");
}
#endif

void dsk2img(const char* filename, int no_lzw)
{
    FILE *inf, *outf;
    struct dskheader header;
    char comment[600];
    char *c, *outname;
    unsigned long rres;
    int sres;
    unsigned long int final_sectors, image_sectors, sector;
    unsigned long int final_size;
    long int start_offset;
    uint8_t* buffer;
    uint32_t checksum;
    int i;
    int type_geom;
    char *rbuf = NULL, *wbuf = NULL;

    inf = fopen(filename, "rb");
    if (inf == NULL)
        return;

    rbuf = (char*)malloc(4096);
    if (rbuf)
        setvbuf(inf, rbuf, _IOFBF, 4096);
    else
        printf("WARNING: can't allocate read buffer.");

    rres = fread(&header, sizeof(header), 1, inf);
    if (rres != 1) {
        printf("Short read\n");
        fclose(inf);
        if (rbuf)
            free(rbuf);
        return;
    }

    /* dump_dsk_header(&header); */

    if (!valid_header(&header)) {
        fclose(inf);
        puts("Not a valid DSK file!");
        if (rbuf)
            free(rbuf);
        return;
    }

    final_sectors = header.cylinders * header.heads * header.sectors;
    final_size = final_sectors * header.sectorsize;

    printf("DSK format: ");
    switch (header.magic) {
    case MAGIC_DSK_OLD:
        puts("Old");
        break;
    case MAGIC_DSK_UNCOMPRESSED:
        puts("Regular");
        break;
    case MAGIC_DSK_COMPRESSED:
        puts("Compressed");
        break;
    default:
        puts("Unknown");
    }
    for (type_geom = 0; media_types[type_geom].sectors > 0; type_geom++)
        if ((media_types[type_geom].cylinders == header.cylinders) && (media_types[type_geom].heads == header.heads) && (media_types[type_geom].sectors == header.sectors) && (media_types[type_geom].sectorsize == header.sectorsize))
            break;
    printf("C=%d H=%d S=%d, %lu sectors of %dB (total %luB) [%s]\n", header.cylinders, header.heads, header.sectors, final_sectors, header.sectorsize, (unsigned long)final_size, media_types[type_geom].name);

    sres = fseek(inf, header.commentoffset, SEEK_SET);
    if (sres != 0) {
        perror("comment seek()");
        fclose(inf);
        if (rbuf)
            free(rbuf);
        return;
    }

    c = fgets(comment, sizeof(comment), inf);
    if (c == NULL) {
        perror("fgets()");
        fclose(inf);
        if (rbuf)
            free(rbuf);
        return;
    }
    if (strlen(c) > 0)
        printf("comment: %s", c);

    start_offset = header.firstoffset ? header.firstoffset : 0x200;
    sres = fseek(inf, start_offset, SEEK_SET);
    if (sres != 0) {
        perror("start seek()");
        fclose(inf);
        if (rbuf)
            free(rbuf);
        return;
    }

    outname = guess_output_filename(filename, (header.magic == MAGIC_DSK_COMPRESSED) && (no_lzw != 0));

    outf = fopen(outname, "wb+");
    if (outf == NULL) {
        perror("fopen()");
        fclose(inf);
        free(outname);
        if (rbuf)
            free(rbuf);
        return;
    }

    wbuf = (char*)malloc(4096);
    if (wbuf)
        setvbuf(outf, wbuf, _IOFBF, 4096);
    else
        printf("WARNING: can't allocate write buffer.");

    buffer = (uint8_t*)malloc(header.sectorsize);
    if (buffer == NULL) {
        perror("malloc()");
        fclose(outf);
        fclose(inf);
        free(outname);
        if (rbuf)
            free(rbuf);
        if (wbuf)
            free(wbuf);
        return;
    }

    /* image_sectors = (header.imageclusters << header.clustershift) +
       header.firstclustersector - 1; */
    image_sectors = header.imagesectors;
    checksum = 0;
    if (header.magic != MAGIC_DSK_COMPRESSED) {
        for (sector = 0; sector < image_sectors; sector++) {
            rres = fread(buffer, header.sectorsize, 1, inf);
            if (rres != 1) {
                perror("fread()");
                fclose(outf);
                fclose(inf);
                free(buffer);
                free(outname);
                if (rbuf)
                    free(rbuf);
                if (wbuf)
                    free(wbuf);
                return;
            }
            for (i = 0; i < header.sectorsize / 2; i++)
                checksum += buffer[i * 2] + 256 * buffer[i * 2 + 1];
            rres = fwrite(buffer, header.sectorsize, 1, outf);
            if (rres != 1) {
                perror("fwrite()");
                fclose(outf);
                fclose(inf);
                free(buffer);
                free(outname);
                if (rbuf)
                    free(rbuf);
                if (wbuf)
                    free(wbuf);
                return;
            }
        }
        pad_file_to_size(outf, final_size);
        if (checksum != header.checksum)
            printf("ERROR: image checksum does not match. Expected 0x%08x but got 0x%08x\n", header.checksum, checksum);
        else
            puts("Done.");
    } else {
        if (no_lzw) {
            size_t wres;

            for (sector = 0;; sector++) {
                rres = fread(buffer, 1, header.sectorsize, inf);
                for (i = 0; i < rres; i++)
                    checksum += buffer[i];
                wres = fwrite(buffer, 1, rres, outf);
                if (wres != rres) {
                    perror("fwrite()");
                    fclose(outf);
                    fclose(inf);
                    free(buffer);
                    free(outname);
                    if (rbuf)
                        free(rbuf);
                    if (wbuf)
                        free(wbuf);
                    return;
                }
                if (rres < header.sectorsize)
                    break;
            }
        } else {
            /* TODO: checksum */
            lzw_decompress_file(inf, outf);
            pad_file_to_size(outf, final_size);
        }
    }

    fclose(outf);
    fclose(inf);
    free(buffer);
    free(outname);
    if (rbuf)
        free(rbuf);
    if (wbuf)
        free(wbuf);
}

void help(void)
{
    puts("Syntax: dsk2img [-h] [-n] [FILE...]");
    puts("\t-n\tDo not attempt to decompress compressed images");
    puts("\t-h\tShow this help");
}

int main(int argc, char* const argv[])
{
    int i;
    int no_lzw = 0;
    int c;

    while ((c = getopt(argc, argv, "hn")) != -1)
        switch (c) {
        case 'n':
            no_lzw = 1;
            break;
        case '?':
        case 'h':
            help();
            return EXIT_SUCCESS;
        }
    if (optind >= argc)
        help();
    else
        for (i = optind; i < argc; i++)
            /* TODO: handle DOS file globbing */
            dsk2img(argv[i], no_lzw);

    return EXIT_SUCCESS;
}