/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements.  See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You 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 "MemoryOutputStream.h"

namespace rocketmq {
MemoryOutputStream::MemoryOutputStream(const size_t initialSize)
    : blockToUse(&internalBlock),
      externalData(NULL),
      position(0),
      size(0),
      availableSize(0) {
  internalBlock.setSize(initialSize, false);
}

MemoryOutputStream::MemoryOutputStream(MemoryBlock& memoryBlockToWriteTo,
                                       const bool appendToExistingBlockContent)
    : blockToUse(&memoryBlockToWriteTo),
      externalData(NULL),
      position(0),
      size(0),
      availableSize(0) {
  if (appendToExistingBlockContent)
    position = size = memoryBlockToWriteTo.getSize();
}

MemoryOutputStream::MemoryOutputStream(void* destBuffer, size_t destBufferSize)
    : blockToUse(NULL),
      externalData(destBuffer),
      position(0),
      size(0),
      availableSize(destBufferSize) {}

MemoryOutputStream::~MemoryOutputStream() { trimExternalBlockSize(); }

void MemoryOutputStream::flush() { trimExternalBlockSize(); }

void MemoryOutputStream::trimExternalBlockSize() {
  if (blockToUse != &internalBlock && blockToUse != NULL)
    blockToUse->setSize(size, false);
}

void MemoryOutputStream::preallocate(const size_t bytesToPreallocate) {
  if (blockToUse != NULL) blockToUse->ensureSize(bytesToPreallocate + 1);
}

void MemoryOutputStream::reset() {
  position = 0;
  size = 0;
}

char* MemoryOutputStream::prepareToWrite(size_t numBytes) {
  size_t storageNeeded = position + numBytes;

  char* data;

  if (blockToUse != NULL) {
    if (storageNeeded >= (unsigned int)(blockToUse->getSize()))
      blockToUse->ensureSize(
          (storageNeeded + std::min(storageNeeded / 2, (size_t)(1024 * 1024)) +
           32) &
          ~31u);

    data = static_cast<char*>(blockToUse->getData());
  } else {
    if (storageNeeded > availableSize) return NULL;

    data = static_cast<char*>(externalData);
  }

  char* const writePointer = data + position;
  position += numBytes;
  size = std::max(size, position);
  return writePointer;
}

bool MemoryOutputStream::write(const void* const buffer, size_t howMany) {
  if (howMany == 0) return true;

  if (char* dest = prepareToWrite(howMany)) {
    memcpy(dest, buffer, howMany);
    return true;
  }

  return false;
}

bool MemoryOutputStream::writeRepeatedByte(uint8 byte, size_t howMany) {
  if (howMany == 0) return true;

  if (char* dest = prepareToWrite(howMany)) {
    memset(dest, byte, howMany);
    return true;
  }

  return false;
}

MemoryBlock MemoryOutputStream::getMemoryBlock() const {
  return MemoryBlock(getData(), getDataSize());
}

const void* MemoryOutputStream::getData() const {
  if (blockToUse == NULL) return externalData;

  if ((unsigned int)blockToUse->getSize() > size)
    static_cast<char*>(blockToUse->getData())[size] = 0;

  return blockToUse->getData();
}

bool MemoryOutputStream::setPosition(int64 newPosition) {
  if (newPosition <= (int64)size) {
    // ok to seek backwards
    if (newPosition < 0)
      position = 0;
    else
      position = (int64)size < newPosition ? size : newPosition;
    return true;
  }

  // can't move beyond the end of the stream..
  return false;
}

int64 MemoryOutputStream::writeFromInputStream(InputStream& source,
                                               int64 maxNumBytesToWrite) {
  // before writing from an input, see if we can preallocate to make it more
  // efficient..
  int64 availableData = source.getTotalLength() - source.getPosition();

  if (availableData > 0) {
    if (maxNumBytesToWrite > availableData || maxNumBytesToWrite < 0)
      maxNumBytesToWrite = availableData;

    if (blockToUse != NULL)
      preallocate(blockToUse->getSize() + (size_t)maxNumBytesToWrite);
  }

  return OutputStream::writeFromInputStream(source, maxNumBytesToWrite);
}

OutputStream& operator<<(OutputStream& stream,
                         const MemoryOutputStream& streamToRead) {
  const size_t dataSize = streamToRead.getDataSize();

  if (dataSize > 0) stream.write(streamToRead.getData(), dataSize);

  return stream;
}
}