/*
 * Copyright (C) 2021 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

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

#include <algorithm>
#include <atomic>
#include <memory>
#include <mutex>
#include <string>
#include <utility>

#include <7zCrc.h>
#include <Xz.h>
#include <XzCrc64.h>

#include <unwindstack/Log.h>

#include "MemoryXz.h"

namespace unwindstack {

// Statistics (used only for optional debug log messages).
static constexpr bool kLogMemoryXzUsage = false;
std::atomic_size_t MemoryXz::total_used_ = 0;
std::atomic_size_t MemoryXz::total_size_ = 0;
std::atomic_size_t MemoryXz::total_open_ = 0;

MemoryXz::MemoryXz(Memory* memory, uint64_t addr, uint64_t size, const std::string& name)
    : compressed_memory_(memory), compressed_addr_(addr), compressed_size_(size), name_(name) {
  total_open_ += 1;
}

bool MemoryXz::Init() {
  static std::once_flag crc_initialized;
  std::call_once(crc_initialized, []() {
    CrcGenerateTable();
    Crc64GenerateTable();
  });
  if (compressed_size_ >= kMaxCompressedSize) {
    return false;
  }
  if (!ReadBlocks()) {
    return false;
  }

  // All blocks (except the last one) must have the same power-of-2 size.
  if (blocks_.size() > 1) {
    size_t block_size_log2 = __builtin_ctz(blocks_.front().decompressed_size);
    auto correct_size = [=](XzBlock& b) { return b.decompressed_size == (1 << block_size_log2); };
    if (std::all_of(blocks_.begin(), std::prev(blocks_.end()), correct_size) &&
        blocks_.back().decompressed_size <= (1 << block_size_log2)) {
      block_size_log2_ = block_size_log2;
    } else {
      // Inconsistent block-sizes.  Decompress and merge everything now.
      std::unique_ptr<uint8_t[]> data(new uint8_t[size_]);
      size_t offset = 0;
      for (XzBlock& block : blocks_) {
        if (!Decompress(&block)) {
          return false;
        }
        memcpy(data.get() + offset, block.decompressed_data.get(), block.decompressed_size);
        offset += block.decompressed_size;
      }
      blocks_.clear();
      blocks_.push_back(XzBlock{
          .decompressed_data = std::move(data),
          .decompressed_size = size_,
      });
      block_size_log2_ = 31;  // Because 32 bits is too big (shift right by 32 is not allowed).
    }
  }

  return true;
}

MemoryXz::~MemoryXz() {
  total_used_ -= used_;
  total_size_ -= size_;
  total_open_ -= 1;
}

size_t MemoryXz::Read(uint64_t addr, void* buffer, size_t size) {
  if (addr >= size_) {
    return 0;  // Read past the end.
  }
  uint8_t* dst = reinterpret_cast<uint8_t*>(buffer);  // Position in the output buffer.
  for (size_t i = addr >> block_size_log2_; i < blocks_.size(); i++) {
    XzBlock* block = &blocks_[i];
    if (block->decompressed_data == nullptr) {
      if (!Decompress(block)) {
        break;
      }
    }
    size_t offset = (addr - (i << block_size_log2_));  // Start inside the block.
    size_t copy_bytes = std::min<size_t>(size, block->decompressed_size - offset);
    memcpy(dst, block->decompressed_data.get() + offset, copy_bytes);
    dst += copy_bytes;
    addr += copy_bytes;
    size -= copy_bytes;
    if (size == 0) {
      break;
    }
  }
  return dst - reinterpret_cast<uint8_t*>(buffer);
}

bool MemoryXz::ReadBlocks() {
  static ISzAlloc alloc;
  alloc.Alloc = [](ISzAllocPtr, size_t size) { return malloc(size); };
  alloc.Free = [](ISzAllocPtr, void* ptr) { return free(ptr); };

  // Read the compressed data, so we can quickly scan through the headers.
  std::unique_ptr<uint8_t[]> compressed_data(new (std::nothrow) uint8_t[compressed_size_]);
  if (compressed_data.get() == nullptr) {
    return false;
  }
  if (!compressed_memory_->ReadFully(compressed_addr_, compressed_data.get(), compressed_size_)) {
    return false;
  }

  // Implement the required interface for communication
  // (written in C so we can not use virtual methods or member functions).
  struct XzLookInStream : public ILookInStream, public ICompressProgress {
    static SRes LookImpl(const ILookInStream* p, const void** buf, size_t* size) {
      auto* ctx = reinterpret_cast<const XzLookInStream*>(p);
      *buf = ctx->data + ctx->offset;
      *size = std::min(*size, ctx->size - ctx->offset);
      return SZ_OK;
    }
    static SRes SkipImpl(const ILookInStream* p, size_t len) {
      auto* ctx = reinterpret_cast<XzLookInStream*>(const_cast<ILookInStream*>(p));
      ctx->offset += len;
      return SZ_OK;
    }
    static SRes ReadImpl(const ILookInStream* p, void* buf, size_t* size) {
      auto* ctx = reinterpret_cast<const XzLookInStream*>(p);
      *size = std::min(*size, ctx->size - ctx->offset);
      memcpy(buf, ctx->data + ctx->offset, *size);
      return SZ_OK;
    }
    static SRes SeekImpl(const ILookInStream* p, Int64* pos, ESzSeek origin) {
      auto* ctx = reinterpret_cast<XzLookInStream*>(const_cast<ILookInStream*>(p));
      switch (origin) {
        case SZ_SEEK_SET:
          ctx->offset = *pos;
          break;
        case SZ_SEEK_CUR:
          ctx->offset += *pos;
          break;
        case SZ_SEEK_END:
          ctx->offset = ctx->size + *pos;
          break;
      }
      *pos = ctx->offset;
      return SZ_OK;
    }
    static SRes ProgressImpl(const ICompressProgress*, UInt64, UInt64) { return SZ_OK; }
    size_t offset;
    uint8_t* data;
    size_t size;
  };
  XzLookInStream callbacks;
  callbacks.Look = &XzLookInStream::LookImpl;
  callbacks.Skip = &XzLookInStream::SkipImpl;
  callbacks.Read = &XzLookInStream::ReadImpl;
  callbacks.Seek = &XzLookInStream::SeekImpl;
  callbacks.Progress = &XzLookInStream::ProgressImpl;
  callbacks.offset = 0;
  callbacks.data = compressed_data.get();
  callbacks.size = compressed_size_;

  // Iterate over the internal XZ blocks without decompressing them.
  CXzs xzs;
  Xzs_Construct(&xzs);
  Int64 end_offset = compressed_size_;
  if (Xzs_ReadBackward(&xzs, &callbacks, &end_offset, &callbacks, &alloc) == SZ_OK) {
    blocks_.reserve(Xzs_GetNumBlocks(&xzs));
    size_t dst_offset = 0;
    for (int s = xzs.num - 1; s >= 0; s--) {
      const CXzStream& stream = xzs.streams[s];
      size_t src_offset = stream.startOffset + XZ_STREAM_HEADER_SIZE;
      for (size_t b = 0; b < stream.numBlocks; b++) {
        const CXzBlockSizes& block = stream.blocks[b];
        blocks_.push_back(XzBlock{
            .decompressed_data = nullptr,  // Lazy allocation and decompression.
            .decompressed_size = static_cast<uint32_t>(block.unpackSize),
            .compressed_offset = static_cast<uint32_t>(src_offset),
            .compressed_size = static_cast<uint32_t>((block.totalSize + 3) & ~3u),
            .stream_flags = stream.flags,
        });
        dst_offset += blocks_.back().decompressed_size;
        src_offset += blocks_.back().compressed_size;
      }
    }
    size_ = dst_offset;
    total_size_ += dst_offset;
  }
  Xzs_Free(&xzs, &alloc);
  return !blocks_.empty();
}

bool MemoryXz::Decompress(XzBlock* block) {
  static ISzAlloc alloc;
  alloc.Alloc = [](ISzAllocPtr, size_t size) { return malloc(size); };
  alloc.Free = [](ISzAllocPtr, void* ptr) { return free(ptr); };

  // Read the compressed data for this block.
  std::unique_ptr<uint8_t[]> compressed_data(new (std::nothrow) uint8_t[block->compressed_size]);
  if (compressed_data.get() == nullptr) {
    return false;
  }
  if (!compressed_memory_->ReadFully(compressed_addr_ + block->compressed_offset,
                                     compressed_data.get(), block->compressed_size)) {
    return false;
  }

  // Allocate decompressed memory.
  std::unique_ptr<uint8_t[]> decompressed_data(new uint8_t[block->decompressed_size]);
  if (decompressed_data == nullptr) {
    return false;
  }

  // Decompress.
  CXzUnpacker state{};
  XzUnpacker_Construct(&state, &alloc);
  state.streamFlags = block->stream_flags;
  XzUnpacker_PrepareToRandomBlockDecoding(&state);
  size_t decompressed_size = block->decompressed_size;
  size_t compressed_size = block->compressed_size;
  ECoderStatus status;
  XzUnpacker_SetOutBuf(&state, decompressed_data.get(), decompressed_size);
  int return_val =
      XzUnpacker_Code(&state, /*decompressed_data=*/nullptr, &decompressed_size,
                      compressed_data.get(), &compressed_size, true, CODER_FINISH_END, &status);
  XzUnpacker_Free(&state);
  if (return_val != SZ_OK || status != CODER_STATUS_FINISHED_WITH_MARK) {
    Log::Error("Cannot decompress \"%s\"", name_.c_str());
    return false;
  }

  used_ += block->decompressed_size;
  total_used_ += block->decompressed_size;
  if (kLogMemoryXzUsage) {
    Log::Info("decompressed memory: %zi%% of %ziKB (%zi files), %i%% of %iKB (%s)",
              100 * total_used_ / total_size_, total_size_ / 1024, total_open_.load(),
              100 * used_ / size_, size_ / 1024, name_.c_str());
  }

  block->decompressed_data = std::move(decompressed_data);
  return true;
}

}  // namespace unwindstack