#include <frostbite2D/resource/npk_archive.h>
#include <algorithm>
#include <cctype>
#include <SDL.h>

namespace frostbite2D {

const uint8 NpkArchive::NPK_KEY[256] = {
    112, 117, 99, 104, 105, 107, 111, 110, 64, 110, 101, 111, 112, 108, 101, 32,
    100, 117, 110, 103, 101, 111, 110, 32, 97, 110, 100, 32, 102, 105, 103, 104,
    116, 101, 114, 32, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78,
    70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70,
    68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68,
    78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78,
    70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70,
    68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68,
    78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78,
    70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70,
    68, 78, 70, 68, 78, 70, 68, 78, 70, 68, 78, 70, 0
};

NpkArchive& NpkArchive::get() {
  static NpkArchive instance;
  return instance;
}

void NpkArchive::setImagePackDirectory(const std::string& dir) {
  imagePackDirectory_ = dir;
}

const std::string& NpkArchive::getImagePackDirectory() const {
  return imagePackDirectory_;
}

void NpkArchive::init() {
  if (initialized_) {
    close();
  }

  scanNpkFiles();
  initialized_ = true;
}

void NpkArchive::close() {
  imgIndex_.clear();
  imageCache_.clear();
  lruList_.clear();
  currentCacheSize_ = 0;
  initialized_ = false;
}

bool NpkArchive::isOpen() const {
  return initialized_;
}

std::string NpkArchive::normalizePath(const std::string& path) const {
  std::string result = path;
  std::transform(result.begin(), result.end(), result.begin(),
                 [](unsigned char c) { return std::tolower(c); });
  return result;
}

void NpkArchive::scanNpkFiles() {
  Asset& asset = Asset::get();
  std::string npkDir = asset.resolvePath(imagePackDirectory_);
  if (!asset.isDirectory(npkDir)) {
    SDL_Log("NPK directory not found: %s", npkDir.c_str());
    return;
  }

  std::vector<std::string> files = asset.listFilesWithExtension(npkDir, ".npk");
  SDL_Log("Scanning %d NPK files...", static_cast<int>(files.size()));
  for (const auto &file : files) {
    parseNpkFile(file);
  }
}

bool NpkArchive::parseNpkFile(const std::string& npkPath) {
  Asset& asset = Asset::get();
  BinaryReader reader(npkPath);

  if (!reader.isOpen()) {
    SDL_Log("Failed to open NPK file: %s", npkPath.c_str());
    return false;
  }

  std::string npkFileName = asset.getFileName(npkPath);
  std::string header = reader.readNullTerminatedString();

  if (header.find("NeoplePack_Bill") == std::string::npos) {
    return false;
  }

  int32 imageCount = reader.readInt32();

  for (int32 i = 0; i < imageCount; ++i) {
    int32 offset = reader.readInt32();
    int32 length = reader.readInt32();
    std::string imgPath = readNpkInfoString(reader);

    ImgRef img;
    img.path = normalizePath(imgPath);
    img.npkFile = npkFileName;
    img.offset = static_cast<uint32>(offset);
    img.size = static_cast<uint32>(length);
    img.loaded = false;

    imgIndex_[img.path] = img;
  }

  return true;
}

bool NpkArchive::hasImg(const std::string& path) const {
  return imgIndex_.find(normalizePath(path)) != imgIndex_.end();
}

std::optional<ImgRef> NpkArchive::getImg(const std::string& path) {
  auto it = imgIndex_.find(normalizePath(path));
  if (it == imgIndex_.end()) {
    return std::nullopt;
  }
  return it->second;
}

std::vector<std::string> NpkArchive::listImgs() const {
  std::vector<std::string> result;
  result.reserve(imgIndex_.size());
  for (const auto& pair : imgIndex_) {
    result.push_back(pair.first);
  }
  return result;
}

std::optional<ImageFrame> NpkArchive::getImageFrame(const ImgRef& img, size_t index) {
  std::string imgPath = normalizePath(img.path);

  auto cacheIt = imageCache_.find(imgPath);
  if (cacheIt == imageCache_.end()) {
    auto it = imgIndex_.find(imgPath);
    if (it == imgIndex_.end()) {
      return std::nullopt;
    }

    if (!loadImgData(it->second)) {
      return std::nullopt;
    }
    cacheIt = imageCache_.find(imgPath);
    if (cacheIt == imageCache_.end()) {
      return std::nullopt;
    }
  }

  updateCacheUsage(imgPath);

  if (index >= cacheIt->second.frames.size()) {
    return std::nullopt;
  }

  return cacheIt->second.frames[index];
}

size_t NpkArchive::getFrameCount(const ImgRef& img) const {
  std::string imgPath = normalizePath(img.path);
  auto it = imgIndex_.find(imgPath);
  if (it == imgIndex_.end()) {
    return 0;
  }
  return it->second.frameCount;
}

bool NpkArchive::loadImgData(ImgRef& img) {
  Asset& asset = Asset::get();
  std::string npkPath = asset.combinePath(imagePackDirectory_, img.npkFile);
  npkPath = asset.resolvePath(npkPath);

  BinaryReader reader(npkPath);
  if (!reader.isOpen()) {
    SDL_Log("Failed to open NPK for IMG: %s", npkPath.c_str());
    return false;
  }

  reader.seek(img.offset);
  std::string flag = reader.readNullTerminatedString();

  if (flag.find("Neople Img File") == std::string::npos) {
    return false;
  }

  int64 tableLength = reader.readInt64();
  reader.readInt32();
  int32 indexCount = reader.readInt32();

  img.frameCount = static_cast<size_t>(indexCount);

  CachedImageData cachedData;
  cachedData.frames.resize(indexCount);

  size_t dataStartPos = img.offset + static_cast<uint32>(tableLength) + 32;

  for (int32 i = 0; i < indexCount; ++i) {
    ImageFrame& frame = cachedData.frames[i];
    frame.type = reader.readInt32();

    if (frame.type == 17) {
      int32 refIndex = reader.readInt32();
      frame.compressionType = refIndex;
      frame.width = cachedData.frames[refIndex].width;
      frame.height = cachedData.frames[refIndex].height;
      frame.xPos = cachedData.frames[refIndex].xPos;
      frame.yPos = cachedData.frames[refIndex].yPos;
      frame.frameXPos = cachedData.frames[refIndex].frameXPos;
      frame.frameYPos = cachedData.frames[refIndex].frameYPos;
      frame.offset = cachedData.frames[refIndex].offset;
      frame.size = cachedData.frames[refIndex].size;
      continue;
    }

    frame.compressionType = reader.readInt32();
    frame.width = reader.readInt32();
    frame.height = reader.readInt32();
    int32 size = reader.readInt32();
    frame.xPos = reader.readInt32();
    frame.yPos = reader.readInt32();
    frame.frameXPos = reader.readInt32();
    frame.frameYPos = reader.readInt32();

    frame.size = size;
    if (i == 0) {
      frame.offset = static_cast<uint32>(dataStartPos);
    } else {
      frame.offset = cachedData.frames[i - 1].offset + cachedData.frames[i - 1].size;
    }
  }

  for (int32 i = 0; i < indexCount; ++i) {
    ImageFrame& frame = cachedData.frames[i];

    if (frame.type == 17) {
      int32 refIndex = frame.compressionType;
      frame.data = cachedData.frames[refIndex].data;
      continue;
    }

    reader.seek(frame.offset);
    std::vector<uint8> compressedData = reader.readBytes(static_cast<size_t>(frame.size));

    int32 deSize = frame.width * frame.height * 4;
    std::vector<uint8> decompressedData(deSize);
    unsigned long realSize = static_cast<unsigned long>(deSize);
    
    int uncompressResult = uncompress(
      decompressedData.data(), 
      &realSize, 
      compressedData.data(), 
      static_cast<unsigned long>(frame.size)
    );

    if (uncompressResult != Z_OK) {
      SDL_Log("Failed to uncompress image data: %d", uncompressResult);
      continue;
    }

    if (frame.type != 16) {
      int32 pngByteSize = deSize * 2;
      frame.data.resize(pngByteSize);

      for (int32 e = 0; e < pngByteSize; e += 4) {
        uint8 needData[2] = {0, 0};
        if ((e / 4) * 2 + 1 < static_cast<int32>(decompressedData.size())) {
          needData[0] = decompressedData[(e / 4) * 2];
          needData[1] = decompressedData[(e / 4) * 2 + 1];
        }
        parseColor(needData, frame.type, frame.data.data(), e);
      }
    } else {
      frame.data = std::move(decompressedData);
    }

    cachedData.memoryUsage += frame.data.size();
  }

  imageCache_[img.path] = std::move(cachedData);
  currentCacheSize_ += imageCache_[img.path].memoryUsage;
  lruList_.push_front(img.path);

  img.loaded = true;
  imgIndex_[img.path] = img;

  return true;
}

void NpkArchive::parseColor(const uint8* tab, int type, uint8* saveByte, int offset) {
  uint8 a = 0, r = 0, g = 0, b = 0;

  switch (type) {
    case 0x0e:
      a = static_cast<uint8>(tab[1] >> 7);
      r = static_cast<uint8>((tab[1] >> 2) & 0x1f);
      g = static_cast<uint8>((tab[0] >> 5) | ((tab[1] & 3) << 3));
      b = static_cast<uint8>(tab[0] & 0x1f);
      a = static_cast<uint8>(a * 0xff);
      r = static_cast<uint8>((r << 3) | (r >> 2));
      g = static_cast<uint8>((g << 3) | (g >> 2));
      b = static_cast<uint8>((b << 3) | (b >> 2));
      break;

    case 0x0f:
      a = static_cast<uint8>(tab[1] & 0xf0);
      r = static_cast<uint8>((tab[1] & 0xf) << 4);
      g = static_cast<uint8>(tab[0] & 0xf0);
      b = static_cast<uint8>((tab[0] & 0xf) << 4);
      break;
  }

  saveByte[offset + 0] = b;
  saveByte[offset + 1] = g;
  saveByte[offset + 2] = r;
  saveByte[offset + 3] = a;
}

void NpkArchive::setCacheSize(size_t maxBytes) {
  maxCacheSize_ = maxBytes;
  evictCacheIfNeeded(0);
}

void NpkArchive::clearCache() {
  imageCache_.clear();
  lruList_.clear();
  currentCacheSize_ = 0;
}

size_t NpkArchive::getCacheUsage() const {
  return currentCacheSize_;
}

void NpkArchive::evictCacheIfNeeded(size_t requiredSize) {
  while (currentCacheSize_ + requiredSize > maxCacheSize_ && !lruList_.empty()) {
    std::string oldest = lruList_.back();
    lruList_.pop_back();

    auto it = imageCache_.find(oldest);
    if (it != imageCache_.end()) {
      currentCacheSize_ -= it->second.memoryUsage;
      imageCache_.erase(it);
    }
  }
}

void NpkArchive::updateCacheUsage(const std::string& imgPath) {
  auto it = std::find(lruList_.begin(), lruList_.end(), imgPath);
  if (it != lruList_.end()) {
    lruList_.erase(it);
  }
  lruList_.push_front(imgPath);

  auto cacheIt = imageCache_.find(imgPath);
  if (cacheIt != imageCache_.end()) {
    cacheIt->second.lastUseTime = SDL_GetTicks();
  }
}

std::string NpkArchive::readNpkInfoString(BinaryReader& reader) {
  if (reader.eof()) {
    return "";
  }

  std::vector<uint8> encrypted = reader.readBytes(256);
  if (encrypted.size() < 256) {
    return "";
  }

  std::vector<char> decrypted(256);
  for (int i = 0; i < 256; ++i) {
    decrypted[i] = static_cast<char>(encrypted[i] ^ NPK_KEY[i]);
  }

  return std::string(decrypted.data());
}

void NpkArchive::setDefaultImg(const std::string& imgPath, size_t frameIndex) {
  defaultImgPath_ = normalizePath(imgPath);
  defaultImgFrame_ = frameIndex;
}

const std::string& NpkArchive::getDefaultImgPath() const {
  return defaultImgPath_;
}

size_t NpkArchive::getDefaultImgFrame() const {
  return defaultImgFrame_;
}

}
