#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "ImageGifc.h"

/***
 * Routines to write a GIF image file.
 * After Sverre Huseby's GifSave, 1992
 */


typedef unsigned long ulong;

typedef unsigned int Word;      /* At least two bytes (16 bits) */
typedef unsigned char Byte;     /* Exactly one byte (8 bits) */


#define RES_CODES 2

#define MAXBITS 12
#define MAXSTR (1 << MAXBITS)

#define HASHSIZE 9973
#define HASHSTEP 2039

#define HASH(index, lastbyte) (((lastbyte << 8) ^ index) % HASHSIZE)

struct bitfile
{
    unsigned char buf[256];
    int idx;
    int bitsLeft;
    FILE *fp;
};

struct strtab
{
    unsigned char *Chr;
    int *Nxt;
    int *Hsh;
    int NumStrings;
};


static void InitBitFile(struct bitfile *p, FILE *fp)
{
    p->idx = 0;
    p->buf[p->idx] = 0;
    p->bitsLeft = 8;
    p->fp = fp;
}


static int ResetOutBitFile(struct bitfile *p)
{
    unsigned numbytes = p->idx + (p->bitsLeft != 8);

    if (numbytes) {
        if (putc(numbytes, p->fp) == EOF)
          return -1;
        if (fwrite(p->buf, 1, numbytes, p->fp) != numbytes)
          return -1;
        InitBitFile(p, p->fp);
    }

    return 0;
}


static int WriteBits(struct bitfile *p, unsigned int bits,
                     int nbits)
{
    int bitswritten = 0;

    do {
        /* If the buffer is full, write it. */

        if ((p->idx == 254 && !p->bitsLeft) || p->idx > 254) {
            if (putc(255, p->fp) == EOF)
              return -2;

            if (fwrite(p->buf, 1, 255, p->fp) != 255)
              return -2;

            memset(p->buf, 0x00, 256);
            p->idx = 0;
            p->bitsLeft = 8;
        }

        /*
         *  Now take care of the two cases
         */
        if (nbits <= p->bitsLeft) {
            p->buf[p->idx] |= (bits & ((1 << nbits) - 1))
                          << (8 - p->bitsLeft);
            bitswritten += nbits;
            p->bitsLeft -= nbits;
            nbits = 0;
        } else {
            p->buf[p->idx] |= (bits & ((1 << p->bitsLeft) - 1))
                          << (8 - p->bitsLeft);
            bitswritten += p->bitsLeft;
            bits >>= p->bitsLeft;
            nbits -= p->bitsLeft;
            ++(p->idx);
            p->buf[p->idx] = 0;
            p->bitsLeft = 8;
        }
    } while (nbits != 0);

    return bitswritten;
}


static void FreeStrtab(struct strtab *p)
{
    if (p->Hsh) free(p->Hsh);
    if (p->Nxt) free(p->Nxt);
    if (p->Chr) free(p->Chr);
}

static int AllocStrtab(struct strtab *p)
{
    p->Chr = malloc(MAXSTR * sizeof *p->Chr);
    p->Nxt = malloc(MAXSTR * sizeof *p->Nxt);
    p->Hsh = malloc(HASHSIZE * sizeof *p->Hsh);

    if (p->Chr && p->Nxt && p->Hsh)
      return 0;

    FreeStrtab(p);
    return -3;
}


static int AddCharString(struct strtab *p, int index, Byte b)
{
    int hshidx;

    if (p->NumStrings >= MAXSTR)
      return 0xFFFF;

    hshidx = HASH(index, b);
    while (p->Hsh[hshidx] != 0xFFFF)
      hshidx = (hshidx + HASHSTEP) % HASHSIZE;

    p->Hsh[hshidx] = p->NumStrings;
    p->Chr[p->NumStrings] = b;
    p->Nxt[p->NumStrings] = (index != 0xFFFF) ? index : 0xFFFF;

    return (p->NumStrings)++;
}



/*-----------------------------------------------------------------
 *  DESCRIPTION: Find index of string consisting of the string of
 *               index plus the byte b.
 *  PARAMETERS:  index - Index to first part of string, or 0xFFFF
 *                       if only 1 byte is wanted
 *               b     - Last byte in string
 *  RETURNS:     Index to string, or 0xFFFF if not found
 */
static int FindCharString(struct strtab *p, int index, Byte b)
{
    int hshidx;
    int nxtidx;

    if (index == 0xFFFF)
      return b;

    /*
     *  Search the string table until the string is found, or
     *  we find HASH_FREE. In that case the string does not
     *  exist.
     */
    hshidx = HASH(index, b);
    while ((nxtidx = p->Hsh[hshidx]) != 0xFFFF) {
        if (p->Nxt[nxtidx] == index && p->Chr[nxtidx] == b)
          return nxtidx;
        hshidx = (hshidx + HASHSTEP) % HASHSIZE;
    }

    return 0xFFFF;
}


static void ClearStrtab(struct strtab *p, int codesize)
{
    int maxcodes = (1 << codesize) + RES_CODES;
    int i;

    p->NumStrings = 0;

    for (i=0; i < HASHSIZE; ++i)
      p->Hsh[i] = 0xFFFF;

    for (i=0; i < maxcodes; ++i)
      AddCharString(p, 0xFFFF, i);
}



/*=================================================================*
 =                     LZW compression routine                     =
 ==================================================================*/

static int LZW_Compress(unsigned char *data, int w, int h, FILE *fp)
{
    struct strtab st;
    struct bitfile bf;
    int c;
    int index;
    int codesize = 8;
    int clearcode = (1 << codesize);
    int endofinfo = clearcode+1;
    int numbits = codesize+1;
    int limit = (1 << numbits) - 1;
    unsigned int prefix = 0xFFFF;
    ulong picsize = w*h;
    ulong i;

    InitBitFile(&bf, fp);

    if (AllocStrtab(&st) != 0)
      return -3;
    ClearStrtab(&st, codesize);

    /* Tell the unpacker to clear the stringtable */
    WriteBits(&bf, clearcode, numbits);

    /* Pack the image */
    for (i=0; i < picsize; ++i)
    {
        c = data[i];

        if ((index = FindCharString(&st, prefix, c)) != 0xFFFF)
        {
            prefix = index;
        }
        else {
            /*
             *  The string does not exist in the table.
             *  First write code of the old prefix.
             *  Then add new string (prefix + new character)
             *  to the stringtable.
             */
            WriteBits(&bf, prefix, numbits);

            if (AddCharString(&st, prefix, c) > limit) {
                ++numbits;
                if (numbits > MAXBITS) {
                    WriteBits(&bf, clearcode, numbits-1);
                    ClearStrtab(&st, codesize);
                    numbits = codesize + 1;
                }
                limit = (1 << numbits) - 1;
            }

            prefix = c;
        }
    }

    if (prefix != 0xFFFF)
      WriteBits(&bf, prefix, numbits);

    WriteBits(&bf, endofinfo, numbits);
    ResetOutBitFile(&bf);
    FreeStrtab(&st);
    return 0;
}



/*=================================================================*
 =                           Other routines                        =
 ==================================================================*/

static int bwrite_endian(unsigned long x, int bytes,
    unsigned char *buf, int endian)
{
    unsigned char u1, u2, u3, u4;
    if (bytes==1) {
        u1 = x&0xFF;
        *buf++ = u1;
        return 1;
    }
    else if (bytes==2) {
        u1 = x&0xFF;
        u2 = (x>>8)&0xFF;
        if (endian) {
            *buf++ = u1;
            *buf++ = u2;
        }
        else {
            *buf++ = u2;
            *buf++ = u1;
        }
        return 2;
    }
    else if (bytes==4) {
        u1 = x&0xFF;
        u2 = (x>>8)&0xFF;
        u3 = (x>>16)&0xFF;
        u4 = (x>>24)&0xFF;
        if (endian) {
            *buf++ = u1;
            *buf++ = u2;
            *buf++ = u3;
            *buf++ = u4;
        }
        else {
            *buf++ = u4;
            *buf++ = u3;
            *buf++ = u2;
            *buf++ = u1;
        }
        return 4;
    }
    else {
        return -1;
    }
}


static int GIF_WriteID(FILE *fp, unsigned char separator,
                       int w, int h,
                       int sort_flag, int interlace_flag,
                       int lct_flag, int lct_sz)
{
    unsigned char flags;
    unsigned char buf[100];
    int Endian = 1; /* little-endian */

    flags = ((lct_flag & 0x01) << 7)
          | ((interlace_flag & 1) << 6)
          | ((sort_flag & 1) << 5)
          | ((0 & 0x03) << 3)
          | (lct_sz & 0x07);

    /* mark start of record */
    bwrite_endian(separator, 1, buf+0, Endian);
    /* left "margin" */
    bwrite_endian(0, 2, buf+1, Endian);
    /* top "margin" */
    bwrite_endian(0, 2, buf+3, Endian);
    bwrite_endian(w, 2, buf+5, Endian);
    bwrite_endian(h, 2, buf+7, Endian);
    bwrite_endian(flags, 1, buf+9, Endian);

    return (fwrite(buf, 1, 10, fp) == 10)? 0: -2;
}


/***
 * Functions and structures related to color table manipulation.
 * No optimization here; just a linear search in findColorEntry().
 * Maybe some sort of hash table would speed it up?
 */

struct ColorTable
{
    unsigned char entries[256][3];
    int nentries;
};


static void initColorTable(struct ColorTable *tab)
{
    tab->nentries = 0;
}

static int findColorEntry(struct ColorTable *tab, unsigned char data[3])
{
    int i;
    for (i=0; i < tab->nentries; ++i) {
        if (memcmp(tab->entries[i], data, 3) == 0)
          return i;
    }
    return -1;
}

static int addColorEntry(struct ColorTable *tab, unsigned char data[3])
{
    int tmp = findColorEntry(tab, data);
    if (tmp != -1)
      return tmp;
    if (tab->nentries >= 256)
      return -1;
    memcpy(tab->entries[tab->nentries], data, 3);
    return tab->nentries++;
}





/*******************************************************************
 *                 P U B L I C    F U N C T I O N S                *
 *******************************************************************/


int WritePalGIF(const char *fname, unsigned char *paldata,
             int w, int h,
             unsigned char (*palette)[3], int npalette)
{
    FILE *fp = fopen(fname, "wb");
    int Endian = 1; /* little-endian */
    unsigned char tmp;
    unsigned char buf[100];
    int rc;

    if (fp == NULL)
      return -1;
    if (npalette > 256) {
        fclose(fp);
        return -4;
    }

    tmp = (1 << 7) /* color table flag */
          | (7 << 4) /* color resolution - 1 */
          | (0 << 3) /* sort flag */
          | (7); /* color table size - 1 */

    memcpy(buf, "GIF87a", 6);
    bwrite_endian(w, 2, buf+6, Endian);
    bwrite_endian(h, 2, buf+8, Endian);
    bwrite_endian(tmp, 1, buf+10, Endian);
    bwrite_endian(0, 1, buf+11, Endian); /* background color index */
    bwrite_endian(0, 1, buf+12, Endian); /* pixel aspect ratio */

    if (fwrite(buf, 1, 13, fp) != 13)
      goto close_fail;


    /* Write global color table */
    {
        unsigned char buffer[3*256] = {0};
        memcpy(buffer, palette, 3*npalette);
        if ((fwrite(buffer, 3, 256, fp)) != 256)
          goto close_fail;
    }


    /* Write image descriptor with ',' */
    GIF_WriteID(fp, 0x2C, w, h, 0, 0, 0, 0);

    /* Code size */
    if (putc(8, fp) == EOF)
      goto close_fail;
    if ((rc = LZW_Compress(paldata, w, h, fp)) != 0)
      goto rc_fail;
    if (putc(0, fp) == EOF)
      goto close_fail;

    /* Write image descriptor with ';' */
    GIF_WriteID(fp, 0x3B, w, h, 0, 0, 0, 0);

    fclose(fp);
    return 0;

    close_fail:
    fclose(fp);
    return -2;

    rc_fail:
    fclose(fp);
    return rc;
}


int WriteGIF(const char *fname, unsigned char (*data)[3],
             int w, int h)
{
    int rc;
    unsigned char *paldata = NULL;
    struct ColorTable colorTable;
    ulong picsize = w*h;
    ulong i;

    paldata = malloc(picsize);
    if (paldata == NULL)
      return -3;

    initColorTable(&colorTable);

    for (i=0; i < picsize; ++i) {
        int elem = addColorEntry(&colorTable, data[i]);
        if (elem == -1) {
            free(paldata);
            return -4;
        }
        paldata[i] = elem;
    }

    rc = WritePalGIF(fname, paldata, w, h,
                     colorTable.entries, colorTable.nentries);
    if (rc != 0)
      goto rc_fail;

    free(paldata);
    return 0;

    rc_fail:
    if (paldata) free(paldata);
    return rc;
}