/*
  VitaShell
  Copyright (C) 2015-2018, TheFloW

  This program 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.

  This program 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 <http://www.gnu.org/licenses/>.
*/

#include "main.h"
#include "archive.h"
#include "file.h"
#include "utils.h"
#include "elf.h"

static char archive_file[MAX_PATH_LENGTH];
static int archive_path_start = 0;
struct archive *archive_fd = NULL;

int checkForUnsafeImports(void *buffer);
char *uncompressBuffer(const Elf32_Ehdr *ehdr, const Elf32_Phdr *phdr, const segment_info *segment,
                       const char *buffer);

void waitpid() {}
void __archive_create_child() {}
void __archive_check_child() {}

typedef struct ArchiveFileNode {
  struct ArchiveFileNode *parent;
  struct ArchiveFileNode *child;
  struct ArchiveFileNode *next;
  char name[MAX_NAME_LENGTH];
  int name_length;
  int is_folder;
  SceOff size;
  SceDateTime ctime;
  SceDateTime mtime;
  SceDateTime atime;
} ArchiveFileNode;

static ArchiveFileNode *archive_root = NULL;

char *serializePathName(char *name, char **p) {
  if (!p)
    return name;
  
  if (*p)
    name = *p + 1;
  
  *p = strchr(name, '/');
  if (*p)
    **p = '\0';
  
  return name;
}

ArchiveFileNode *createArchiveNode(const char *name, const struct stat *stat, int is_folder, ArchiveFileNode *parent) {
  ArchiveFileNode *node = malloc(sizeof(ArchiveFileNode));
  if (!node)
    return NULL;
  
  memset(node, 0, sizeof(ArchiveFileNode));
  
  node->parent = parent;
  node->child = NULL;
  node->next = NULL;
  if (is_folder || stat->st_mode & S_IFDIR)
    node->is_folder = 1;
  node->name_length = strlen(name);
  strcpy(node->name, name);
  
  if (stat) {
    SceDateTime time;

    node->size = stat->st_size;
        
    sceRtcSetTime_t(&time, stat->st_ctime);
    convertLocalTimeToUtc(&node->ctime, &time);
    
    sceRtcSetTime_t(&time, stat->st_mtime);
    convertLocalTimeToUtc(&node->mtime, &time);
    
    sceRtcSetTime_t(&time, stat->st_atime);
    convertLocalTimeToUtc(&node->atime, &time);
  }
  
  return node;
}

ArchiveFileNode *_findArchiveNode(ArchiveFileNode *parent, char *name, ArchiveFileNode **parent_out,
                                  ArchiveFileNode **prev_out, char **name_out, char **p_out) {
  char *p = NULL;
  ArchiveFileNode *prev = NULL;

  // Remove trailing slash
  removeEndSlash(name);
  
  // Serialize path name
  name = serializePathName(name, &p);
  
  // Root
  if (name[0] == '\0')
    return parent;
  
  // Traverse
  ArchiveFileNode *curr = parent->child;
  while (curr) {
    // Found name
    if (strcmp(curr->name, name) == 0) {
      // Found node
      if (!p)
        return curr;
      
      // Is folder
      if (curr->is_folder) {
        // Serialize path name
        name = serializePathName(name, &p);
        
        // Get child entry of this directory
        parent = curr;
        curr = curr->child;
        continue;
      }
    }
    
    // Get next entry in this directory
    prev = curr;
    curr = curr->next;
  }
  
  // Out
  if (parent_out)
    *parent_out = parent;
  if (prev_out)
    *prev_out = prev;
  if (name_out)
    *name_out = name;
  if (p_out)
    *p_out = p;
  
  return NULL;
}

ArchiveFileNode *findArchiveNode(const char *path) {
  char name[MAX_PATH_LENGTH];
  strcpy(name, path);
  return _findArchiveNode(archive_root, name, NULL, NULL, NULL, NULL);
}

int addArchiveNodeRecursive(ArchiveFileNode *parent, char *name, const struct stat *stat) {  
  char *p = NULL;
  ArchiveFileNode *prev = NULL;
  
  if (!parent)
    return -1;
  
  ArchiveFileNode *res = _findArchiveNode(parent, name, &parent, &prev, &name, &p);
  
  // Already exist
  // TODO: update node
  if (res)
    return 0;
  
  // Create new node
  ArchiveFileNode *node = createArchiveNode(name, stat, !!p, parent);  
  if (!parent->child) { // First child
    parent->child = node;
  } else {              // Neighbour
    prev->next = node;
  }
  
  // Recursion
  if (p)
    addArchiveNodeRecursive(node, p + 1, stat);
  
  return 0;
}

void addArchiveNode(const char *path, const struct stat *stat) {  
  char name[MAX_PATH_LENGTH];
  strcpy(name, path);
  addArchiveNodeRecursive(archive_root, name, stat);
}

void freeArchiveNodes(ArchiveFileNode *curr) {
  // Traverse
  while (curr) {
    // Enter directory
    if (curr->is_folder) {
      // Traverse child entry of this directory
      freeArchiveNodes(curr->child);
    }
    
    // Get next entry in this directory
    ArchiveFileNode *next = curr->next;
    free(curr);
    curr = next;
  }
}

int archiveCheckFilesForUnsafeFself() {
  int res;
  
  // Initialize
  struct archive *archive = archive_read_new();
  archive_read_support_filter_all(archive);
  archive_read_support_format_all(archive);

  // Open archive file
  res = archive_read_open_filename(archive, archive_file, 10240);
  if (res)
    return 0;
  
  // Traverse
  while (1) {
    struct archive_entry *archive_entry;
    res = archive_read_next_header(archive, &archive_entry);
    if (res == ARCHIVE_EOF)
      break;
    
    if (res != ARCHIVE_OK) {
      archive_read_free(archive);
      return 0;
    }
    
    // Get entry information
    const char *name = archive_entry_pathname(archive_entry);
    const struct stat *stat = archive_entry_stat(archive_entry);
        
    // Read magic
    uint32_t magic = 0;
    archive_read_data(archive, &magic, sizeof(uint32_t));
    
    // SCE magic
    if (magic == 0x00454353) {
      char sce_header[0x84];
      archive_read_data(archive, sce_header, sizeof(sce_header));

      uint64_t elf1_offset = *(uint64_t *)(sce_header + 0x3C);
      uint64_t phdr_offset = *(uint64_t *)(sce_header + 0x44);
      uint64_t section_info_offset = *(uint64_t *)(sce_header + 0x54);

      // jump to elf1
      // Until here we have read 0x88 bytes
      int i;
      for (i = 0; i < elf1_offset - 0x88; i += sizeof(uint32_t)) {
        uint32_t dummy = 0;
        archive_read_data(archive, &dummy, sizeof(uint32_t));
      }

      // Check imports
      char *buffer = malloc(stat->st_size);
      if (buffer) {
        archive_read_data(archive, buffer, stat->st_size);

        Elf32_Ehdr *elf1 = (Elf32_Ehdr*)buffer;
        Elf32_Phdr *phdr = (Elf32_Phdr*)(buffer + phdr_offset - elf1_offset);
        segment_info *info = (segment_info*)(buffer + section_info_offset - elf1_offset);
        
        // segment is elf2 section
        char *segment = buffer + info->offset - elf1_offset;

        // zlib compress magic
        char *uncompressed_buffer = NULL;
        if (segment[0] == 0x78) {
          // uncompressedBuffer will return elf2 section
          uncompressed_buffer = uncompressBuffer(elf1, phdr, info, segment);
          if (uncompressed_buffer) {
            segment = uncompressed_buffer;
          }
        }
        
        int unsafe = checkForUnsafeImports(segment);

        if (uncompressed_buffer)
          free(uncompressed_buffer);
        free(buffer);

        if (unsafe) {
          archive_read_free(archive);
          return unsafe;
        }
      }

      // Check authid flag
      uint64_t authid = *(uint64_t *)(sce_header + 0x7C);
      if (authid != 0x2F00000000000002) {
        archive_read_free(archive);
        return 1; // Unsafe
      }
    }
  }

  archive_read_free(archive);
  return 0;
}

int fileListGetArchiveEntries(FileList *list, const char *path, int sort) {
  int res;

  if (!list)
    return -1;

  FileListEntry *entry = malloc(sizeof(FileListEntry));
  strcpy(entry->name, DIR_UP);
  entry->name_length = strlen(entry->name);
  entry->is_folder = 1;
  entry->type = FILE_TYPE_UNKNOWN;
  fileListAddEntry(list, entry, sort);
  
  // Traverse
  ArchiveFileNode *curr = findArchiveNode(path + archive_path_start);
  if (curr)
    curr = curr->child;
  while (curr) {
    FileListEntry *entry = malloc(sizeof(FileListEntry));

    strcpy(entry->name, curr->name);
    
    entry->is_folder = curr->is_folder;
    if (entry->is_folder) {
      addEndSlash(entry->name);
      entry->type = FILE_TYPE_UNKNOWN;
      list->folders++;
    } else {
      entry->type = getFileType(entry->name);
      list->files++;
    }

    entry->name_length = curr->name_length;
    entry->size = curr->size;
    
    memcpy(&entry->ctime, (SceDateTime *)&curr->ctime, sizeof(SceDateTime));
    memcpy(&entry->mtime, (SceDateTime *)&curr->mtime, sizeof(SceDateTime));
    memcpy(&entry->atime, (SceDateTime *)&curr->atime, sizeof(SceDateTime));
    
    fileListAddEntry(list, entry, sort);
    
    // Get next entry in this directory
    curr = curr->next;
  }

  return 0;
}

int getArchivePathInfo(const char *path, uint64_t *size, uint32_t *folders, uint32_t *files) {
  SceIoStat stat;
  memset(&stat, 0, sizeof(SceIoStat));
  if (archiveFileGetstat(path, &stat) < 0) {
    FileList list;
    memset(&list, 0, sizeof(FileList));
    fileListGetArchiveEntries(&list, path, SORT_NONE);

    FileListEntry *entry = list.head->next; // Ignore ..

    int i;
    for (i = 0; i < list.length - 1; i++) {
      char *new_path = malloc(strlen(path) + strlen(entry->name) + 2);
      snprintf(new_path, MAX_PATH_LENGTH - 1, "%s%s", path, entry->name);

      getArchivePathInfo(new_path, size, folders, files);

      free(new_path);

      entry = entry->next;
    }

    if (folders)
      (*folders)++;

    fileListEmpty(&list);
  } else {
    if (size)
      (*size) += stat.st_size;

    if (files)
      (*files)++;
  }

  return 0;
}

int extractArchivePath(const char *src, const char *dst, FileProcessParam *param) {
  SceIoStat stat;
  memset(&stat, 0, sizeof(SceIoStat));
  if (archiveFileGetstat(src, &stat) < 0) {
    FileList list;
    memset(&list, 0, sizeof(FileList));
    fileListGetArchiveEntries(&list, src, SORT_NONE);

    int ret = sceIoMkdir(dst, 0777);
    if (ret < 0 && ret != SCE_ERROR_ERRNO_EEXIST) {
      fileListEmpty(&list);
      return ret;
    }

    if (param) {
      if (param->value)
        (*param->value) += DIRECTORY_SIZE;

      if (param->SetProgress)
        param->SetProgress(param->value ? *param->value : 0, param->max);
      
      if (param->cancelHandler && param->cancelHandler()) {
        fileListEmpty(&list);
        return 0;
      }
    }

    FileListEntry *entry = list.head->next; // Ignore ..

    int i;
    for (i = 0; i < list.length - 1; i++) {
      char *src_path = malloc(strlen(src) + strlen(entry->name) + 2);
      snprintf(src_path, MAX_PATH_LENGTH - 1, "%s%s", src, entry->name);

      char *dst_path = malloc(strlen(dst) + strlen(entry->name) + 2);
      snprintf(dst_path, MAX_PATH_LENGTH - 1, "%s%s", dst, entry->name);

      int ret = extractArchivePath(src_path, dst_path, param);

      free(dst_path);
      free(src_path);

      if (ret <= 0) {
        fileListEmpty(&list);
        return ret;
      }

      entry = entry->next;
    }

    fileListEmpty(&list);
  } else {
    SceUID fdsrc = archiveFileOpen(src, SCE_O_RDONLY, 0);
    if (fdsrc < 0)
      return fdsrc;

    SceUID fddst = sceIoOpen(dst, SCE_O_WRONLY | SCE_O_CREAT | SCE_O_TRUNC, 0777);
    if (fddst < 0) {
      archiveFileClose(fdsrc);
      return fddst;
    }

    void *buf = malloc(TRANSFER_SIZE);

    uint64_t seek = 0;

    while (1) {
      int read = archiveFileRead(fdsrc, buf, TRANSFER_SIZE);
      if (read < 0) {
        free(buf);

        sceIoClose(fddst);
        archiveFileClose(fdsrc);

        return read;
      }

      if (read == 0)
        break;

      int written = sceIoWrite(fddst, buf, read);

      if (written < 0) {
        free(buf);

        sceIoClose(fddst);
        archiveFileClose(fdsrc);

        return written;
      }

      seek += written;

      if (param) {
        if (param->value)
          (*param->value) += read;

        if (param->SetProgress)
          param->SetProgress(param->value ? *param->value : 0, param->max);

        if (param->cancelHandler && param->cancelHandler()) {
          free(buf);

          sceIoClose(fddst);
          archiveFileClose(fdsrc);

          return 0;
        }
      }
    }

    free(buf);

    sceIoClose(fddst);
    archiveFileClose(fdsrc);
  }

  return 1;
}

int archiveFileGetstat(const char *file, SceIoStat *stat) {
  // Is directory
  if (hasEndSlash(file + archive_path_start))
    return -1;
  
  ArchiveFileNode *node = findArchiveNode(file + archive_path_start);
  if (!node)
    return -1;
  
  if (stat) {
    stat->st_size = node->size;
    memcpy(&stat->st_ctime, &node->ctime, sizeof(SceDateTime));
    memcpy(&stat->st_mtime, &node->mtime, sizeof(SceDateTime));
    memcpy(&stat->st_atime, &node->atime, sizeof(SceDateTime));
  }
  
  return 0;
}

int archiveFileOpen(const char *file, int flags, SceMode mode) {
  int res;
  
  // A file is already open
  if (archive_fd)
    return -1;

  // Initialize
  archive_fd = archive_read_new();
  archive_read_support_filter_all(archive_fd);
  archive_read_support_format_all(archive_fd);

  // Open archive file
  res = archive_read_open_filename(archive_fd, archive_file, 10240);
  if (res)
    return -1;

  // Traverse
  while (1) {
    struct archive_entry *archive_entry;
    res = archive_read_next_header(archive_fd, &archive_entry);
    if (res == ARCHIVE_EOF)
      break;
    
    if (res != ARCHIVE_OK) {
      archive_read_free(archive_fd);
      return -1;
    }
    
    // Compare pathname
    const char *name = archive_entry_pathname(archive_entry);
    if (strcmp(name, file + archive_path_start) == 0) {
      return ARCHIVE_FD;
    }
  }

  archive_read_free(archive_fd);
  archive_fd = NULL;

  return -1;
}

int archiveFileRead(SceUID fd, void *data, SceSize size) {
  if (!archive_fd || fd != ARCHIVE_FD)
    return -1;

  return archive_read_data(archive_fd, data, size);
}

int archiveFileClose(SceUID fd) {
  if (!archive_fd || fd != ARCHIVE_FD)
    return -1;

  archive_read_free(archive_fd);
  archive_fd = NULL;

  return 0;
}

int ReadArchiveFile(const char *file, void *buf, int size) {
  SceUID fd = archiveFileOpen(file, SCE_O_RDONLY, 0);
  if (fd < 0)
    return fd;

  int read = archiveFileRead(fd, buf, size);
  archiveFileClose(fd);
  return read;
}

int archiveClose() {
  freeArchiveNodes(archive_root->child);
  free(archive_root);
  return 0;
}

int archiveOpen(const char *file) {
  int res;

  // Start position of the archive path
  archive_path_start = strlen(file) + 1;
  strcpy(archive_file, file);
  
  // Initialize
  struct archive *archive = archive_read_new();
  archive_read_support_filter_all(archive);
  archive_read_support_format_all(archive);

  // Open archive file
  res = archive_read_open_filename(archive, file, 10240);
  if (res)
    return -1;

  // Create archive root
  archive_root = createArchiveNode("/", NULL, 1, NULL);
  
  // Traverse
  while (1) {
    struct archive_entry *archive_entry;
    res = archive_read_next_header(archive, &archive_entry);
    if (res == ARCHIVE_EOF)
      break;
    
    if (res != ARCHIVE_OK) {
      archive_read_free(archive);
      return -1;
    }
    
    // Get entry information
    const char *name = archive_entry_pathname(archive_entry);
    const struct stat *stat = archive_entry_stat(archive_entry);
    
    // Add node
    addArchiveNode(name, stat);
  }

  archive_read_free(archive);
  return 0;
}