#include "JSystem/JKernel/JKRDecomp.h"

#include "JSystem/JKernel/JKRAram.h"
#include "JSystem/JKernel/JKRHeap.h"
#include "JSystem/JSystem.h"
#include "dolphin/os/OSMessage.h"
#include "types.h"

OSMessage JKRDecomp::sMessageBuffer[JKRDECOMP_MSG_BUF_COUNT] = { 0 };
OSMessageQueue JKRDecomp::sMessageQueue = { 0 };
JKRDecomp* JKRDecomp::sDecompObject;

JKRDecomp* JKRDecomp::create(s32 decompPriority) {
  if (JKRDecomp::sDecompObject == nullptr) {
    JKRDecomp::sDecompObject =
      new (JKRGetSystemHeap(), 0) JKRDecomp(decompPriority);
  }

  return JKRDecomp::sDecompObject;
}

JKRDecomp::JKRDecomp(s32 priority)
  : JKRThread(JKRDECOMP_STACK_SIZE, JKRDECOMP_THREAD_MSG_BUF_COUNT,
    priority) {
  OSResumeThread(this->mThreadRecord);
}

JKRDecomp::~JKRDecomp() {}

void* JKRDecomp::run() {
  OSMessage recMesg;
  JKRDecompCommand* cmd;
  OSInitMessageQueue(&JKRDecomp::sMessageQueue, JKRDecomp::sMessageBuffer,
    JKRDECOMP_MSG_BUF_COUNT);

  while (true) {
    while (true) {
      while (true) {
        OSReceiveMessage(&JKRDecomp::sMessageQueue, &recMesg, OS_MESSAGE_BLOCK);
        cmd = static_cast<JKRDecompCommand*>(recMesg);
        JKRDecomp::decode(cmd->mSrcBuffer, cmd->mDstBuffer, cmd->mSrcLength,
          cmd->mSkipCount);

        if (cmd->transferType == JKRDecompCommand::MRAM) {
          break;
        }

        if (cmd->transferType == JKRDecompCommand::ARAM) {
          JKRAramPcs_SendCommand(cmd->mAMCommand);
        }
      }

      if (cmd->mCallback == nullptr) {
        break;
      }

      cmd->mCallback((u32)cmd);
    }

    if (cmd->pMesgQueue1C != nullptr) {
      OSSendMessage(cmd->pMesgQueue1C, (OSMessage)1, OS_MESSAGE_NOBLOCK);
    }
    else {
      OSSendMessage(&cmd->mMesgQueue, (OSMessage)1, OS_MESSAGE_NOBLOCK);
    }
  }
}

JKRDecompCommand* JKRDecomp::prepareCommand(u8* srcBuffer, u8* dstBuffer,
  u32 srcLength, u32 skipCount,
  DecompCallback* callback) {
  JKRDecompCommand* cmd = new (JKRGetSystemHeap(), -4) JKRDecompCommand();

  cmd->mSrcBuffer = srcBuffer;
  cmd->mDstBuffer = dstBuffer;
  cmd->mSrcLength = srcLength;
  cmd->mSkipCount = skipCount;
  cmd->mCallback = callback;

  return cmd;
}

BOOL JKRDecomp::sendCommand(JKRDecompCommand* cmd) {
  BOOL res = OSSendMessage(&JKRDecomp::sMessageQueue, (OSMessage)cmd,
    OS_MESSAGE_BLOCK);

#ifdef JSYSTEM_DEBUG
  if (res == FALSE) {
    JPANIC(142, "Decomp MesgBuf FULL!");
  }
#endif

  return res;
}

JKRDecompCommand* JKRDecomp::orderAsync(u8* srcBuffer, u8* dstBuffer,
  u32 srcLength, u32 skipCount,
  DecompCallback* callback) {
  JKRDecompCommand* cmd = JKRDecomp::prepareCommand(
    srcBuffer, dstBuffer, srcLength, skipCount, callback);
  JKRDecomp::sendCommand(cmd);
  return cmd;
}

bool JKRDecomp::sync(JKRDecompCommand* cmd, BOOL noBlock) {
  OSMessage msg;

  if (!noBlock) {
    OSReceiveMessage(&cmd->mMesgQueue, &msg, OS_MESSAGE_BLOCK);
    return true;
  }
  else {
    return OSReceiveMessage(&cmd->mMesgQueue, &msg, OS_MESSAGE_NOBLOCK) !=
      FALSE;
  }
}

bool JKRDecomp::orderSync(u8* srcBuffer, u8* dstBuffer, u32 srcLength,
  u32 skipCount) {
  JKRDecompCommand* cmd = JKRDecomp::orderAsync(srcBuffer, dstBuffer, srcLength,
    skipCount, nullptr);
  bool res = JKRDecomp::sync(cmd, FALSE);
  delete cmd;
  return res;
}

void JKRDecomp::decode(u8* srcBuffer, u8* dstBuffer, u32 srcLength,
  u32 skipCount) {
  CompressionMode mode = JKRDecomp::checkCompressed(srcBuffer);
  if (mode == SZP) {
    JKRDecomp::decodeSZP(srcBuffer, dstBuffer, srcLength, skipCount);
  }
  else if (mode == SZS) {
    JKRDecomp::decodeSZS(srcBuffer, dstBuffer, srcLength, skipCount);
  }
}

void JKRDecomp::decodeSZP(u8* src, u8* dst, u32 srcLength, u32 skipCount) {
  int srcChunkOffset;
  int count;
  int dstOffset;
  u32 length;
  int linkInfo;
  int offset;
  int i;

  int decodedSize = JKRDECOMP_READU32BE(src, 4);
  int linkTableOffset = JKRDECOMP_READU32BE(src, 8);
  int srcDataOffset = JKRDECOMP_READU32BE(src, 12);

  dstOffset = 0;
  u32 counter =
    0;  // curently counter gets assembled before the READ_U32 operations
  srcChunkOffset = 16;

  u32 chunkBits;
  if (srcLength == 0) return;
  if (skipCount > decodedSize) return;

  length = srcLength;
  do {
    if (counter == 0) {
      chunkBits = JKRDECOMP_READU32BE(src, srcChunkOffset);
      srcChunkOffset += sizeof(u32);
      counter = sizeof(u32) * 8;
    }

    if (chunkBits & 0x80000000) {
      if (skipCount == 0) {
        dst[dstOffset] = src[srcDataOffset];
        length--;
        if (length == 0) return;
      }
      else {
        skipCount--;
      }
      dstOffset++;
      srcDataOffset++;
    }
    else {
      linkInfo = src[linkTableOffset] << 8 | src[linkTableOffset + 1];
      linkTableOffset += sizeof(u16);

      offset = dstOffset - (linkInfo & 0xFFF);
      count = (linkInfo >> 12);
      if (count == 0) {
        count = (u32)src[srcDataOffset++] + 0x12;
      }
      else
        count += 2;

      if ((int)count > decodedSize - dstOffset) count = decodedSize - dstOffset;

      for (i = 0; i < (int)count; i++, dstOffset++, offset++) {
        if (skipCount == 0) {
          dst[dstOffset] = dst[offset - 1];
          length--;
          if (length == 0) return;
        }
        else
          skipCount--;
      }
    }

    chunkBits <<= 1;
    counter--;
  } while (dstOffset < decodedSize);
}

void JKRDecomp::decodeSZS(u8* src_buffer, u8* dst_buffer, u32 srcSize,
  u32 skipCount) {
  u8* decompEnd = dst_buffer + *(u32*)(src_buffer + 4) - skipCount;
  u8* copyStart;
  s32 copyByteCount;
  s32 chunkBitsLeft = 0;
  s32 chunkBits;

  if (srcSize == 0) return;
  if (skipCount > *(u32*)src_buffer) return;

  u8* curSrcPos = src_buffer + 0x10;
  do {
    if (chunkBitsLeft == 0) {
      chunkBits = *curSrcPos++;
      chunkBitsLeft = 8;
    }
    if ((chunkBits & 0x80) != 0) {
      if (skipCount == 0) {
        *dst_buffer = *curSrcPos;
        srcSize--;
        dst_buffer++;
        if (srcSize == 0) return;
      }
      else {
        skipCount--;
      }
      curSrcPos++;
    }
    else {
      u8 curVal = *curSrcPos;
      copyStart = dst_buffer - (curSrcPos[1] | (curVal & 0xF) << 8);
      curSrcPos += 2;
      if (curVal >> 4 == 0) {
        copyByteCount = *curSrcPos + 0x12;
        curSrcPos++;
      }
      else {
        copyByteCount = (curVal >> 4) + 2;
      }
      do {
        if (skipCount == 0) {
          *dst_buffer = *(copyStart - 1);
          srcSize--;
          dst_buffer++;
          if (srcSize == 0) return;
        }
        else {
          skipCount--;
        }
        copyByteCount--;
        copyStart++;
      } while (copyByteCount != 0);
    }
    chunkBits <<= 1;
    chunkBitsLeft--;
  } while (dst_buffer != decompEnd);
}

JKRDecomp::CompressionMode JKRDecomp::checkCompressed(u8* buf) {
  if (buf[0] == 'Y' && buf[1] == 'a' && buf[3] == '0') {
    if (buf[2] == 'y') {
      return SZP;
    }

    if (buf[2] == 'z') {
      return SZS;
    }
  }

  return NONE;
}

JKRDecompCommand::JKRDecompCommand() {
  OSInitMessageQueue(&this->mMesgQueue, this->mMesgBuffer, 1);
  this->mCallback = nullptr;
  this->pMesgQueue1C = nullptr;
  this->mCmd = this;
  this->transferType = MRAM;
}

JKRDecompCommand::~JKRDecompCommand() {}