// ultra alpha code
// not even known to work :)
//
// who knows :)


#include <prcore/extlib/zlib/zlib.h>
#include <prcore/prcore.hpp>

using namespace prcore;

//////////////////////////////////////////////////////
// zip structures                                  //
////////////////////////////////////////////////////

#pragma pack(1)

struct LocalFileHeader
{
  uint32	signature;			// 0x04034b50 ("PK..")
  uint16	versionNeeded;		// version needed to extract
  uint16	flags;				// 	
  uint16	compression;		// compression method
  uint16	lastModTime;		// last mod file time
  uint16	lastModDate;		// last mod file date
  uint32	crc;				//
  uint32	compressedSize;		// 
  uint32	uncompressedSize;	// 
  uint16	filenameLen;		// length of the filename field following this structure
  uint16	extraFieldLen;		// length of the extra field following the filename field
};

struct DataDesc
{
  uint32	signature;			// 0x08074b50
  uint32	crc;				//
  uint32	compressedSize;		//
  uint32	uncompressedSize;	//
};

struct DirFileHeader
{
  uint32	signature;			// 0x02014b50
  uint16	versionUsed;		//
  uint16	versionNeeded;		//
  uint16	flags;				//
  uint16	compression;		// compression method
  uint16	lastModTime;		//
  uint16	lastModDate;		//
  uint32	crc;				//
  uint32	compressedSize;		//
  uint32	uncompressedSize;	//
  uint16	filenameLen;		// length of the filename field following this structure
  uint16	extraFieldLen;		// length of the extra field following the filename field
  uint16	commentLen;			// length of the file comment field following the extra field
  uint16	diskStart;			// the number of the disk on which this file begins
  uint16	internal;			// internal file attributes
  uint32	external;			// external file attributes
  uint32	localOffset;		// relative offset of the local file header
};

struct DirEndRecord
{
  uint32	signature;			// 0x06054b50
  uint16	thisDisk;			// number of this disk
  uint16	dirStartDisk;		// number of the disk containing the start of the central directory
  uint16	numEntriesOnDisk;	// # of entries in the central directory on this disk
  uint16	numEntriesTotal;	// total # of entries in the central directory
  uint32	dirSize;			// size of the central directory
  uint32	dirStartOffset;		// offset of the start of central directory on the disk where the central directory begins
  uint16	commentLen;			// zip file comment length
};

#pragma pack()

static const char dirFileHeaderSig[4] = { 0x50, 0x4b, 0x01, 0x02 };
static const char localFileHeaderSig[4]	= { 0x50, 0x4b, 0x03, 0x04 };
static const char dirEndRecordSig[4] = { 0x50, 0x4b, 0x05, 0x06 };
static const char dataDescSig[4] = { 0x50, 0x4b, 0x07, 0x08 };
static const unsigned long* l_crcTable = get_crc_table();


/////////////////////////////////////////////////////
// zip functions                                  //
////////////////////////////////////////////////////

static char zip_password[256];
static int zip_password_length = 0;


static uint32 Crc32(uint32 crc, unsigned char ch)
{
	return l_crcTable[((int)crc ^ (ch)) & 0xff] ^ (crc >> 8);
}


static void UpdateKeys(uint32* keys, int ch)
{
	keys[0] = Crc32( keys[0], ch );
	keys[1] = (keys[1] + (keys[0] & 0xff)) * 134775813L + 1;
	keys[2] = Crc32( keys[2], keys[1] >> 24 );
}


static int DecryptByte(uint32* keys)
{
	uint32 temp = (keys[2] & 0xffff) | 2;
	return (((temp * (temp ^ 1)) >> 8) & 0xff);
}


static uint8* zipDecrypt(const uint8* pSrc, uint32 srcLen, int version, uint32 crc)
{
	if ( zip_password_length < 1 )
		return NULL;

	// init keys with password
	uint32 key[3] = { 305419896L, 591751049L, 878082192L };
	for( int i=0; i < zip_password_length;i++ )
	{
		UpdateKeys(key,zip_password[i]);
	}

	// decrypt the 12 byte encryption header
	uint8 fileKey[12];
	memcpy( fileKey, pSrc, sizeof(fileKey) );
	for( i=0; i < 12; i++ )
	{
		uint8 ch = fileKey[i] ^ DecryptByte( key );
		UpdateKeys( key, ch );
		fileKey[i] = ch;
	}

	// check that password is correct one
	if( version < 20 )
	{
		// version prior 2.0
		if( *((uint16*)&fileKey[10]) != (crc>>16) )
		{
			return NULL;
		}
	} 
	else 
	{
		// version 2.0+
		if( fileKey[11] != (crc>>24) ) 
		{
			return NULL;
		}
	}

	// read compressed data & decrypt
	uint8* pBuffer = new uint8[ srcLen ];
	memcpy( pBuffer, &pSrc[ sizeof(fileKey) ], srcLen );

	for( uint32 n=0; n < srcLen; n++ )
	{
		uint8 ch = pBuffer[n];
		ch ^= DecryptByte( key );
		pBuffer[n] = ch;
		UpdateKeys( key, ch );
	}

	return pBuffer;
}


uint32 zipDecompress(uint8* compressed, uint8* uncompressed, uint32 compressedLen, uint32 uncompressedLen)
{
    int err;
    z_stream d_stream; // decompression stream

	memset( &d_stream, 0, sizeof(d_stream) );

    d_stream.zalloc = (alloc_func)0;
    d_stream.zfree	= (free_func)0;
    d_stream.opaque = (voidpf)0;
	
    d_stream.next_in  = compressed;
    d_stream.avail_in = compressedLen + 1;//TODO: correct?

	err = inflateInit2( &d_stream, -MAX_WBITS );
	if( err != Z_OK )
	{
		PRCORE_EXCEPTION("Zip","ZIP file decompression (inflateInit2) failed.")
		return 0;
	}

	for (;;) 
	{
		d_stream.next_out	= uncompressed;
		d_stream.avail_out	= uncompressedLen;

		err = inflate( &d_stream, Z_FINISH );
		if (err == Z_STREAM_END) 
			break;
      
		if( err != Z_OK )
		{
			switch( err )
			{
				case Z_MEM_ERROR:	PRCORE_EXCEPTION("Zip","ZIP file decompression failed (memory error).")
				case Z_BUF_ERROR:	PRCORE_EXCEPTION("Zip","ZIP file decompression failed (buffer error).")
				case Z_DATA_ERROR:	PRCORE_EXCEPTION("Zip","ZIP file decompression failed (data error).")
				default:			PRCORE_EXCEPTION("Zip","ZIP file decompression failed (unknown error).")
			}
			return 0;
		}
	}

	err = inflateEnd( &d_stream );
	if( err != Z_OK )
	{
		PRCORE_EXCEPTION("Zip","ZIP file decompression (inflateEnd) failed.")
		return 0;
	}

	return d_stream.total_out;
}



//////////////////////////////////////////////////////
// unzip                                           //
////////////////////////////////////////////////////

class unzip
{
public:
  
  unzip(const char* zippath);
  ~unzip();
  
  int GetFileIndex(const char* filepath);
  uint8* LoadFileFromZIP(int fileindex, uint32& size);
  bool IsOpen() const {return mIsOpen;}
  
private:

  struct FileInfo
  {
    DirFileHeader header;
    char filename[256];

  };

  bool mIsOpen;
  FILE *mFile;
  int mNumFiles;
  int mFileSize;

  FileInfo* mFileInfos;
  DirEndRecord mEndRecord;

  void CloseFile();
  bool ReadEndRecord();
  void ReadDirFileHeaders();
  uint32 ReadLocalHeader(const FileInfo& info);
};


unzip::unzip(const char* zippath)
{
 
	mIsOpen	= false;

	// open the zip file

	mFile=fopen(zippath,"ro");
	if(mFile==NULL)
	  {
	    PRCORE_EXCEPTION("Zip","Unable to open zip  file (unzip::unzip(const char *zippath) failed")
	      CloseFile();
	      return;
	  }

	fseek(mFile,0,SEEK_END);
	mFileSize=ftell(mFile);
	
	if(mFileSize < 0)
	  {
	    PRCORE_EXCEPTION("Zip","Corrupt zip file")
	      CloseFile();
	    return;
	  }

	if(!ReadEndRecord())
	  {
	    CloseFile();
	    return;
	  }
	ReadDirFileHeaders();
	
	mIsOpen = true;
}


unzip::~unzip()
{
  CloseFile();
}


void unzip::CloseFile()
{
  if(mIsOpen)
    {
      fclose(mFile);
      delete mFileInfos;
      mIsOpen = false;
      delete mFileInfos;

    }
}

bool unzip::ReadEndRecord()
{
  char sig[4];

  for(uint32 pos=mFileSize-sizeof(DirEndRecord);pos > 0;pos--)
   {
     fseek(mFile,pos,SEEK_SET); // seek to position
     fread(sig,sizeof(uint8),4,mFile);
     //     printf("%c%c%c%c\n",sig[0],sig[1],sig[2],sig[3]);

     if(memcmp(sig,dirEndRecordSig,4) == 0)
       {
	 
	 fseek(mFile,-4,SEEK_CUR); // go back 4 
	 fread(&mEndRecord,sizeof(DirEndRecord),1,mFile);

      // fail, if multi-volume zip
      if( mEndRecord.thisDisk		!= 0	|| 
	  mEndRecord.dirStartDisk != 0	|| 
	  mEndRecord.numEntriesOnDisk != mEndRecord.numEntriesTotal )
       	return false;
      
      return true;
       }
   }
  
  return false;
}     

void unzip::ReadDirFileHeaders()
{
  int n;
  mNumFiles = mEndRecord.numEntriesTotal;
  printf("mNumFiles: %d\n",mNumFiles);
  mFileInfos = new FileInfo[mNumFiles];
  
  fseek(mFile,mEndRecord.dirStartOffset,SEEK_SET);
  
  printf("start offset %d\n",mEndRecord.dirStartOffset);
  for(n=0;n<mNumFiles;n++)
    {
      
      printf("%d\n",n);
      
      FileInfo &info = mFileInfos[n];
      
      fread(&info.header,sizeof(DirFileHeader),1,mFile);
       
      char *pBuffer = new char[info.header.filenameLen + 1];
      fread(pBuffer,sizeof(uint8),info.header.filenameLen,mFile);
      
      pBuffer[info.header.filenameLen]=0;
      //strupr(pBuffer);
      strcpy(info.filename,pBuffer);
      delete[] pBuffer;
     
      // seek over extra field + comment
      fseek(mFile,info.header.extraFieldLen,SEEK_CUR);
      fseek(mFile,info.header.commentLen,SEEK_CUR);
      
    }

}

int unzip::GetFileIndex(const char* filePath)
{
  int i;

  // case 'insensitive search

  for(i=0;i<mNumFiles;i++)
    if(stricmp(filePath,mFileInfos[i].filename)==0)
      return i; // found

  return -1; // not found!?

}


uint8* unzip::LoadFileFromZIP(int fileIndex, uint32& rSize)
{
  if( fileIndex < 0 || fileIndex >= mNumFiles )
    return 0;
  
  // integrity check
  FileInfo& info = mFileInfos[ fileIndex ];
  uint32 sourcePos = ReadLocalHeader( info );
	
  bool	srcAllocated	= false;

  // source 
  uint32 srcLen = info.header.compressedSize;
 
 fseek(mFile,sourcePos,SEEK_SET); // load up compressed data into memory
 uint8* pSrc = new uint8[srcLen];
 fread(pSrc,sizeof(uint8),srcLen,mFile);
 
 // destination
 uint32 destLen = info.header.uncompressedSize;
 uint8* pDest = new uint8[ destLen ]; // allocate memory for destination
  
 // encrypted? 
 if( info.header.flags & 1 )
   {
     // NOTE: the compressed data stream must be read to memory and decrypted
     pSrc = zipDecrypt( pSrc, srcLen, (info.header.versionUsed & 0xff), info.header.crc );
     srcAllocated = true;
   }
  
 // decompress
 switch( info.header.compression )
   {
   case 0:	// no compression - stored 
     {		
       memcpy( pDest, pSrc, destLen );
       break;
     }
      
   case 8: // deflated
     {
       if( zipDecompress( pSrc, pDest, srcLen, destLen ) != destLen )
	 return 0;
       break;
     }
     
   default:
     return 0;
   }
 
 // CRC check
 if( crc32( 0, pDest, destLen ) != info.header.crc )
   return 0;
 
 // if we allocated source buffer to memory (in decrypting), free it
 if( srcAllocated ) 
   delete[] pSrc;
 
  // done
 rSize = destLen;
 return pDest;

}

uint32 unzip::ReadLocalHeader(const FileInfo& info)
{
  const DirFileHeader& dirHdr = info.header;
  LocalFileHeader lhdr;
  
  // Read the local header
  fseek(mFile,dirHdr.localOffset,SEEK_SET);
  fread(&lhdr,sizeof(LocalFileHeader),1,mFile);
    
  // Consistency check
  if( lhdr.compressedSize != dirHdr.compressedSize ) 
    return 0;
   
  // Skip local header + name + extra field
  return dirHdr.localOffset + sizeof(LocalFileHeader) + lhdr.filenameLen + lhdr.extraFieldLen;
}

////////////////////////////////////////////////////
// factory
////////////////////////////////////////////////////

static int NumEXT()
{
  return 2;
}

static const char* GetEXT(int index)
{
  switch(index)
    {
    case 0: return "zip";
    case 1: return "pak";

    default: return "";
    }
}

static bool IsReader()
{
	return true;
}


static bool IsWriter()
{
	return false;
}

static Stream* CreateFile(const char* path, const char* filename, Stream::AccessMode mode, void *password)
{
  if ( mode != Stream::READ )
    return NULL;
  
  // create unzip
  unzip* u = new unzip(path);
  if ( !u->IsOpen() )
    {
      delete u;
      return NULL;
    }
  
  
  // get fileindex
  int index = u->GetFileIndex(filename);
  if( index < 0 )
    {
      delete u;
      return NULL;
    }
  
  
  // set decrypt password
  if ( password )
    {
      strcpy(zip_password,(char*)password);
      zip_password_length = strlen((char*)password);
    }

  
  // decompress whole file to memory
  uint32 size = 0;
  uint8* data = u->LoadFileFromZIP(index,size);
  if( !data || !size )
    {
      delete u;
      return NULL;
    }
	
  // release unzip
  delete u;

  // create zip interface
  MemoryStream* z = new MemoryStream(data,size);
  return z;
}

FileFactory CreateFileFactoryZIP()
{
  FileFactory factory;
  
  factory.NumEXT = NumEXT;
  factory.GetEXT = GetEXT;
  factory.IsReader = IsReader;
  factory.IsWriter = IsWriter;
  factory.CreateFile = CreateFile;

  return factory;
}
