// rewrite png file with offsets in grAb chunk
// and optionally insert or remove alPh chunk

#include <fstream>
#include "grabpng.h"


PngGrabber::PngInfo PngGrabber::getInfo(const char *pngpath)
{  
  PngGrabber p;
  p.open(pngpath);
  PngGrabber::PngInfo info = p.pngInfo;
  p.close();
  return info;
}


void PngGrabber::staticGrab(const char *pngpath, int x, int y)
{  
  PngGrabber p;
  p.open(pngpath);
  p.grab(pngpath, x, y);
  p.close();
}

void PngGrabber::staticAlph(const char *pngpath, int alph)
{  
  PngGrabber p;
  p.open(pngpath);
  p.setAlph(pngpath, alph);
  p.close();
}

bool PngGrabber::open (const char *pngpath)
{
  if (isOpen) close();
  std::ifstream infile(pngpath,std::ios::binary|std::ios::in);
  
  infile.seekg (0, std::ios::end);
  int size = infile.tellg();
  dataSize = size;
  
  imageData = new char[size];
  
  // read first 4 bytes
  infile.seekg (0, std::ios::beg);
  infile.read((char *)imageData,4);
  
  // return false if file doesn't have png header
  if (!(imageData[1]=='P' && imageData[2]=='N' && imageData[3]=='G')) 
  {
    // fprintf(stderr,"File is not a png: %s\n", pngpath);
    infile.close();
    close();
    return false;
  }
  
  // read rest of file
  infile.read((char *)imageData+4,size-4);
  infile.close();
  
  // get size
  for(int i=4; i<size; i++)
  {
    if (imageData[i]=='I' && imageData[i+1]=='D' && imageData[i+2]=='A' && imageData[i+3]=='T') 
      break;  // stop when we reach the image data
    if (imageData[i]=='I' && imageData[i+1]=='H' && imageData[i+2]=='D' && imageData[i+3]=='R') 
    {
      imageSize.width = (int)byteSwapLong((long &)imageData[i+4]);
      imageSize.height = (int)byteSwapLong((long &)imageData[i+8]);
      pngInfo.w = imageSize.width;
      pngInfo.h = imageSize.height;
    }
    if (imageData[i]=='g' && imageData[i+1]=='r' && imageData[i+2]=='A' && imageData[i+3]=='b') 
    {
      pngInfo.x = (int)byteSwapLong((long &)imageData[i+4]);
      pngInfo.y = (int)byteSwapLong((long &)imageData[i+8]);
    }
    if (imageData[i]=='a' && imageData[i+1]=='l' && imageData[i+2]=='P' && imageData[i+3]=='h') 
    {
      pngInfo.alph = 1;
    }
  }
  pngInfo.ok = 1;
  isOpen = true;
  return true;
}
  
void PngGrabber::close()
{
  isOpen = false;
  imageSize.width = imageSize.height = 0; 
  if (imageData) delete imageData;
  imageData = 0;
}

bool PngGrabber::grab(const char *outpath, int x, int y)
{
  // look for grab and alph chunk
  int grabPos = 33;
  int oldGrabSize = 0;
  for(unsigned i=4; i<dataSize; i++)
  {
    if (imageData[i]=='I' && imageData[i+1]=='D' && imageData[i+2]=='A' && imageData[i+3]=='T') 
      break;  // stop when we reach the image data
    if (imageData[i]=='g' && imageData[i+1]=='r' && imageData[i+2]=='A' && imageData[i+3]=='b')
    {
      grabPos = i-4;
      oldGrabSize = 20;
    }
  }
  
  // prepare new grab chunk
  grabChunk chunk = { 
    {0,0,0,8},
    {'g', 'r', 'A', 'b'}, 
    byteSwapLong((long)x), 
    byteSwapLong((long)y) ,
    0
  };
  chunk.crc = byteSwapULong(crc((unsigned char*)&chunk+4, 12));
  
  // create output stream overwriting the file
  std::ofstream outfile(outpath,std::ios::binary|std::ios::out);
  
  // Write PNG header and IHDR chunk
  outfile.write((char *)imageData, grabPos);

  int filePos= grabPos;
  
  // Write grAb chunk
  outfile.write((char *)&chunk, 20);
  filePos = grabPos + oldGrabSize;
  
  // Write the rest of the file
  outfile.write((char *)imageData + filePos, dataSize - filePos);
  
  outfile.close();

  return true;
}


bool PngGrabber::setAlph(const char *outpath, int doAlph) // doAlph: 1=add, -1=remove
{
  // look for alph chunk
  int alphPos = 33;
  int oldAlphSize = 0;
  for(unsigned i=4; i<dataSize; i++)
  {
    if (imageData[i]=='I' && imageData[i+1]=='D' && imageData[i+2]=='A' && imageData[i+3]=='T') 
      break;  // stop when we reach the image data
    if (imageData[i]=='a' && imageData[i+1]=='l' && imageData[i+2]=='P' && imageData[i+3]=='h')
    {
      alphPos = i-4;
      oldAlphSize = 12;
    }
  }
  
  // create output stream overwriting the file
  std::ofstream outfile(outpath,std::ios::binary|std::ios::out);
  
  // Write PNG header and IHDR chunk
  outfile.write((char *)imageData, alphPos);

  int filePos= alphPos;

  if (doAlph == 1 || oldAlphSize > 1) // insert, overwrite, or remove alph
  {
    if (doAlph == 1)
    {
      alphChunk a =  {{0,0,0,0},{'a','l','P','h'},{0x1a,0xc1,0xc9,0}};
      outfile.write((char *)&a, 12);
    }
    filePos = (oldAlphSize ? alphPos + oldAlphSize : filePos);
  }
  
  // Write the rest of the file
  outfile.write((char *)imageData + filePos, dataSize - filePos);
  
  outfile.close();

  return true;
}

// swap bytes for big endian format

 inline unsigned long PngGrabber::byteSwapULong(unsigned long value)
{
  return (((value&0x000000FF)<<24)+((value&0x0000FF00)<<8)+
  ((value&0x00FF0000)>>8)+((value&0xFF000000)>>24));
}
 inline unsigned long PngGrabber::byteSwapLong(long value)
{
  return (long)byteSwapULong((unsigned long)value) ;
}

// wild and crazy crc stuff from zlib

int PngGrabber::crc_table_computed = 0;
unsigned long PngGrabber::crc_table[256];
void PngGrabber::make_crc_table(void)
{
  unsigned long c;
  int n, k;
  for (n = 0; n < 256; n++)
  {
    c = (unsigned long) n;
    for (k = 0; k < 8; k++)
    {
      if (c & 1)
        c = 0xedb88320L ^ (c >> 1);
      else
        c = c >> 1;
    }
    crc_table[n] = c;
  }
  crc_table_computed = 1;
}
unsigned long PngGrabber::update_crc(unsigned long crc, unsigned char *buf, int len)
{
  unsigned long c = crc;
  int n;
  if (!crc_table_computed)
    make_crc_table();
  for (n = 0; n < len; n++)
    c = crc_table[(c ^ buf[n]) & 0xff] ^ (c >> 8);
  return c;
}
unsigned long PngGrabber::crc(unsigned char *buf, int len)
{
  return update_crc(0xffffffffL, buf, len) ^ 0xffffffffL;
}
