/*
PSFLIB - PSF2FS implementation

Copyright (c) 2012-2015 Christopher Snowhill

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

#include "psf2fs.h"

#include <ctype.h>
#include <string.h>
#include <stdlib.h>

#include <zlib.h>

/////////////////////////////////////////////////////////////////////////////

#define MYMAXPATH (1024)

struct SOURCE_FILE {
  uint8_t * reserved_data;
  int reserved_size;
  struct SOURCE_FILE *next;
};

struct DIR_ENTRY {
  char name[37];
  struct DIR_ENTRY *subdir;
  int length;
  int block_size;
  struct SOURCE_FILE *source;
  int *offset_table;
  struct DIR_ENTRY *next;
};

struct CACHEBLOCK {
  struct SOURCE_FILE *from_source;
  int   from_offset;
  char *uncompressed_data;
  int   uncompressed_size;
};

struct PSF2FS {
  struct SOURCE_FILE *sources;
  struct DIR_ENTRY *dir;
  struct CACHEBLOCK cacheblock;

  int adderror;
};

/////////////////////////////////////////////////////////////////////////////

static void source_cleanup_free(struct SOURCE_FILE *source) {
  while(source) {
    struct SOURCE_FILE *next = source->next;
    if(source->reserved_data) free( source->reserved_data );
    free( source );
    source = next;
  }
}

static void dir_cleanup_free(struct DIR_ENTRY *dir) {
  while(dir) {
    struct DIR_ENTRY *next = dir->next;
    if(dir->subdir) dir_cleanup_free(dir->subdir);
    free( dir );
    dir = next;
  }
}

static void cache_cleanup(struct CACHEBLOCK *cacheblock) {
  if(!cacheblock) return;
  if(cacheblock->uncompressed_data) free( cacheblock->uncompressed_data );
}

/////////////////////////////////////////////////////////////////////////////

void *psf2fs_create(void) {
  struct PSF2FS *fs;
  fs = ( struct PSF2FS * ) malloc( sizeof( struct PSF2FS ) );
  if(!fs) return NULL;
  memset(fs, 0, sizeof(struct PSF2FS));
  return fs;
}

/////////////////////////////////////////////////////////////////////////////

void psf2fs_delete(void *psf2fs) {
  struct PSF2FS *fs = (struct PSF2FS*)psf2fs;
  if(fs->sources) source_cleanup_free(fs->sources);
  if(fs->dir) dir_cleanup_free(fs->dir);
  cache_cleanup(&(fs->cacheblock));
  free( fs );
}

/////////////////////////////////////////////////////////////////////////////

static int isdirsep(char c) { return (c == '/' || c == '\\' || c == '|' || c == ':'); }

/////////////////////////////////////////////////////////////////////////////

static unsigned read32lsb(const uint8_t * foo) {
  return (
    ((foo[0] & 0xFF) <<  0) |
    ((foo[1] & 0xFF) <<  8) |
    ((foo[2] & 0xFF) << 16) |
    ((foo[3] & 0xFF) << 24)
  );
}

/////////////////////////////////////////////////////////////////////////////

static int __memicmp(const char * a, const char * b, int length)
{
  int o, p;
  for (o = 0; o < length; o++) {
    p = tolower(a[o]) - tolower(b[o]);
    if (p) return p;
  }
  return 0;
}

/////////////////////////////////////////////////////////////////////////////

static struct DIR_ENTRY *finddirentry(
  struct DIR_ENTRY *dir,
  const char *name,
  int name_l
) {
  if(name_l > 36) return NULL;
  while(dir) {
    if(!__memicmp(dir->name, name, name_l) && dir->name[name_l] == 0) return dir;
    dir = dir->next;
  }
  return NULL;
}

/////////////////////////////////////////////////////////////////////////////
//
// Make a DIR_ENTRY list for a given file and Reserved offset.
// Recurses subdirectories also.
// All entries are set to point to the given SOURCE_FILE.
//
static struct DIR_ENTRY *makearchivedir(
  struct PSF2FS *fs,
  int offset,
  struct SOURCE_FILE *source
) {
  struct DIR_ENTRY *dir = NULL;
  const uint8_t *file = source->reserved_data;
  int n, num;
  if(offset < 0) goto corrupt;
  if(offset >= source->reserved_size) { goto corrupt; }
  if((offset + 4) > source->reserved_size) { goto corrupt; }
  num = read32lsb(file + offset);
  offset += 4;
  if(num < 0) goto corrupt;
  for(n = 0; n < num; n++) {
    int o, u, b;
    if((offset + 48) > source->reserved_size) { goto corrupt; }
    { struct DIR_ENTRY *entry = ( struct DIR_ENTRY * ) malloc( sizeof( struct DIR_ENTRY ) );
      if(!entry) goto outofmemory;
      memset(entry, 0, sizeof(struct DIR_ENTRY));
      entry->next = dir;
      dir = entry;
    }
    memcpy(dir->name, file + offset, 36);
    o = read32lsb(file + offset + 36);
    u = read32lsb(file + offset + 40);
    b = read32lsb(file + offset + 44);
    offset += 48;
    if(o < 0) goto corrupt;
    if(u < 0) goto corrupt;
    if(b < 0) goto corrupt;
    if(o && o < offset) {
//      char s[100];
//      sprintf(s,"q[o=%08X offset=%08X]",o,offset);
//      errormessageadd(fs, s);
      goto corrupt;
    }
    // if this new entry describes a subdirectory:
    if(u == 0 && b == 0 && o != 0) {
      dir->subdir = makearchivedir(fs, o, source);
      if(fs->adderror) goto error;
//      if(!dir->subdir) goto error;
    // if this new entry describes a zero-length file:
    } else if(u == 0 || b == 0 || o == 0) {
      // fields were zero anyway
    // if this new entry describes a real source file:
    } else {
      int i;
      int blocks = (u + (b-1)) / b;
      int dataofs = o + 4 * blocks;
      if(dataofs >= source->reserved_size) { goto corrupt; }
      // record the info
      dir->length = u;
      dir->block_size = b;
      dir->source = source;
      dir->offset_table = (int *) malloc( ( blocks + 1 ) * sizeof( int ) );
      if(!dir->offset_table) goto outofmemory;
      for(i = 0; i < blocks; i++) {
        int cbs;
        if((o + 4) > source->reserved_size) { goto corrupt; }
        cbs = read32lsb(file + o);
        o += 4;
        dir->offset_table[i] = dataofs;
        dataofs += cbs;
      }
      dir->offset_table[i] = dataofs;
    }
  }
//success:
  return dir;

corrupt:
  goto error;
outofmemory:
  goto error;
error:
  dir_cleanup_free(dir);
  fs->adderror = 1;
  return NULL;
}

/////////////////////////////////////////////////////////////////////////////
//
// Merge two SOURCE_FILE lists.
// Guaranteed to succeed and not to free anything.
//
static struct SOURCE_FILE *mergesource(
  struct SOURCE_FILE *to,
  struct SOURCE_FILE *from
) {
  struct SOURCE_FILE *to_tail;
  if(!to && !from) return NULL;
  if(!to) {
    struct SOURCE_FILE *t;
    t = to; to = from; from = t;
  }
  to_tail = to;
  while(to_tail->next) { to_tail = to_tail->next; }
  to_tail->next = from;
  return to;
}

/////////////////////////////////////////////////////////////////////////////
//
// Merge two DIR_ENTRY lists.
// Guaranteed to succeed. May free some structures.
//
static struct DIR_ENTRY *mergedir(
  struct DIR_ENTRY *to,
  struct DIR_ENTRY *from
) {
  // will traverse "from", and add to "to".
  while(from) {
    struct DIR_ENTRY *entry_to;
    struct DIR_ENTRY *entry_from;
    entry_from = from;
    from = from->next;
    // delink entry_from
    entry_from->next = NULL;
    // look for a duplicate entry in "to"
    entry_to = finddirentry(to, entry_from->name, (int)strlen(entry_from->name));
    // if there is one, do something fancy and then free entry_from.
    if(entry_to) {
      // if both are subdirs, merge the subdirs
      if((entry_to->subdir) && (entry_from->subdir)) {
        entry_to->subdir = mergedir(entry_to->subdir, entry_from->subdir);
        entry_from->subdir = NULL;
      // if both are files, copy over the info
      } else if((!(entry_to->subdir)) && (!(entry_from->subdir))) {
        entry_to->length = entry_from->length;
        entry_to->block_size = entry_from->block_size;
        entry_to->source = entry_from->source;
        if(entry_to->offset_table) free( entry_to->offset_table );
        entry_to->offset_table = entry_from->offset_table;
        entry_from->offset_table = NULL;
      // if one's a subdir but the other's not, we lose "from". this is fine.
      }
      dir_cleanup_free(entry_from);
      entry_from = NULL;
    // otherwise, just relink to the top of "to"
    } else {
      entry_from->next = to;
      to = entry_from;
    }
  }
  return to;
}

/////////////////////////////////////////////////////////////////////////////
//
// only modifies *psource and *pdir on success
//
static int addarchive(
  struct PSF2FS *fs,
  const uint8_t *reserved_data,
  int reserved_size,
  struct SOURCE_FILE **psource,
  struct DIR_ENTRY **pdir
) {
  struct SOURCE_FILE *source = *psource;
  struct DIR_ENTRY *dir = *pdir;
  // these relate to the current file
  struct SOURCE_FILE *this_source = NULL;
  struct DIR_ENTRY *this_dir = NULL;

  // default to no error
  fs->adderror = 0;

  // create a source entry for this psf2
  this_source = ( struct SOURCE_FILE * ) malloc( sizeof( struct SOURCE_FILE ) );
  if(!this_source) goto outofmemory;
  this_source->reserved_data = ( uint8_t * ) malloc( reserved_size );
  if(!this_source->reserved_data) goto outofmemory;
  memcpy(this_source->reserved_data, reserved_data, reserved_size);
  this_source->reserved_size = reserved_size;
  this_source->next = NULL;
  this_dir = makearchivedir(fs, 0, this_source);
  if(fs->adderror) goto error;

  // success
  // now merge everything
  *psource = mergesource(source, this_source);
  *pdir = mergedir(dir, this_dir);
//success:
  return 0;

outofmemory:
  goto error;
error:
  if(dir) dir_cleanup_free(dir);
  if(source) source_cleanup_free(source);
  if(this_dir) dir_cleanup_free(this_dir);
  if(this_source) source_cleanup_free(this_source);
  return -1;
}

/////////////////////////////////////////////////////////////////////////////
//
//
//
int psf2fs_load_callback(void * psf2fs, const uint8_t * exe, size_t exe_size,
                                  const uint8_t * reserved, size_t reserved_size) {
  struct PSF2FS *fs = (struct PSF2FS*)psf2fs;
  (void)exe;
  (void)exe_size;
  return addarchive(fs, reserved, (int)reserved_size, &(fs->sources), &(fs->dir));
}

/////////////////////////////////////////////////////////////////////////////
//
//
//
static int virtual_read(struct PSF2FS *fs, struct DIR_ENTRY *entry, int offset, char *buffer, int length) {
  int length_read = 0;
  int r;
  unsigned long destlen;
  if(offset >= entry->length) return 0;
  if((offset + length) > entry->length) length = entry->length - offset;
  while(length_read < length) {
    // get info on the current block
    int blocknum = offset / entry->block_size;
    int ofs_within_block = offset % entry->block_size;
    int canread;
    int block_zofs  = entry->offset_table[blocknum];
    int block_zsize = entry->offset_table[blocknum+1] - block_zofs;
    int block_usize;
    if(block_zofs <= 0 || block_zofs >= entry->source->reserved_size) goto bounds;
    if((block_zofs+block_zsize) > entry->source->reserved_size) goto bounds;

    // get the actual uncompressed size of this block
    block_usize = entry->length - (blocknum * entry->block_size);
    if(block_usize > entry->block_size) block_usize = entry->block_size;

    // if it's not already in the cache block, read it
    if(
      (fs->cacheblock.from_offset != block_zofs) ||
      (fs->cacheblock.from_source != entry->source)
    ) {
      // invalidate cache without freeing buffer
      fs->cacheblock.from_source = NULL;

      // make sure there's a buffer allocated
      // but only reallocate if the size is different
      if(fs->cacheblock.uncompressed_size != block_usize) {
        fs->cacheblock.uncompressed_size = 0;
        if(fs->cacheblock.uncompressed_data) {
          free( fs->cacheblock.uncompressed_data );
          fs->cacheblock.uncompressed_data = NULL;
        }
        fs->cacheblock.uncompressed_data = ( char * ) malloc( block_usize );
        if(!fs->cacheblock.uncompressed_data) goto outofmemory;
        fs->cacheblock.uncompressed_size = block_usize;
      }
      destlen = block_usize;
      // attempt decompress
      r = uncompress((unsigned char *) fs->cacheblock.uncompressed_data, &destlen, (const unsigned char *) entry->source->reserved_data + block_zofs, block_zsize);
      if(r != Z_OK || destlen != (unsigned long)block_usize) {
//        char s[999];
//        sprintf(s,"zdata=%02X %02X %02X blockz=%d blocku=%d destlenout=%d", zdata[0], zdata[1], zdata[2], block_zsize, block_usize, destlen);
//        errormessageadd(fs, s);
        goto error;
      }
    }

    // at this point, we can read whatever we want out of the cacheblock
    canread = fs->cacheblock.uncompressed_size - ofs_within_block;
    if(canread > (length - length_read)) canread = length - length_read;

    // copy
    memcpy(buffer, fs->cacheblock.uncompressed_data + ofs_within_block, canread);

    // advance pointers/counters
    offset += canread;
    length_read += canread;
    buffer += canread;
  }

//success:
  return length_read;

bounds:
  goto error;
outofmemory:
  goto error;
error:
  // if cacheblock was invalidated, we can free it
  if(!fs->cacheblock.from_source) {
    fs->cacheblock.uncompressed_size = 0;
    if(fs->cacheblock.uncompressed_data) {
      free( fs->cacheblock.uncompressed_data );
      fs->cacheblock.uncompressed_data = NULL;
    }
  }
  return -1;
}

/////////////////////////////////////////////////////////////////////////////
//
//
//
int psf2fs_virtual_readfile(void *psf2fs, const char *path, int offset, char *buffer, int length) {
  struct PSF2FS *fs = (struct PSF2FS*)psf2fs;
  struct DIR_ENTRY *entry = fs->dir;


  if(!path) goto invalidarg;
  if(offset < 0) goto invalidarg;
  if(!buffer) goto invalidarg;
  if(length < 0) goto invalidarg;

  for(;;) {
    int l;
    int need_dir;
    if(!entry) goto pathnotfound;
    while(isdirsep(*path)) path++;
    for(l = 0;; l++) {
      if(!path[l]) { need_dir = 0; break; }
      if(isdirsep(path[l])) { need_dir = 1; break; }
    }
    entry = finddirentry(entry, path, l);
    if(!entry) goto pathnotfound;
    if(!need_dir) break;
    entry = entry->subdir;
    path += l;
  }

  // if we "found" a file but it's a directory, then we didn't find it
  if(entry->subdir) goto pathnotfound;

  // special case: if requested length is 0, return the total file length
  if(!length) return entry->length;

  // otherwise, read from source
  return virtual_read(fs, entry, offset, buffer, length);

pathnotfound:
  goto error;
invalidarg:
  goto error;
error:
  return -1;
}

/////////////////////////////////////////////////////////////////////////////
