// SPDX-License-Identifier: Unlicense

#include "psarc.h"

#ifdef SHR3D_PSARC

#include "file.h"
#include "global.h"
#include "inflate.h"
#include "rijndael.h"
#include "wem.h"

#include <string.h>
#include <iterator>

static constexpr u8 psarcKey[32] = {
  0xC5, 0x3D, 0xB2, 0x38, 0x70, 0xA1, 0xA2, 0xF7,
  0x1C, 0xAE, 0x64, 0x06, 0x1F, 0xDD, 0x0E, 0x11,
  0x57, 0x30, 0x9D, 0xC8, 0x52, 0x04, 0xD4, 0xC5,
  0xBF, 0xDF, 0x25, 0x09, 0x0D, 0xF2, 0x57, 0x2C
};

static std::vector<u8> decryptPsarc(const u8* psarcData, u64 /*psarcDataSize*/, u32 startPos, const u32 length) {

  // This is a bit strange not sure if the padding stuff is really required.

  const u32 pad = 512 - (length % 512);
  const u32 lengthPadded = length + pad;

  u8* encText = reinterpret_cast<u8*>(malloc(lengthPadded * sizeof(u8)));
  memcpy(encText, &psarcData[startPos], length);
  memset(&encText[length], 0, pad);

  std::vector<u8> plainText(lengthPadded);

  Rijndael::decrypt(psarcKey, encText, plainText.data(), lengthPadded);
  free(encText);

  plainText.resize(length - startPos);

  return plainText;
}

static i32 tocBNum(const u32 blockSizeAlloc) {
  switch (blockSizeAlloc) {
  case 65536:
    return 2;
  }

  ASSERT(false);

  return {};
}

static void inflateTocEntry(Psarc::Info::TOCEntry& tocEntry, const u32 blockSizeAlloc, const u8* psarcData,
  const std::vector<u32>& zBlockSizeList) {
  if (tocEntry.length == 0)
    return;

  tocEntry.content.resize(tocEntry.length);

  const i32 zHeader = 0x78DA;
  u32 zChunkId = tocEntry.zIndexBegin;

  u64 inCur = 0;
  i32 outCur = 0;
  do {
    const u32 blockSize = zBlockSizeList[zChunkId];
    if (blockSize == 0) { // raw. full cluster used.
      memcpy(&tocEntry.content[outCur], &psarcData[tocEntry.offset + inCur], blockSizeAlloc);
      inCur += blockSizeAlloc;
      outCur += blockSizeAlloc;
    }
    else {
      const u16 num = u16_be(&psarcData[tocEntry.offset + inCur]);
      if (num == zHeader) {
        const i32 size = Inflate::inflate(&psarcData[tocEntry.offset + inCur], blockSize, &tocEntry.content[outCur], i32(tocEntry.length - outCur));
        inCur += blockSize;
        outCur += size;
      }
      else { // raw. used only for data(chunks) smaller than 64 kb
        memcpy(&tocEntry.content[outCur], &psarcData[tocEntry.offset + inCur], blockSize);
        inCur += blockSize;
        outCur += blockSize;
      }
    }
    zChunkId += 1;
  } while (outCur < i32(tocEntry.length));
}

static void readManifest(std::vector<Psarc::Info::TOCEntry>& tocEnties, u32 blockSizeAlloc, const u8* psarcData, const std::vector<u32>& zBlockSizeList)
{
  Psarc::Info::TOCEntry& tocEntry = tocEnties[0];

  ASSERT(tocEntry.name.empty());

  tocEntry.name = "NameBlock.bin";
  inflateTocEntry(tocEntry, blockSizeAlloc, psarcData, zBlockSizeList);

  { // read names of toc entries.
    i32 begin = 0;
    i32 tocIndex = 1;
    for (i32 i = 0; i < i32(tocEntry.content.size()); ++i) {
      if (tocEntry.content[i] == '\n') {
        tocEnties[tocIndex++].name.assign(reinterpret_cast<const char*>(&tocEntry.content[begin]), i - begin);
        begin = i + 1;
      }
    }
    tocEnties[tocIndex].name.assign(reinterpret_cast<const char*>(&tocEntry.content[begin]), tocEntry.content.size() - begin);
  }
}

Psarc::Info Psarc::parse(const u8* psarcData, u64 psarcDataSize) {
  Info psarcInfo;

  { // parse Header
    psarcInfo.header.magicNumber = u32_be(&psarcData[0]);
    ASSERT(psarcInfo.header.magicNumber == 1347633490_u32 && "Invalid Psarc content");
    psarcInfo.header.version = u32_be(&psarcData[4]);
    psarcInfo.header.compressMethod = u32_be(&psarcData[8]);
    psarcInfo.header.totalTocSize = u32_be(&psarcData[12]);
    psarcInfo.header.TOCEntrySize = u32_be(&psarcData[16]);
    psarcInfo.header.numFiles = u32_be(&psarcData[20]);
    psarcInfo.header.blockSizeAlloc = u32_be(&psarcData[24]);
    psarcInfo.header.archiveFlags = u32_be(&psarcData[28]);
  }

  { // parse TOC
    psarcInfo.tocRaw = decryptPsarc(psarcData, psarcDataSize, 32, psarcInfo.header.totalTocSize);
    for (u32 i = 0; i < psarcInfo.header.numFiles; ++i) {
      const u64 offset = i * 30;
      Info::TOCEntry tocEntry;
      memcpy(tocEntry.md5, &psarcInfo.tocRaw[offset], 16);
      tocEntry.zIndexBegin = u32_be(&psarcInfo.tocRaw[offset + 16]);
      tocEntry.length = u40_be(&psarcInfo.tocRaw[offset + 20]);
      tocEntry.offset = u40_be(&psarcInfo.tocRaw[offset + 25]);
      psarcInfo.tocEntries.push_back(tocEntry);
    }
  }

  std::vector<u32> zBlockSizeList;
  { // parse zBlockSizeList
    const i32 tocSize = psarcInfo.header.totalTocSize - 32;
    const i32 tocChunkSize = (int)(psarcInfo.header.numFiles * psarcInfo.header.TOCEntrySize);
    const i32 bNum = tocBNum(psarcInfo.header.blockSizeAlloc);
    const i32 zNum = (tocSize - tocChunkSize) / bNum;
    const u64 offset = psarcInfo.header.numFiles * 30;
    zBlockSizeList.resize(zNum);

    for (int i = 0; i < zNum; ++i) {
      switch (bNum) {
      case 2: // 64KB
        zBlockSizeList[i] = u16_be(&psarcInfo.tocRaw[offset + i * 2]);
        break;
      default:
        unreachable();
      }
    }
  }

  { // decompress
    switch (psarcInfo.header.compressMethod) {
    case 2053925218:
      readManifest(psarcInfo.tocEntries, psarcInfo.header.blockSizeAlloc, psarcData,
        zBlockSizeList);
      break;
    default:
      unreachable();
    }
  }

  { // inflate entries
    for (Info::TOCEntry& tocEntry : psarcInfo.tocEntries) {
      if (tocEntry.name == "NameBlock.bin")
        continue;

      inflateTocEntry(tocEntry, psarcInfo.header.blockSizeAlloc, psarcData, zBlockSizeList);
    }
  }

  return psarcInfo;
}

#endif // SHR3D_PSARC
