/* Copyright 2017 - 2026 R. Thomas
 * Copyright 2017 - 2026 Quarkslab
 *
 * 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 "PE/checksum.hpp"
#include "LIEF/BinaryStream/SpanStream.hpp"

namespace LIEF {
ChecksumStream& ChecksumStream::write(const uint8_t* s, size_t n) {
  SpanStream stream(s, n);
  return write(stream);
}

ChecksumStream& ChecksumStream::write(BinaryStream& chk_stream) {
  if (has_leftover()) {
    uint16_t chunk = leftover();
    if (auto res = chk_stream.read<uint8_t>()) {
      chunk = (*res << 8) | chunk;
    }
    clear_leftover();
    partial_sum_ += chunk;
    partial_sum_ = (partial_sum_ >> 16) + (partial_sum_ & 0xffff);
  }

  const uint64_t file_length = chk_stream.size() - chk_stream.pos();
  uint64_t nb_chunk =
      (file_length + 1) / sizeof(uint16_t); // Number of uint16_t chunks

  while (chk_stream) {
    uint16_t chunk = 0;
    if (auto res = chk_stream.read<uint16_t>()) {
      chunk = *res;
    } else {
      break;
    }
    --nb_chunk;
    partial_sum_ += chunk;
    partial_sum_ = (partial_sum_ >> 16) + (partial_sum_ & 0xffff);
  }

  if (nb_chunk > 0) {
    if (auto res = chk_stream.read<uint8_t>()) {
      set_leftover(*res);
    }
  }

  size_ += chk_stream.size();
  return *this;
}

uint32_t ChecksumStream::finalize() {
  if (has_leftover()) {
    uint8_t chunk = leftover();
    clear_leftover();
    partial_sum_ += chunk;
    partial_sum_ = (partial_sum_ >> 16) + (partial_sum_ & 0xffff);
  }

  auto partial_sum_res =
      static_cast<uint16_t>(((partial_sum_ >> 16) + partial_sum_) & 0xffff);
  const uint32_t binary_checksum = checksum_;
  const uint32_t adjust_sum_lsb = binary_checksum & 0xFFFF;
  const uint32_t adjust_sum_msb = binary_checksum >> 16;

  partial_sum_res -= static_cast<int>(partial_sum_res < adjust_sum_lsb);
  partial_sum_res -= adjust_sum_lsb;

  partial_sum_res -= static_cast<int>(partial_sum_res < adjust_sum_msb);
  partial_sum_res -= adjust_sum_msb;

  return static_cast<uint32_t>(partial_sum_res) + size_;
}
}