/* This file is part of REWise.
 *
 * REWise is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * REWise is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <getopt.h>
#include <time.h>
#include <libgen.h> // dirname
#include <errno.h>
#include <sys/stat.h> // mkdir
#include <utime.h>
#include <sys/statvfs.h>

// PATH_MAX, NAME_MAX
#ifdef __linux__
#include <linux/limits.h>
#else
#include <limits.h>
#endif

#include "print.h"
#include "reader.h"
#include "inflate.h"
#include "pefile.h"
#include "errors.h"
#include "wiseoverlay.h"
#include "wisescript.h"
#include "version.h"


#ifndef REWISE_DEFAULT_TMP_PATH
#define REWISE_DEFAULT_TMP_PATH "/tmp/"
#endif


#define SIZE_KiB 1024
#define SIZE_MiB 1048576    // 1024^2
#define SIZE_GiB 1073741824 // 1024^3


enum Operation {
  OP_NONE         = 0,
  OP_EXTRACT      = 1,
  OP_EXTRACT_RAW  = 2,
  OP_LIST         = 3,
  OP_VERIFY       = 4,
  OP_HELP         = 5,
  OP_SCRIPT_DEBUG = 6
};


void printPrettySize(size_t size) {
  if (size > SIZE_GiB) {
    printf("%.2f GiB", (float)size / SIZE_GiB);
  }
  else
  if (size > SIZE_MiB) {
    printf("%.2f MiB", (float)size / SIZE_MiB);
  }
  else
  if (size > SIZE_KiB) {
    printf("%.2f KiB", (float)size / SIZE_KiB);
  }
  else {
    printf("%zu bytes", size);
  }
}


unsigned long getFreeDiskSpace(char * path) {
  struct statvfs fsStats;
  if (statvfs((const char *)path, &fsStats) != 0) {
    printError("Failed to determine free disk space for '%s'. Errno: %s\n",
               strerror(errno));
    return 0;
  }
  return fsStats.f_bsize * fsStats.f_bavail;
}


void convertMsDosTime(struct tm * destTime, uint16_t date, uint16_t time) {
  destTime->tm_year = (int)((date >> 9) + 80);
  destTime->tm_mon  = (int)((date >> 5) & 0b0000000000001111);
  destTime->tm_mday = (int)(date & 0b0000000000011111);
  destTime->tm_hour = (int)(time >> 11);
  destTime->tm_min  = (int)((time >> 5) & 0b0000000000111111);
  destTime->tm_sec  = (int)(time & 0b0000000000011111) * 2;
}

static InflateObject * InflateObjPtr;
static long ScriptDeflateOffset;

#define MAX_OUTPUT_PATH (PATH_MAX - WIN_PATH_MAX) - 2
static char OutputPath[MAX_OUTPUT_PATH]; // should be absolute and end with a '/'
static char TempPath[MAX_OUTPUT_PATH] = REWISE_DEFAULT_TMP_PATH;
static char PreserveTmp = 0;
static char NoExtract   = 0;


void printHelp(void) {
  printf("==============================================================\n");
  printf("              Welcome to REWise version %s\n", REWISE_VERSION_STR);
  printf("==============================================================\n\n");
  printf(" Usage: rewise [OPERATION] [OPTIONS] INPUT_FILE\n\n");
  printf("  OPERATIONS\n");
  printf("   -x --extract      OUTPUT_PATH  Extract files.\n");
  printf("   -r --raw          OUTPUT_PATH  Extract all files in the overlay "
                                           "data. This does not move/rename "
                                           "files!\n");
  printf("   -l --list                      List files.\n");
  printf("   -V --verify                    Run extract without actually "
                                           "outputting files, crc32s will be "
                                           "checked.\n");
  printf("   -z --script-debug              Print parsed WiseScript.bin\n");
  printf("   -v --version                   Print version and exit.\n");
  printf("   -h --help                      Display this HELP.\n");
  printf("\n");
  printf("  OPTIONS\n");
  printf("   -p --preserve                  Don't delete TMP files.\n");
  printf("   -t --tmp-path     TMP_PATH     Set temporary path, default: %s\n",
         REWISE_DEFAULT_TMP_PATH);
  printf("   -d --debug                     Print debug info.\n");
  printf("   -s --silent                    Be silent, don't print anything.\n");
  printf("   -n --no-extract                Don't extract anything. This will "
                                           "be ignored with -x or -r. It also "
                                           "will not try to remove TMP files, "
                                           "so -p won't do anything.\n");
  printf("\n");
  printf("  NOTES\n");
  printf("    - Path to directory OUTPUT_PATH and TMP_PATH should exist and "
         "be writable.\n");
}


void printFile(WiseScriptFileHeader * data) {
  struct tm fileDatetime;
  convertMsDosTime(&fileDatetime, data->date, data->time);
  printf("% 12u %02d-%02d-%04d %02d:%02d:%02d '%s'\n", data->inflatedSize,
         fileDatetime.tm_mday, fileDatetime.tm_mon, fileDatetime.tm_year + 1900,
         fileDatetime.tm_hour, fileDatetime.tm_min, fileDatetime.tm_sec,
         data->destFile);
}


/* preparePath() - Joins the two given paths to dest and tries to create the
 *                 directories that don't exist yet.
 * param subPath: Rest of the filepath (including file) from WiseScript.bin
 *                Should not be larger then (WIN_PATH_MAX + 1)
 * param dest   : A pre-allocated char buffer with size PATH_MAX */
bool preparePath(char * basePath, char * subPath, char * dest) {
  // Join paths
  if ((strlen(basePath) + strlen(subPath) + 1) > PATH_MAX) {
    printError("Overflow of final path > PATH_MAX\n");
    return false;
  }
  strcpy(dest, basePath);
  strcat(dest, subPath);

  // Try to create directories as needed
  char * outputFilePath;
  char * currentSubPath;
  char * separator;

  // make a copy which strchr may manipulate.
  outputFilePath = strdup(dest);

  if (outputFilePath == NULL) {
    printError("Errno: %s\n", strerror(errno));
    return false;
  }

  // get the path without filename
  currentSubPath = dirname(outputFilePath);

  // get the path its root (string until first '/')
  separator = strchr(currentSubPath, '/');

  // This should not happen because the given path by the user should exist.
  if (separator == NULL) {
    printError("This should not happen, please report if it does! (1)\n");
    return false;
  }

  // iterate through all sub-directories from root
  while (separator != NULL) {
    // terminate the dirName string on next occurance of '/'
    separator[0] = 0x00;

    // do not create root
    if (currentSubPath[0] != 0x00) {
      // stat currentSubPath
      if (access(currentSubPath, F_OK) != 0) {
        // currentSubPath exists but is not a directory
        if (errno == ENOTDIR) {
          printError("Extract subpath '%s' exists but is not a directory!\n",
                     currentSubPath);
          free(outputFilePath);
          return false;
        }

        // currentSubPath does not exist, try to create a new directory
        if (errno == ENOENT) {
          errno = 0;
          if (mkdir(currentSubPath, 0777) != 0) {
            printError("Failed to create subpath (1): '%s'\n", currentSubPath);
            printError("Errno: %s\n", strerror(errno));
            free(outputFilePath);
            return false;
          }
        }
      }
    }

    // reset the previous set terminator
    separator[0] = '/';

    // set separator to next occurrence of '/' (will be set to NULL when
    // there are no more occurrences of '/'.
    separator = strchr(separator + 1, '/');
  }

  // last subdir
  if (access(currentSubPath, F_OK) != 0) {
    // currentSubPath exists but is not a directory
    if (errno == ENOTDIR) {
      printError("Extract path '%s' exists but is not a directory!\n",
                 currentSubPath);
      free(outputFilePath);
      return false;
    }

    // currentSubPath does not exist, try to create a new directory
    if (errno == ENOENT) {
      if (mkdir(currentSubPath, 0777) != 0) {
        printError("Failed to create subpath (2): '%s'\n", currentSubPath);
        printError("Errno: %s\n", strerror(errno));
        free(outputFilePath);
        return false;
      }
    }
  }

  // cleanup
  free(outputFilePath);

  return true;
}


void extractFile(WiseScriptFileHeader * data) {
  bool result;
  char outputFilePath[PATH_MAX];

  // Create the final absolute filepath and make sure the path exists (will be
  // created when it doesn't exist).
  if (preparePath(OutputPath, data->destFile, outputFilePath) == false) {
    printError("preparePath failed.\n");
    stopWiseScriptParse();
    return;
  }

  // Seek to deflated file start
  if (fseek(InflateObjPtr->inputFile, ((long)data->deflateStart) + ScriptDeflateOffset, SEEK_SET) != 0) {
    printError("Failed seek to file offset 0x%08X\n", data->deflateStart);
    printError("Errno: %s\n", strerror(errno));
    stopWiseScriptParse();
    return;
  }

  // Inflate/extract the file
  result = inflateExtractNextFile(InflateObjPtr, outputFilePath);
  if (result == false) {
    printError("Failed to extract '%s'\n", outputFilePath);
    stopWiseScriptParse();
    return;
  }

  // Set file access/modification datetime
  struct tm fileCreation;
  time_t creationSeconds;
  convertMsDosTime(&fileCreation, data->date, data->time);
  creationSeconds = mktime(&fileCreation);
  const struct utimbuf times = {
    .actime  = creationSeconds,
    .modtime = creationSeconds
  };
  if (utime(outputFilePath, &times) != 0) {
    printWarning("Failed to set access and modification datetime for file "
                 "'%s'\n", outputFilePath);
  }

  printInfo("Extracted %s\n", data->destFile);
}


void noExtractFile(WiseScriptFileHeader * data) {
  // Inflate/extract the file
  bool result = inflateExtractNextFile(InflateObjPtr, NULL);
  if (result == false) {
    printError("Failed to no-extract '%s'\n", data->destFile);
    stopWiseScriptParse();
    return;
  }
  printInfo("CRC32 success for '%s'\n", data->destFile);
}


bool setPath(const char * optarg, char * dest) {
  // Resolve absolute path
  char * outputPath = realpath(optarg, dest);
  if (outputPath == NULL) {
    printError("Invalid PATH given, could not resolve absolute path for "
               "'%s'. Errno: %s\n", optarg, strerror(errno));
    return false;
  }

  size_t outputPathLen = strlen(outputPath);
  // -2 for the potential '/' we may add
  if (outputPathLen >= (MAX_OUTPUT_PATH - 1)) {
    printError("Absolute path of PATH is to large.\n");
    return false;
  }

  // Make sure the path ends with a '/'
  if (dest[outputPathLen - 1] != '/') {
    strcat(dest, "/");
  }

  // Make sure the path exists
  if (access(dest, F_OK) != 0) {
    // dest exists but is not a directory
    if (errno == ENOTDIR) {
      printError("'%s' is not a directory.\n", dest);
      return false;
    }
    // NOTE: realpath would have failed when the directory does not exist.
    // dest does not exist
    /*if (errno == ENOENT) {
      printError("'%s' does not exist.\n", dest);
      return false;
    }*/
  }

  return true;
}


int main(int argc, char *argv[]) {
  char inputFile[PATH_MAX];
  long overlayOffset;
  FILE * fp;
  REWError status;
  enum Operation operation = OP_NONE;
  inputFile[0] = 0x00;

  // https://www.gnu.org/software/libc/manual/html_node/Getopt-Long-Options.html
  // https://www.gnu.org/software/libc/manual/html_node/Getopt-Long-Option-Example.html
  struct option long_options[] = {
    // OPERATIONS
    {"extract"     , required_argument, NULL, 'x'},
    {"raw"         , required_argument, NULL, 'r'},
    {"list"        , no_argument      , NULL, 'l'},
    {"verify"      , no_argument      , NULL, 'V'},
    {"script-debug", no_argument      , NULL, 'z'},
    {"version"     , no_argument      , NULL, 'v'},
    {"help"        , no_argument      , NULL, 'h'},
    // OPTIONS
    {"temp"        , required_argument, NULL, 't'},
    {"debug"       , no_argument      , NULL, 'd'},
    {"preserve"    , no_argument      , NULL, 'p'},
    {"silent"      , no_argument      , NULL, 's'},
    {"no-extract"  , no_argument      , NULL, 'n'},
    {NULL          , 0                , NULL, 0}
  };

  int option_index = 0;
  for (;;) {
    int opt = getopt_long(argc, argv, "x:r:t:lhdspznVv",
                          long_options, &option_index);

    if (opt == -1) {
      break;
    }

    switch (opt) {
      // OPERATIONS
      case 'x':
      {
        if (operation != OP_NONE) {
          printError("More then one operation is set! Do set only one.\n");
          return 1;
        }
        operation = OP_EXTRACT;
        if (setPath(optarg, OutputPath) == false) {
          return 1;
        }
      }
        break;

      case 'r':
        if (operation != OP_NONE) {
          printError("More then one operation is set! Do set only one.\n");
          return 1;
        }
        operation = OP_EXTRACT_RAW;
        if (setPath(optarg, OutputPath) == false) {
          return 1;
        }
        break;

      case 'l':
        if (operation != OP_NONE) {
          printError("More then one operation is set! Do set only one.\n");
          return 1;
        }
        operation = OP_LIST;
        break;

      case 'V':
        if (operation != OP_NONE) {
          printError("More then one operation is set! Do set only one.\n");
          return 1;
        }
        operation = OP_VERIFY;
        break;

      case 'z':
        if (operation != OP_NONE) {
          printError("More then one operation is set! Do set only one.\n");
          return 1;
        }
        operation = OP_SCRIPT_DEBUG;
        break;

      case 'v':
        printf("REWise v%s\n", REWISE_VERSION_STR);
        return 0;

      case 'h':
        printHelp();
        return 0;

      // OPTIONS
      case 'd':
        setPrintFlag(PRINT_DEBUG);
        break;

      case 's':
        setPrintFlags(PRINT_SILENT);
        break;

      case 't':
        if (setPath(optarg, TempPath) == false) {
          printError("Invalid TMP_PATH given.\n");
          return 1;
        }
        break;

      case 'p':
        PreserveTmp = 1;
        break;

      case 'n':
        NoExtract = 1;
        break;

      case '?':
        // invalid option
        printError("Invalid operation or option\n");
        return 1;

      default:
        printError("default\n");
        break;
    }
  }

  if ((argc - 1 ) < optind) {
    printError("Please supply a input file\n");
    return 1;
  }
  if ((argc - 1 ) > optind) {
    printError("Please supply only one input file\n");
    return 1;
  }

  if (strlen(argv[optind]) > (PATH_MAX - 1)) {
    printError("What are you trying to do? INPUT_FILE is larger then PATH_MAX\n");
    return 1;
  }
  strcpy(inputFile, argv[optind]);

  if (operation == OP_NONE) {
    printError("Please specify a operation.\n");
    return 1;
  }

  /* Check if input file exists */
  if (access(inputFile, F_OK) != 0) {
    printError("InputFile '%s' not found. Errno: %s\n", inputFile,
               strerror(errno));
    return 1;
  }

  // Get offset to overlay data
  overlayOffset = pefileGetOverlayOffset(inputFile);
  if (overlayOffset == -1) {
    printError("Failed to find overlay offset.\n", inputFile);
    return 1;
  }

  printDebug("InputFile: %s\n", inputFile);
  printDebug("OverlayOffset: %ld\n", overlayOffset);

  /* Open inputFile */
  fp = fopen(inputFile, "rb");

  if (fp == NULL) {
    printError("Failed to open inputFile '%s'\n", inputFile);
    printError("Errno: %s\n", strerror(errno));
    return 1;
  };

  // Seek to overlayData
  if (fseek(fp, overlayOffset, SEEK_SET) != 0) {
    printError("Failed to seek to overlayData. Offset: 0x%08X\n", overlayOffset);
    printError("Errno: %s\n", strerror(errno));
    fclose(fp);
    return 1;
  }

  // Read Wise overlay header
  WiseOverlayHeader overlayHeader;
  if ((status = readWiseOverlayHeader(fp, &overlayHeader)) != REWERR_OK) {
    printError("Failed to read WiseOverlayHeader.\n");
    fclose(fp);
    return 1;
  }
  freeWiseOverlayHeader(&overlayHeader);

  // Here we arrived at the delated data, each entry followed by a CRC32
  // https://en.wikipedia.org/wiki/DEFLATE
  if (huffmanInitFixedTrees() == false) {
    printError("Failed to huffmanInitFixedTrees, out of mem?\n");
    fclose(fp);
    return 1;
  }

  // Initial check on free disk space (TMP_PATH)
  unsigned long tmpFreeDiskSpace = getFreeDiskSpace(TempPath);
  if (tmpFreeDiskSpace == 0) { // failed to determine free disk space
    fclose(fp);
    return 1;
  }
  // make sure at-least 1 MiB is available at the TMP path
  if (tmpFreeDiskSpace < SIZE_MiB) {
    printError("At-least 1 MiB of free space is required in the TMP_PATH.\n");
    fclose(fp);
    return 1;
  }

  bool result;
  InflateObject inflateObj;
  inflateInit(&inflateObj, fp);
  InflateObjPtr = &inflateObj;

  // Raw extract
  if (operation == OP_EXTRACT_RAW) {
    uint32_t extractCount = 0;
    char extractFilePath[PATH_MAX];

    // Start inflating and outputting files
    while (ftell(fp) < inflateObj.inputFileSize) {
      char fileName[21];
      if (snprintf(fileName, 20, "EXTRACTED_%09u", extractCount) > 20) {
        // truncated
        printError("Failed to format filename, it truncated.\n");
        fclose(fp);
        huffmanFreeFixedTrees();
        return 1;
      }
      if (preparePath(OutputPath, fileName, extractFilePath) == false) {
        printError("Failed to create directories for '%s'.\n", fileName);
        fclose(fp);
        huffmanFreeFixedTrees();
        return 1;
      }

      result = inflateExtractNextFile(&inflateObj, (const char *)extractFilePath);
      if (result == false) {
        printError("Failed to extract '%s'.\n", extractFilePath);
        fclose(fp);
        huffmanFreeFixedTrees();
        return 1;
      }

      printInfo("Extracted '%s'\n", extractFilePath);
      extractCount++;
    }

    printInfo("Extracted %d files.\n", extractCount);
  }

  else {
    char tmpFileScript[PATH_MAX];

    // Skip WiseColors.dib
    if (NoExtract == 0) {
      result = inflateExtractNextFile(&inflateObj, NULL);
      if (result == false) {
        printError("Failed to extract 'WiseColors.dib'.\n");
        fclose(fp);
        huffmanFreeFixedTrees();
        return 1;
      }
    }

    // Create filepath for WiseScript.bin
    if (preparePath(TempPath, "WiseScript.bin", tmpFileScript) == false) {
      fclose(fp);
      huffmanFreeFixedTrees();
      printf("Failed to create filepath for WiseScript.bin.\n");
      return 1;
    }
    // Extract WiseScript.bin
    if (NoExtract == 0) {
      result = inflateExtractNextFile(&inflateObj, tmpFileScript);
      if (result == false) {
        printError("Failed to extract '%s'.\n", tmpFileScript);
        fclose(fp);
        huffmanFreeFixedTrees();
        return 1;
      }
    }

    // Determine the inflate data offset inside WiseScript.bin (this needs to
    // be added to the inflateStart we got for files from WiseScript to get to
    // the real inflateStart offset in the PE file.)
    WiseScriptParsedInfo * parsedInfo = wiseScriptGetParsedInfo(tmpFileScript);
    ScriptDeflateOffset = inflateObj.inputFileSize - parsedInfo->inflateStartOffset;
    printDebug("scriptDeflateOffset: %ld (0x%08X).\n", parsedInfo->inflateStartOffset);

    WiseScriptCallbacks callbacks;
    initWiseScriptCallbacks(&callbacks);

    // LIST
    if (operation == OP_LIST) {
      callbacks.cb_0x00 = &printFile;
      printf("    FILESIZE FILEDATE   FILETIME FILEPATH\n");
      printf("------------ ---------- -------- ----------------------------\n");
      status = parseWiseScript(tmpFileScript, &callbacks);
      if (status != REWERR_OK) {
        printError("Parsing WiseScript failed.\n");
      }
      printf("------------ ---------- -------- ----------------------------\n");
      printf("Total size: ");
      printPrettySize(parsedInfo->inflatedSize0x00);
      printf(" (%zu bytes)\n", parsedInfo->inflatedSize0x00);
    }
    // EXTRACT
    else
    if (operation == OP_EXTRACT) {
      // Check if there is enough free disk space
      unsigned long outputFreeDiskSpace = getFreeDiskSpace(OutputPath);
      if (outputFreeDiskSpace == 0) { // failed to determine free disk space
        fclose(fp);
        return 1;
      }
      if (outputFreeDiskSpace <= parsedInfo->inflatedSize0x00) {
        printError("Not enough free disk space at '%s'. Required: %ld Left: "
                   "%ld\n", OutputPath, parsedInfo->inflatedSize0x00,
                   outputFreeDiskSpace);
        fclose(fp);
        return 1;
      }

      // Start inflating and outputting files
      callbacks.cb_0x00 = &extractFile;
      status = parseWiseScript(tmpFileScript, &callbacks);

      // Something went wrong
      if (status != REWERR_OK) {
        printError("Parsing WiseScript failed.\n");
      }
    }
    // SCRIPT_DEBUG
    else
    if (operation == OP_SCRIPT_DEBUG) {
      status = wiseScriptDebugPrint(tmpFileScript);
      if (status != REWERR_OK) {
        printError("Debug print WiseScript failed.\n");
      }
    }
    else
    if (operation == OP_VERIFY) {
      callbacks.cb_0x00 = &noExtractFile;
      status = parseWiseScript(tmpFileScript, &callbacks);
      if (status != REWERR_OK) {
        printError("Parsing WiseScript failed.\n");
      }
      printInfo("All looks good!\n");
    }

    // remove tmp files
    if (PreserveTmp == 0 && NoExtract == 0) {
      if (remove(tmpFileScript) != 0) {
        printError("Failed to remove '%s'. Errno: %s\n", tmpFileScript,
                   strerror(errno));
      }
    }
  }

  // Cleanup
  huffmanFreeFixedTrees();
  fclose(fp);

  return status;
}