// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "puffin/src/puffin_stream.h"

#include <algorithm>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "puffin/src/bit_reader.h"
#include "puffin/src/bit_writer.h"
#include "puffin/src/include/puffin/common.h"
#include "puffin/src/include/puffin/huffer.h"
#include "puffin/src/include/puffin/puffer.h"
#include "puffin/src/include/puffin/stream.h"
#include "puffin/src/logging.h"
#include "puffin/src/puff_reader.h"
#include "puffin/src/puff_writer.h"

using std::shared_ptr;
using std::unique_ptr;
using std::vector;

namespace puffin {

namespace {

bool CheckArgsIntegrity(uint64_t puff_size,
                        const vector<BitExtent>& deflates,
                        const vector<ByteExtent>& puffs) {
  TEST_AND_RETURN_FALSE(puffs.size() == deflates.size());
  // Check if the |puff_size| is actually greater than the last byte of the last
  // puff in |puffs|.
  if (!puffs.empty()) {
    TEST_AND_RETURN_FALSE(puff_size >=
                          puffs.back().offset + puffs.back().length);
  }

  // Check to make sure |puffs| and |deflates| are sorted and non-overlapping.
  auto is_overlap = [](const auto& a, const auto& b) {
    return (a.offset + a.length) > b.offset;
  };
  TEST_AND_RETURN_FALSE(deflates.end() == std::adjacent_find(deflates.begin(),
                                                             deflates.end(),
                                                             is_overlap));
  TEST_AND_RETURN_FALSE(puffs.end() == std::adjacent_find(puffs.begin(),
                                                          puffs.end(),
                                                          is_overlap));
  return true;
}

}  // namespace

UniqueStreamPtr PuffinStream::CreateForPuff(UniqueStreamPtr stream,
                                            shared_ptr<Puffer> puffer,
                                            uint64_t puff_size,
                                            const vector<BitExtent>& deflates,
                                            const vector<ByteExtent>& puffs,
                                            size_t max_cache_size) {
  TEST_AND_RETURN_VALUE(CheckArgsIntegrity(puff_size, deflates, puffs),
                        nullptr);
  TEST_AND_RETURN_VALUE(stream->Seek(0), nullptr);

  UniqueStreamPtr puffin_stream(new PuffinStream(std::move(stream), puffer,
                                                 nullptr, puff_size, deflates,
                                                 puffs, max_cache_size));
  TEST_AND_RETURN_VALUE(puffin_stream->Seek(0), nullptr);
  return puffin_stream;
}

UniqueStreamPtr PuffinStream::CreateForHuff(UniqueStreamPtr stream,
                                            shared_ptr<Huffer> huffer,
                                            uint64_t puff_size,
                                            const vector<BitExtent>& deflates,
                                            const vector<ByteExtent>& puffs) {
  TEST_AND_RETURN_VALUE(CheckArgsIntegrity(puff_size, deflates, puffs),
                        nullptr);
  TEST_AND_RETURN_VALUE(stream->Seek(0), nullptr);

  UniqueStreamPtr puffin_stream(new PuffinStream(
      std::move(stream), nullptr, huffer, puff_size, deflates, puffs, 0));
  TEST_AND_RETURN_VALUE(puffin_stream->Seek(0), nullptr);
  return puffin_stream;
}

PuffinStream::PuffinStream(UniqueStreamPtr stream,
                           shared_ptr<Puffer> puffer,
                           shared_ptr<Huffer> huffer,
                           uint64_t puff_size,
                           const vector<BitExtent>& deflates,
                           const vector<ByteExtent>& puffs,
                           size_t max_cache_size)
    : stream_(std::move(stream)),
      puffer_(puffer),
      huffer_(huffer),
      puff_stream_size_(puff_size),
      deflates_(deflates),
      puffs_(puffs),
      puff_pos_(0),
      skip_bytes_(0),
      deflate_bit_pos_(0),
      last_byte_(0),
      extra_byte_(0),
      is_for_puff_(puffer_ ? true : false),
      closed_(false),
      max_cache_size_(max_cache_size),
      cur_cache_size_(0) {
  // Building upper bounds for faster seek.
  upper_bounds_.reserve(puffs.size());
  for (const auto& puff : puffs) {
    upper_bounds_.emplace_back(puff.offset + puff.length);
  }
  upper_bounds_.emplace_back(puff_stream_size_ + 1);

  // We can pass the size of the deflate stream too, but it is not necessary
  // yet. We cannot get the size of stream from itself, because we might be
  // writing into it and its size is not defined yet.
  uint64_t deflate_stream_size = puff_stream_size_;
  if (!puffs.empty()) {
    deflate_stream_size =
        ((deflates.back().offset + deflates.back().length) / 8) +
        puff_stream_size_ - (puffs.back().offset + puffs.back().length);
  }

  deflates_.emplace_back(deflate_stream_size * 8, 0);
  puffs_.emplace_back(puff_stream_size_, 0);

  // Look for the largest puff and deflate extents and get proper size buffers.
  uint64_t max_puff_length = 0;
  for (const auto& puff : puffs) {
    max_puff_length = std::max(max_puff_length, puff.length);
  }
  puff_buffer_.reset(new Buffer(max_puff_length + 1));
  if (max_cache_size_ < max_puff_length) {
    max_cache_size_ = 0;  // It means we are not caching puffs.
  }

  uint64_t max_deflate_length = 0;
  for (const auto& deflate : deflates) {
    max_deflate_length = std::max(max_deflate_length, deflate.length * 8);
  }
  deflate_buffer_.reset(new Buffer(max_deflate_length + 2));
}

bool PuffinStream::GetSize(uint64_t* size) const {
  *size = puff_stream_size_;
  return true;
}

bool PuffinStream::GetOffset(uint64_t* offset) const {
  *offset = puff_pos_ + skip_bytes_;
  return true;
}

bool PuffinStream::Seek(uint64_t offset) {
  TEST_AND_RETURN_FALSE(!closed_);
  if (!is_for_puff_) {
    // For huffing we should not seek, only seek to zero is accepted.
    TEST_AND_RETURN_FALSE(offset == 0);
  }

  TEST_AND_RETURN_FALSE(offset <= puff_stream_size_);

  // We are searching first available puff which either includes the |offset| or
  // it is the next available puff after |offset|.
  auto next_puff_iter =
      std::upper_bound(upper_bounds_.begin(), upper_bounds_.end(), offset);
  TEST_AND_RETURN_FALSE(next_puff_iter != upper_bounds_.end());
  auto next_puff_idx = std::distance(upper_bounds_.begin(), next_puff_iter);
  cur_puff_ = std::next(puffs_.begin(), next_puff_idx);
  cur_deflate_ = std::next(deflates_.begin(), next_puff_idx);

  if (offset < cur_puff_->offset) {
    // between two puffs.
    puff_pos_ = offset;
    auto back_track_bytes = cur_puff_->offset - puff_pos_;
    deflate_bit_pos_ = ((cur_deflate_->offset + 7) / 8 - back_track_bytes) * 8;
    if (cur_puff_ != puffs_.begin()) {
      auto prev_deflate = std::prev(cur_deflate_);
      if (deflate_bit_pos_ < prev_deflate->offset + prev_deflate->length) {
        deflate_bit_pos_ = prev_deflate->offset + prev_deflate->length;
      }
    }
  } else {
    // Inside a puff.
    puff_pos_ = cur_puff_->offset;
    deflate_bit_pos_ = cur_deflate_->offset;
  }
  skip_bytes_ = offset - puff_pos_;
  if (!is_for_puff_ && offset == 0) {
    TEST_AND_RETURN_FALSE(stream_->Seek(0));
    TEST_AND_RETURN_FALSE(SetExtraByte());
  }
  return true;
}

bool PuffinStream::Close() {
  closed_ = true;
  return stream_->Close();
}

bool PuffinStream::Read(void* buffer, size_t count) {
  TEST_AND_RETURN_FALSE(!closed_);
  TEST_AND_RETURN_FALSE(is_for_puff_);
  if (cur_puff_ == puffs_.end()) {
    TEST_AND_RETURN_FALSE(count == 0);
  }
  auto bytes = static_cast<uint8_t*>(buffer);
  uint64_t length = count;
  uint64_t bytes_read = 0;
  while (bytes_read < length) {
    if (puff_pos_ < cur_puff_->offset) {
      // Reading between two deflates. We also read bytes that have at least one
      // bit of a deflate bit stream. The byte which has both deflate and raw
      // data will be shifted or masked off the deflate bits and the remaining
      // value will be saved in the puff stream as an byte integer.
      uint64_t start_byte = (deflate_bit_pos_ / 8);
      uint64_t end_byte = (cur_deflate_->offset + 7) / 8;
      auto bytes_to_read = std::min(length - bytes_read, end_byte - start_byte);
      TEST_AND_RETURN_FALSE(bytes_to_read >= 1);

      TEST_AND_RETURN_FALSE(stream_->Seek(start_byte));
      TEST_AND_RETURN_FALSE(stream_->Read(bytes + bytes_read, bytes_to_read));

      // If true, we read the first byte of the curret deflate. So we have to
      // mask out the deflate bits (which are most significant bits.)
      if ((start_byte + bytes_to_read) * 8 > cur_deflate_->offset) {
        bytes[bytes_read + bytes_to_read - 1] &=
            (1 << (cur_deflate_->offset & 7)) - 1;
      }

      // If true, we read the last byte of the previous deflate and we have to
      // shift it out. The least significat bits belongs to the deflate
      // stream. The order of these last two conditions are important because a
      // byte can contain a finishing deflate and a starting deflate with some
      // bits between them so we have to modify correctly. Keep in mind that in
      // this situation both are modifying the same byte.
      if (start_byte * 8 < deflate_bit_pos_) {
        bytes[bytes_read] >>= deflate_bit_pos_ & 7;
      }

      // Pass |deflate_bit_pos_| for all the read bytes.
      deflate_bit_pos_ -= deflate_bit_pos_ & 7;
      deflate_bit_pos_ += bytes_to_read * 8;
      if (deflate_bit_pos_ > cur_deflate_->offset) {
        // In case it reads into the first byte of the current deflate.
        deflate_bit_pos_ = cur_deflate_->offset;
      }

      bytes_read += bytes_to_read;
      puff_pos_ += bytes_to_read;
      TEST_AND_RETURN_FALSE(puff_pos_ <= cur_puff_->offset);
    } else {
      // Reading the deflate itself. We read all bytes including the first and
      // last byte (which may partially include a deflate bit). Here we keep the
      // |puff_pos_| point to the first byte of the puffed stream and
      // |skip_bytes_| shows how many bytes in the puff we have copied till now.
      auto start_byte = (cur_deflate_->offset / 8);
      auto end_byte = (cur_deflate_->offset + cur_deflate_->length + 7) / 8;
      auto bytes_to_read = end_byte - start_byte;
      // Puff directly to buffer if it has space.
      bool puff_directly_into_buffer =
          max_cache_size_ == 0 && (skip_bytes_ == 0) &&
          (length - bytes_read >= cur_puff_->length);

      auto cur_puff_idx = std::distance(puffs_.begin(), cur_puff_);
      if (max_cache_size_ == 0 ||
          !GetPuffCache(cur_puff_idx, cur_puff_->length, &puff_buffer_)) {
        // Did not find the puff buffer in cache. We have to build it.
        deflate_buffer_->resize(bytes_to_read);
        TEST_AND_RETURN_FALSE(stream_->Seek(start_byte));
        TEST_AND_RETURN_FALSE(
            stream_->Read(deflate_buffer_->data(), bytes_to_read));
        BufferBitReader bit_reader(deflate_buffer_->data(), bytes_to_read);

        BufferPuffWriter puff_writer(puff_directly_into_buffer
                                         ? bytes + bytes_read
                                         : puff_buffer_->data(),
                                     cur_puff_->length);

        // Drop the first unused bits.
        size_t extra_bits_len = cur_deflate_->offset & 7;
        TEST_AND_RETURN_FALSE(bit_reader.CacheBits(extra_bits_len));
        bit_reader.DropBits(extra_bits_len);

        TEST_AND_RETURN_FALSE(
            puffer_->PuffDeflate(&bit_reader, &puff_writer, nullptr));
        TEST_AND_RETURN_FALSE(bytes_to_read == bit_reader.Offset());
        TEST_AND_RETURN_FALSE(cur_puff_->length == puff_writer.Size());
      } else {
        // Just seek to proper location.
        TEST_AND_RETURN_FALSE(stream_->Seek(start_byte + bytes_to_read));
      }
      // Copy from puff buffer to output if needed.
      auto bytes_to_copy =
          std::min(length - bytes_read, cur_puff_->length - skip_bytes_);
      if (!puff_directly_into_buffer) {
        memcpy(bytes + bytes_read, puff_buffer_->data() + skip_bytes_,
               bytes_to_copy);
      }

      skip_bytes_ += bytes_to_copy;
      bytes_read += bytes_to_copy;

      // Move to next puff.
      if (puff_pos_ + skip_bytes_ == cur_puff_->offset + cur_puff_->length) {
        puff_pos_ += skip_bytes_;
        skip_bytes_ = 0;
        deflate_bit_pos_ = cur_deflate_->offset + cur_deflate_->length;
        cur_puff_++;
        cur_deflate_++;
        if (cur_puff_ == puffs_.end()) {
          break;
        }
      }
    }
  }

  TEST_AND_RETURN_FALSE(bytes_read == length);
  return true;
}

bool PuffinStream::Write(const void* buffer, size_t count) {
  TEST_AND_RETURN_FALSE(!closed_);
  TEST_AND_RETURN_FALSE(!is_for_puff_);
  auto bytes = static_cast<const uint8_t*>(buffer);
  uint64_t length = count;
  uint64_t bytes_wrote = 0;
  while (bytes_wrote < length) {
    if (deflate_bit_pos_ < (cur_deflate_->offset & ~7ull)) {
      // Between two puffs or before the first puff. We know that we are
      // starting from the byte boundary because we have already processed the
      // non-deflate bits of the last byte of the last deflate. Here we don't
      // process any byte that has deflate bit.
      TEST_AND_RETURN_FALSE((deflate_bit_pos_ & 7) == 0);
      auto copy_len =
          std::min((cur_deflate_->offset / 8) - (deflate_bit_pos_ / 8),
                   length - bytes_wrote);
      TEST_AND_RETURN_FALSE(stream_->Write(bytes + bytes_wrote, copy_len));
      bytes_wrote += copy_len;
      puff_pos_ += copy_len;
      deflate_bit_pos_ += copy_len * 8;
    } else {
      // We are in a puff. We have to buffer incoming bytes until we reach the
      // size of the current puff so we can huff :). If the last bit of the
      // current deflate does not end in a byte boundary, then we have to read
      // one more byte to fill up the last byte of the deflate stream before
      // doing anything else.

      // |deflate_bit_pos_| now should be in the same byte as
      // |cur_deflate->offset|.
      if (deflate_bit_pos_ < cur_deflate_->offset) {
        last_byte_ |= bytes[bytes_wrote++] << (deflate_bit_pos_ & 7);
        skip_bytes_ = 0;
        deflate_bit_pos_ = cur_deflate_->offset;
        puff_pos_++;
        TEST_AND_RETURN_FALSE(puff_pos_ == cur_puff_->offset);
      }

      auto copy_len = std::min(length - bytes_wrote,
                               cur_puff_->length + extra_byte_ - skip_bytes_);
      TEST_AND_RETURN_FALSE(puff_buffer_->size() >= skip_bytes_ + copy_len);
      memcpy(puff_buffer_->data() + skip_bytes_, bytes + bytes_wrote, copy_len);
      skip_bytes_ += copy_len;
      bytes_wrote += copy_len;

      if (skip_bytes_ == cur_puff_->length + extra_byte_) {
        // |puff_buffer_| is full, now huff into the |deflate_buffer_|.
        auto start_byte = cur_deflate_->offset / 8;
        auto end_byte = (cur_deflate_->offset + cur_deflate_->length + 7) / 8;
        auto bytes_to_write = end_byte - start_byte;

        deflate_buffer_->resize(bytes_to_write);
        BufferBitWriter bit_writer(deflate_buffer_->data(), bytes_to_write);
        BufferPuffReader puff_reader(puff_buffer_->data(), cur_puff_->length);

        // Write last byte if it has any.
        TEST_AND_RETURN_FALSE(
            bit_writer.WriteBits(cur_deflate_->offset & 7, last_byte_));
        last_byte_ = 0;

        TEST_AND_RETURN_FALSE(huffer_->HuffDeflate(&puff_reader, &bit_writer));
        TEST_AND_RETURN_FALSE(bit_writer.Size() == bytes_to_write);
        TEST_AND_RETURN_FALSE(puff_reader.BytesLeft() == 0);

        deflate_bit_pos_ = cur_deflate_->offset + cur_deflate_->length;
        if (extra_byte_ == 1) {
          deflate_buffer_->data()[bytes_to_write - 1] |=
              puff_buffer_->data()[cur_puff_->length] << (deflate_bit_pos_ & 7);
          deflate_bit_pos_ = (deflate_bit_pos_ + 7) & ~7ull;
        } else if ((deflate_bit_pos_ & 7) != 0) {
          // This happens if current and next deflate finish and end on the same
          // byte, then we cannot write into output until we have huffed the
          // next puff buffer, so untill then we cache it into |last_byte_| and
          // we won't write it out.
          last_byte_ = deflate_buffer_->data()[bytes_to_write - 1];
          bytes_to_write--;
        }

        // Write |deflate_buffer_| into output.
        TEST_AND_RETURN_FALSE(
            stream_->Write(deflate_buffer_->data(), bytes_to_write));

        // Move to the next deflate/puff.
        puff_pos_ += skip_bytes_;
        skip_bytes_ = 0;
        cur_puff_++;
        cur_deflate_++;
        if (cur_puff_ == puffs_.end()) {
          break;
        }
        // Find if need an extra byte to cache at the end.
        TEST_AND_RETURN_FALSE(SetExtraByte());
      }
    }
  }

  TEST_AND_RETURN_FALSE(bytes_wrote == length);
  return true;
}

bool PuffinStream::SetExtraByte() {
  TEST_AND_RETURN_FALSE(cur_deflate_ != deflates_.end());
  if ((cur_deflate_ + 1) == deflates_.end()) {
    extra_byte_ = 0;
    return true;
  }
  uint64_t end_bit = cur_deflate_->offset + cur_deflate_->length;
  if ((end_bit & 7) && ((end_bit + 7) & ~7ull) <= (cur_deflate_ + 1)->offset) {
    extra_byte_ = 1;
  } else {
    extra_byte_ = 0;
  }
  return true;
}

bool PuffinStream::GetPuffCache(int puff_id,
                                uint64_t puff_size,
                                shared_ptr<Buffer>* buffer) {
  bool found = false;
  // Search for it.
  std::pair<int, shared_ptr<Buffer>> cache;
  // TODO(*): Find a faster way of doing this? Maybe change the data structure
  // that supports faster search.
  for (auto iter = caches_.begin(); iter != caches_.end(); ++iter) {
    if (iter->first == puff_id) {
      cache = std::move(*iter);
      found = true;
      // Remove it so later we can add it to the begining of the list.
      caches_.erase(iter);
      break;
    }
  }

  // If not found, either create one or get one from the list.
  if (!found) {
    // If |caches_| were full, remove last ones in the list (least used), until
    // we have enough space for the new cache.
    while (!caches_.empty() && cur_cache_size_ + puff_size > max_cache_size_) {
      cache = std::move(caches_.back());
      caches_.pop_back();  // Remove it from the list.
      cur_cache_size_ -= cache.second->capacity();
    }
    // If we have not populated the cache yet, create one.
    if (!cache.second) {
      cache.second.reset(new Buffer(puff_size));
    }
    cache.second->resize(puff_size);

    constexpr uint64_t kMaxSizeDifference = 20 * 1024;
    if (puff_size + kMaxSizeDifference < cache.second->capacity()) {
      cache.second->shrink_to_fit();
    }

    cur_cache_size_ += cache.second->capacity();
    cache.first = puff_id;
  }

  *buffer = cache.second;
  // By now we have either removed a cache or created new one. Now we have to
  // insert it in the front of the list so it becomes the most recently used
  // one.
  caches_.push_front(std::move(cache));
  return found;
}

}  // namespace puffin