// unzip.cpp : Defines the entry point for the application.
//

#include <windows.h>
#include <stdlib.h>
#include "unzip.h"
#include "zconf.h"
#include "zlib.h"
#include "../Profiler.h"

// postfix of ZIP central index file name.
#define CENTRAL_DIRECTORY_INDEX_FILE_POSTFIX  _T(".cdindex")
#define CENTRAL_DIRECTORY_INDEX_FILE_VERSION  0x00000001

#define FILE_NAME_LENGTH_OFFSET 28
#define FILE_NAME_OFFSET 46

/**
 * central directory index
 */
struct central_directory_index
{
  DWORD zip_file_size;                      // file size of target ZIP file.
  FILETIME  last_write_time_of_zip_file;    // timestamp of target ZIP file.
  DWORD number_of_items;
  const char** items;
};

/**
 * ZIPファイルエントリ
 */
struct zip_entry
{
  WORD  compression_method;	          // 圧縮方法（0は非圧縮）
  DWORD compressed_size;	          // 圧縮後サイズ
  DWORD uncompressed_size;	          // 非圧縮サイズ
  DWORD relative_offset_of_local_header;  // ローカルヘッダ開始アドレスのオフセット
};

/**
 * ZIPファイルハンドル
 */
struct ZIPFILE
{
  HANDLE hFile;			          // ファイルハンドル
  HANDLE hMapping;		          // マッピングハンドル
  LPVOID data;			          // マップされたアドレス 

  DWORD file_size;		          // ZIPファイルサイズ 

  DWORD central_directory_size;	          // central directoryのサイズ
  DWORD central_directory_offset;	  // central directoryの開始オフセット

  zip_entry current_entry;	          // カレントエントリ

  central_directory_index* index;         // central directory index.
};

/**
 * read 32bit value.
 */
#define READ_INT32(p) \
	(((__int32) (p)[0] & 0x000000ff) \
			| ((__int32) (p)[1] << 8 & 0x0000ff00) \
			| ((__int32) (p)[2] << 16 & 0x00ff0000) \
			| ((__int32) (p)[3] << 24 & 0xff000000))

/**
 * read 16bit value.
 */
#define READ_INT16(p) \
	((__int16) (p)[0] & 0x00ff) \
			| ((__int16) (p)[1] << 8 & 0xff00)


// static functions
static int compare_central_directory_file_name(WORD len1, const char* name1, WORD len2, const char* name2);
static int compare_central_directory_index_item(const void* item1, const void* item2);
static int load_central_directory_index(ZIPFILE* target, HANDLE hFile, central_directory_index** index);
static int free_central_directory_index(central_directory_index* index);
static const char* find_from_central_directory_index(central_directory_index* index, const char *filename);
static int create_central_directory_index_file(ZIPFILE* target, HANDLE hFileDest);


/**
 * compares central directory file name.
 *
 */
static int compare_central_directory_file_name(WORD len1, const char* name1, WORD len2, const char* name2)
{
  if (len1 != len2)
  {
    return ((int)len1 & 0xffff) - ((int) len2 & 0xffff);
  }

  return memcmp(name1, name2, len1);
}

/**
 * compares central directory index item
 */
static int compare_central_directory_index_item(const void* item1, const void* item2)
{
  const char* p1 = *(const char**) item1;
  const char* p2 = *(const char**) item2;
  
  WORD len1 = READ_INT16 (p1 + FILE_NAME_LENGTH_OFFSET);
  WORD len2 = READ_INT16 (p2 + FILE_NAME_LENGTH_OFFSET);
  const char* name1 = (const char*) (p1 + FILE_NAME_OFFSET);
  const char* name2 = (const char*) (p2 + FILE_NAME_OFFSET);
  return compare_central_directory_file_name(len1, name1, len2, name2);
}

/**
 * create central directory index file
 */
static int create_central_directory_index_file(ZIPFILE* target, HANDLE hFileDest)
{
  HANDLE hFileTarget = target->hFile;
  DWORD size = GetFileSize(hFileTarget, NULL);
  DWORD numOfWritten = 0;

  // write index file version
  {
    DWORD version = CENTRAL_DIRECTORY_INDEX_FILE_VERSION;
    if (! WriteFile(hFileDest, &version, sizeof(version), &numOfWritten, NULL))
    {
      return 0;
    }
  }

  // write ZIP file size.
  if (! WriteFile(hFileDest, &size, sizeof(size), &numOfWritten, NULL))
  {
    return 0;
  }

  // write last write time of ZIP file.
  FILETIME lastWriteTime;
  GetFileTime(hFileTarget, NULL, NULL, &lastWriteTime);
  if (! WriteFile(hFileDest, &lastWriteTime, sizeof(lastWriteTime), &numOfWritten, NULL))
  {
    return 0;
  }

  // get central directory's address.
  const char *p =
    (const char *) ((char *) target->data + target->central_directory_offset);

  const char* const start = p;

  // get end address of central directory
  const char *end = p + target->central_directory_size;

  // count entries.
  int numOfItems = 0;
  while (p < end)
    {
      if (!(p[0] == 0x50 && p[1] == 0x4b && p[2] == 0x01 && p[3] == 0x02))
	{
	  // invalid signature.
	  return 0;
	}

      numOfItems++;

      // ファイル名、エクストラフィールド、ファイルコメントの長さを得る
      WORD filename_length = READ_INT16 (p + FILE_NAME_LENGTH_OFFSET);
      WORD extra_field_length = READ_INT16 (p + 30);
      WORD file_comment_length = READ_INT16 (p + 32);
      // 次の位置に移動する
      p += 46			// central directoryの固定部分サイズ
	    + filename_length
            + extra_field_length
            + file_comment_length;
    }

  const char** items = (const char**) malloc(sizeof(char*) * numOfItems);
  if (! items)
  {
    return 0;
  }

  p = start;
  int current_index = 0;
  while (p < end)
  {
      // ファイル名、エクストラフィールド、ファイルコメントの長さを得る
      WORD filename_length = READ_INT16 (p + 28);
      WORD extra_field_length = READ_INT16 (p + 30);
      WORD file_comment_length = READ_INT16 (p + 32);
      // 次の位置に移動する
      p += 46			// central directoryの固定部分サイズ
	    + filename_length
            + extra_field_length
            + file_comment_length;

      items[current_index++] = p;
  }

  // sort items
  qsort((void*) items, (size_t) numOfItems, sizeof(const char*), compare_central_directory_index_item);
  
  // writes number of items.
  if (! WriteFile(hFileDest, &numOfItems, sizeof(numOfItems), &numOfWritten, NULL))
  {
    free(items);
    return 0;
  }

  // adjust offset.
  size_t buffer_size = sizeof(unsigned int) * numOfItems;
  unsigned int* offset = (unsigned int*) malloc(buffer_size);
  for (int i = 0; i < numOfItems; ++i)
  {
    offset[i] = items[i] - (const char*) target->data;
  }

  // write offset
  DWORD totalWritten = 0;
  char* work = (char*) offset;
  while (totalWritten < buffer_size)
  {
    if (! WriteFile(hFileDest, work, buffer_size - totalWritten, &numOfWritten, NULL))
      {
        free(items);
        free(offset);
        return 0;
      }
    totalWritten += numOfWritten;
    work += numOfWritten;
  }
  free(offset);
  free(items);
  return numOfItems;
}

/**
 * loads central directory index file.
 */
static int load_central_directory_index(ZIPFILE* target, HANDLE hFile, central_directory_index** index)
{
  *index = NULL;
  
  // check index file version.
  DWORD numOfRead;
  {
    DWORD version;
    if (! ReadFile(hFile, &version, sizeof(version), &numOfRead, NULL))
    {
      return 0;
    }
    if (version != CENTRAL_DIRECTORY_INDEX_FILE_VERSION)
    {
      return 0;
    }
  }

  // read target ZIP file size.
  DWORD size;
  if (! ReadFile(hFile, &size, sizeof(size), &numOfRead, NULL))
  {
    return 0;
  }
  FILETIME last_write_time;
  if (! ReadFile(hFile, &last_write_time, sizeof(last_write_time), &numOfRead, NULL))
  {
    return 0;
  }

  DWORD number_of_items;
  if (! ReadFile(hFile, &number_of_items, sizeof(number_of_items), &numOfRead, NULL))
  {
    return 0;
  }

  // allocate memory.
  *index = (central_directory_index*) malloc(sizeof(central_directory_index));
  if (! *index)
  {
    return 0;
  }
  central_directory_index* work = *index;
  work->items = (const char**) malloc(sizeof(char*) * number_of_items);
  if (! work->items)
  {
    free(work);
    return 0;
  }

  // setup struct.
  work->zip_file_size = size;
  work->last_write_time_of_zip_file = last_write_time;
  work->number_of_items = number_of_items;

  // load offset from file.
  DWORD total_read = 0;
  const DWORD buffer_size = sizeof(int) * number_of_items;
  unsigned int* buff = (unsigned int*) malloc(buffer_size);
  unsigned int* start_of_buff = buff;
  while (total_read < buffer_size)
  {
    if (! ReadFile(hFile, buff, buffer_size - total_read, &numOfRead, NULL))
    {
      free_central_directory_index(work);
      free(start_of_buff);
      return 0;
    }
    total_read += numOfRead;
    buff += numOfRead;
  }

  // adjust offset
  for (DWORD i = 0; i < number_of_items; ++i)
  {
    work->items[i] = (char*)target->data + start_of_buff[i];
  }
  free(start_of_buff);

  return number_of_items;
}

static int free_central_directory_index(central_directory_index* index)
{
  if (index)
  {
    free(index->items);
    free(index);
    return 1;
  }
  return 0;
}

static const char* find_from_central_directory_index(central_directory_index* index, const char *filename)
{
  // binary search.
  if (! index)
  {
    return NULL;
  }

  DWORD left = 0, right = index->number_of_items - 1;
  DWORD middle;
  size_t len = strlen(filename);

  while (left <= right)
  {
    middle = (left + right) / 2;
    const char* p = index->items[middle];
    WORD len2 = READ_INT16 (p + FILE_NAME_LENGTH_OFFSET);
    const char* name2 = (const char*) (p + FILE_NAME_OFFSET);
    int c = compare_central_directory_file_name(len, filename, len2, name2);
    if (c == 0)
    {
      return p;
    }
    else if (c > 0)
    {
      left = middle + 1;
    }
    else
    {
      right = middle - 1;
    }
  }
  return NULL;
}

/**
 * open spacified ZIP file.
 *
 * @param filename  ZIP file path.
 */
ZIPFILE *
unzip_open (LPTSTR filename)
{
  ZIPFILE *zfile = NULL;

  // open for file mapping.
  HANDLE hFile = CreateFileForMapping (filename,
				       GENERIC_READ,
				       FILE_SHARE_READ,
				       NULL,
				       OPEN_EXISTING,
				       FILE_FLAG_RANDOM_ACCESS,
				       NULL);
  if (hFile != NULL)
    {
      // get file size.
      DWORD size = GetFileSize (hFile, NULL);
      // ファイルマッピングハンドルを作成する
      HANDLE hMapping = CreateFileMapping (hFile,
					   NULL,
					   PAGE_READONLY,
					   0,	        // unlimited file size
					   0,	        // entire file.
					   NULL);	// this file mapping object has no name.
      if (hMapping != NULL)
	{
	  // ファイルをメモリにマップする
	  LPVOID buff = MapViewOfFile (hMapping,
				       FILE_MAP_READ,
				       0,	// map offset.
				       0,	// 
				       0);	// map entire file.
	  if (buff != NULL)
	    {
	      // ファイルの最後尾 - 20バイト目に移動し、"End of central directory record"
	      // の内容を読み込む
	      for (char *p = (char *)buff + (size - 1 - 20); p > buff; --p)
		{
		  // 最後尾からシグニチャに一致する数値を探す
		  if (p[0] == 0x50 && p[1] == 0x4b && p[2] == 0x05
		      && p[3] == 0x06)
		    {
		      // initialize struct ZIPFILE
		      zfile = (ZIPFILE *) malloc (sizeof (ZIPFILE));
		      zfile->hFile = hFile;
		      zfile->hMapping = hMapping;
		      zfile->data = buff;
		      zfile->file_size = size;
		      zfile->central_directory_size = READ_INT32 (p + 12);
		      zfile->central_directory_offset = READ_INT32 (p + 16);

		      zfile->current_entry.relative_offset_of_local_header =
			(DWORD) - 1;

		      break;
		    }
		}
	      if (zfile == NULL)
		{
		  // invalid ZIP file.
		  UnmapViewOfFile (buff);
		}
	    }
	  if (zfile == NULL)
	    {
	      // invalid zip file.
	      CloseHandle (hMapping);
	    }
	}

      if (zfile == NULL)
	{
	  // invalid zip file.
	  CloseHandle (hFile);
	}
    }

  if (zfile)
  {
    // load central directory index.
    zfile->index = NULL;
    _TCHAR index_file_name[MAX_PATH];
    _stprintf(index_file_name, _T("%s%s"), filename, CENTRAL_DIRECTORY_INDEX_FILE_POSTFIX);
    HANDLE hFileIndex = CreateFile(index_file_name,
                                   GENERIC_READ,
                                   FILE_SHARE_READ,
                                   NULL,
                                   OPEN_EXISTING,
                                   0,
                                   NULL);
    if (hFileIndex != INVALID_HANDLE_VALUE)
    {
      int load_success = load_central_directory_index(zfile, hFileIndex, &zfile->index);
      CloseHandle(hFileIndex);
      if (! load_success)
      {
        zfile->index = NULL;
      }
      else
      {
        // get ZIP file size and time.
        DWORD size = GetFileSize (zfile->hFile, NULL);
        FILETIME last_write_time;
        GetFileTime(zfile->hFile, NULL, NULL, &last_write_time);

        // validates index file.
        if (zfile->index->zip_file_size != size
          || last_write_time.dwHighDateTime != zfile->index->last_write_time_of_zip_file.dwHighDateTime
          || last_write_time.dwLowDateTime != zfile->index->last_write_time_of_zip_file.dwLowDateTime)
        {
          free_central_directory_index(zfile->index);
          zfile->index = NULL;
        }
      }
    }

    if (! zfile->index)
    {
      // if index file load failed, create a new central directory index file.
      hFileIndex = CreateFile(index_file_name,
			      GENERIC_READ | GENERIC_WRITE,
			      FILE_SHARE_READ,
			      NULL,
			      OPEN_ALWAYS,
			      0,
			      NULL);
      if (hFileIndex != INVALID_HANDLE_VALUE)
      {
        int result = create_central_directory_index_file(zfile, hFileIndex);
        CloseHandle(hFileIndex);
        // reload index file.
        hFileIndex = CreateFile(index_file_name,
                                GENERIC_READ,
                                FILE_SHARE_READ,
                                NULL,
                                OPEN_EXISTING,
                                0,
                                NULL);
        if (hFileIndex != INVALID_HANDLE_VALUE)
        {
          if (! load_central_directory_index(zfile, hFileIndex, &zfile->index))
          {
            zfile->index = NULL;
          }
          CloseHandle(hFileIndex);
        }
      }
    }
  }

  return zfile;
}

/**
 * 指定されたZIPファイルをクローズする
 */
BOOL
unzip_close (ZIPFILE * zfile)
{
  if (zfile == NULL)
    {
      return FALSE;
    }

  // マップを解除し、ハンドルをクローズする
  BOOL result1 = UnmapViewOfFile (zfile->data);
  BOOL result2 = CloseHandle (zfile->hMapping);
  BOOL result3 = CloseHandle (zfile->hFile);

  // メモリを解放する
  free_central_directory_index(zfile->index);
  free (zfile);

  return result1 && result2 && result3;
}

/**
 * 指定されたファイル名の位置まで移動する
 */
BOOL
unzip_seek (ZIPFILE * zfile, const char *filename)
{
#ifdef ENABLE_PROFILER
  Profiler* profiler = Profiler::UNZIP_SEEK_PROFILER;
  LARGE_INTEGER startTime;
  profiler->queryTime(&startTime);
#endif

  BOOL result = FALSE;

  const char* p;
  if (zfile->index)
  {
    p = find_from_central_directory_index(zfile->index, filename);
    if (p)
    {
      result = TRUE;
    }
  }
  else
  {
    int len = strlen (filename);

    // central directoryのアドレスを計算する
    p = (const char *) ((char *) zfile->data + zfile->central_directory_offset);

    // 検索終了アドレスを計算する
    const char *end = p + zfile->central_directory_size;

    while (p < end)
      {
        if (!(p[0] == 0x50 && p[1] == 0x4b && p[2] == 0x01 && p[3] == 0x02))
	  {
	    // シグニチャがcentral directoryのものではない場合、処理を中止する
	    break;
	  }

        // ファイル名、エクストラフィールド、ファイルコメントの長さを得る
        WORD filename_length = READ_INT16 (p + 28);
        WORD extra_field_length = READ_INT16 (p + 30);
        WORD file_comment_length = READ_INT16 (p + 32);

        if (filename_length == len)
	  {
	    // ファイル名の比較を行う
	    const char *name = (const char *) (p + FILE_NAME_OFFSET);
	    if (memcmp (name, filename, len) == 0)
	      {
	        // ファイル名が一致した場合、エントリを読み込む
	        result = TRUE;
	        break;
	      }
	  }

        // 次の位置に移動する
        p += 46			// central directoryの固定部分サイズ
	  + filename_length + extra_field_length + file_comment_length;
      }
  }
  if (result)
  {
    zip_entry *entry = &zfile->current_entry;
    entry->compression_method = READ_INT16 (p + 10);
    entry->compressed_size = READ_INT32 (p + 20);
    entry->uncompressed_size = READ_INT32 (p + 24);
    entry->relative_offset_of_local_header = READ_INT32 (p + 42);
  }
  else
  {
      // ファイルが見つからなかった
      zfile->current_entry.relative_offset_of_local_header = (DWORD) - 1;
  }
#ifdef ENABLE_PROFILER
  profiler->recordEnd(0, &startTime);
#endif

  return result;
}

/**
 * 現在のエントリに対応するファイルの展開後のサイズを返す。
 */
DWORD
unzip_get_uncompressed_size (ZIPFILE * zfile)
{
  DWORD size = (DWORD) - 1;
  if (zfile->current_entry.relative_offset_of_local_header != (DWORD) - 1)
    {
      size = zfile->current_entry.uncompressed_size;
    }
  return size;
}

/**
 * エントリを解凍し、指定されたバッファに入れる
 * （現在の実装では、バッファサイズ == 展開後のサイズになっていなければならない）
 */
BOOL
unzip_read (ZIPFILE * zfile, char *buff, DWORD bufflen)
{
  zip_entry *entry = &zfile->current_entry;
  if (entry->relative_offset_of_local_header == (DWORD) - 1)
    {
      return FALSE;
    }

  // local header開始アドレスを得る
  const char *header = (const char *) zfile->data
    + entry->relative_offset_of_local_header;
  if (!
      (header[0] == 0x50 && header[1] == 0x4b && header[2] == 0x03
       && header[3] == 0x04))
    {
      // シグニチャがlocal headerのものと一致しない場合
      return FALSE;
    }

  // データを読み込む
  BOOL result = FALSE;

  // local headerのサイズを計算するために
  // ファイル名とエクストラフィールドの長さを取得する
  WORD filename_length = READ_INT16 (header + 26);
  WORD extra_field_length = READ_INT16 (header + 28);
  unsigned char *body = (unsigned char *) header + 30	// local headerの固定部分サイズ
    + filename_length + extra_field_length;
  if (entry->compression_method == 0)
    {
      // 非圧縮の場合、コピーする
      memcpy (buff, body, bufflen);
      result = TRUE;
    }
  else
    {
      // 圧縮
      z_stream stream;
      stream.zalloc = Z_NULL;
      stream.zfree = Z_NULL;
      stream.opaque = Z_NULL;
      if (inflateInit2 (&stream, -MAX_WBITS) == Z_OK)
	{
	  // local header直後に圧縮データがある
	  stream.next_in = body;
	  stream.avail_in =
	    zfile->file_size - entry->relative_offset_of_local_header;
	  stream.next_out = (unsigned char *) buff;
	  stream.avail_out = bufflen;

	  int err = inflate (&stream, Z_FINISH);
	  if (err == Z_STREAM_END)
	    {
	      result = TRUE;
	    }
	  inflateEnd (&stream);
	}
    }

  return result;
}
