// Copyright 2021 The Tint Authors.
//
// 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 "src/tint/fuzzers/tint_spirv_tools_fuzzer/spirv_fuzz_mutator.h"

#include <fstream>
#include <utility>

#include "source/opt/build_module.h"
#include "src/tint/fuzzers/tint_spirv_tools_fuzzer/util.h"

namespace tint::fuzzers::spvtools_fuzzer {

SpirvFuzzMutator::SpirvFuzzMutator(
    spv_target_env target_env,
    std::vector<uint32_t> binary,
    unsigned seed,
    const std::vector<spvtools::fuzz::fuzzerutil::ModuleSupplier>& donors,
    bool enable_all_passes,
    spvtools::fuzz::RepeatedPassStrategy repeated_pass_strategy,
    bool validate_after_each_pass,
    uint32_t transformation_batch_size)
    : transformation_batch_size_(transformation_batch_size),
      errors_(std::make_unique<std::stringstream>()),
      fuzzer_(nullptr),
      validator_options_(),
      original_binary_(std::move(binary)),
      seed_(seed) {
  auto ir_context = spvtools::BuildModule(
      target_env, spvtools::fuzz::fuzzerutil::kSilentMessageConsumer,
      original_binary_.data(), original_binary_.size());
  assert(ir_context && "|binary| is invalid");

  auto transformation_context =
      std::make_unique<spvtools::fuzz::TransformationContext>(
          std::make_unique<spvtools::fuzz::FactManager>(ir_context.get()),
          validator_options_);

  auto fuzzer_context = std::make_unique<spvtools::fuzz::FuzzerContext>(
      std::make_unique<spvtools::fuzz::PseudoRandomGenerator>(seed),
      spvtools::fuzz::FuzzerContext::GetMinFreshId(ir_context.get()), false);
  fuzzer_ = std::make_unique<spvtools::fuzz::Fuzzer>(
      std::move(ir_context), std::move(transformation_context),
      std::move(fuzzer_context), util::GetBufferMessageConsumer(errors_.get()),
      donors, enable_all_passes, repeated_pass_strategy,
      validate_after_each_pass, validator_options_);
}

Mutator::Result SpirvFuzzMutator::Mutate() {
  // The assertion will fail in |fuzzer_->Run| if the previous fuzzing led to
  // invalid module.
  auto result = fuzzer_->Run(transformation_batch_size_);
  switch (result.status) {
    case spvtools::fuzz::Fuzzer::Status::kComplete:
      return {Mutator::Status::kComplete, result.is_changed};
    case spvtools::fuzz::Fuzzer::Status::kModuleTooBig:
    case spvtools::fuzz::Fuzzer::Status::kTransformationLimitReached:
      return {Mutator::Status::kLimitReached, result.is_changed};
    case spvtools::fuzz::Fuzzer::Status::kFuzzerStuck:
      return {Mutator::Status::kStuck, result.is_changed};
    case spvtools::fuzz::Fuzzer::Status::kFuzzerPassLedToInvalidModule:
      return {Mutator::Status::kInvalid, result.is_changed};
  }
}

std::vector<uint32_t> SpirvFuzzMutator::GetBinary() const {
  std::vector<uint32_t> result;
  fuzzer_->GetIRContext()->module()->ToBinary(&result, true);
  return result;
}

std::string SpirvFuzzMutator::GetErrors() const {
  return errors_->str();
}

void SpirvFuzzMutator::LogErrors(const std::string* path,
                                 uint32_t count) const {
  auto message = GetErrors();
  std::cout << count << " | SpirvFuzzMutator (seed: " << seed_ << ")"
            << std::endl;
  std::cout << message << std::endl;

  if (path) {
    auto prefix = *path + std::to_string(count);

    // Write errors to file.
    std::ofstream(prefix + ".fuzzer.log") << "seed: " << seed_ << std::endl
                                          << message << std::endl;

    // Write the invalid SPIR-V binary.
    util::WriteBinary(prefix + ".fuzzer.invalid.spv", GetBinary());

    // Write the original SPIR-V binary.
    util::WriteBinary(prefix + ".fuzzer.original.spv", original_binary_);

    // Write transformations.
    google::protobuf::util::JsonOptions options;
    options.add_whitespace = true;
    std::string json;
    google::protobuf::util::MessageToJsonString(
        fuzzer_->GetTransformationSequence(), &json, options);
    std::ofstream(prefix + ".fuzzer.transformations.json") << json << std::endl;

    std::ofstream binary_transformations(
        prefix + ".fuzzer.transformations.binary",
        std::ios::binary | std::ios::out);
    fuzzer_->GetTransformationSequence().SerializeToOstream(
        &binary_transformations);
  }
}

}  // namespace tint::fuzzers::spvtools_fuzzer