// Copyright 2026 Google LLC
//
// 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
//
//    https://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 "formats/riff/writer.h"

#include <cstdint>
#include <memory>
#include <string>
#include <utility>

#include "absl/log/absl_check.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/strings/str_cat.h"
#include "formats/byte_range.h"
#include "formats/riff/create_riff.h"
#include "formats/riff/extractor.h"
#include "generator/generator.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "proto/generator.pb.h"
#include "proto/manifest.pb.h"
#include "proto/validation_status.pb.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"
#include "testing/nonprod_generator.h"
#include "testing/nonprod_validator.h"
#include "testing/proto_test_utils.h"
#include "validator/asset_validator.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::credentio_testing::ParseTextProtoOrDie;
using ::credentio_testing::PartiallyEqualsProto;
using ::testing::AllOf;
using ::testing::Contains;
using ::testing::Field;
using ::testing::Gt;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::Not;
using ::testing::Optional;

std::string CreateWebPImage() {
  return credentio_riff::RiffChunk("WEBP",
                                   credentio_riff::Chunk("VP8 ", "image_data"));
}

std::string CreateAviVideo() {
  return credentio_riff::RiffChunk(
      "AVI ",
      absl::StrCat(
          credentio_riff::ListChunk(
              "hdrl", credentio_riff::Chunk("avih", std::string(56, '\0'))),
          credentio_riff::ListChunk(
              "movi", credentio_riff::Chunk("00dc", "fake video frame data"))));
}

class RiffWriterTest : public ::testing::Test {
 protected:
  void SetUp() override {
    auto generator = CreateNonProdManifestStoreGenerator({});
    ABSL_CHECK_OK(generator);
    generator_ = *std::move(generator);
    writer_ = std::make_unique<RiffWriter>(generator_.get());

    auto validator = CreateNonProdValidator({});
    ABSL_CHECK_OK(validator);
    validator_ = *std::move(validator);
  }

  std::unique_ptr<ManifestStoreGenerator> generator_;
  std::unique_ptr<RiffWriter> writer_;
  std::unique_ptr<AssetValidator> validator_;
};

TEST_F(RiffWriterTest, Write) {
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");

  std::string image = CreateWebPImage();
  riegeli::StringReader source(image);
  ABSL_ASSERT_OK(source.status());
  std::string result;
  riegeli::StringWriter destination(&result);
  ABSL_ASSERT_OK(destination.status());
  RiffWriter::OutputMetadata output_metadata;
  ABSL_ASSERT_OK(
      writer_->Write(request, &source, &destination, &output_metadata));
  ASSERT_TRUE(source.Close());
  ASSERT_TRUE(destination.Close());
  EXPECT_TRUE(output_metadata.timestamp_included);

  riegeli::StringReader result_reader(result);
  ABSL_ASSERT_OK(result_reader.status());

  EXPECT_THAT(RiffExtractor().ExtractManifestStoreLocation(result_reader, {}),
              IsOkAndHolds(Optional(AllOf(Field(&ByteRange::offset, Gt(0)),
                                          Field(&ByteRange::length, Gt(0))))));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validation_result = validator_->Validate(result_reader, "image/webp");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
  EXPECT_THAT(
      (*validation_result)
          ->proto()
          .active_manifest()
          .validation()
          .informationals(),
      Not(Contains(PartiallyEqualsProto(ParseTextProtoOrDie<ValidationStatus>(
          R"pb(
            code: "assertion.dataHash.additionalExclusionsPresent"
          )pb")))));
}

TEST_F(RiffWriterTest, WriteInPlace) {
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");

  std::string image = CreateWebPImage();
  riegeli::StringWriter<> asset(
      &image, riegeli::StringWriterBase::Options().set_append(true));
  ABSL_ASSERT_OK(asset.status());

  RiffWriter::OutputMetadata output_metadata;
  ABSL_ASSERT_OK(writer_->Write(request, &asset, &output_metadata));
  ASSERT_TRUE(asset.Close());
  EXPECT_TRUE(output_metadata.timestamp_included);

  riegeli::StringReader<> result_reader(image);
  ABSL_ASSERT_OK(result_reader.status());

  EXPECT_THAT(RiffExtractor().ExtractManifestStoreLocation(result_reader, {}),
              IsOkAndHolds(Optional(AllOf(Field(&ByteRange::offset, Gt(0)),
                                          Field(&ByteRange::length, Gt(0))))));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validation_result = validator_->Validate(result_reader, "image/webp");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
  EXPECT_THAT(
      (*validation_result)
          ->proto()
          .active_manifest()
          .validation()
          .informationals(),
      Not(Contains(PartiallyEqualsProto(ParseTextProtoOrDie<ValidationStatus>(
          R"pb(
            code: "assertion.dataHash.additionalExclusionsPresent"
          )pb")))));
}

TEST_F(RiffWriterTest, WriteAvi) {
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test avi");

  std::string video = CreateAviVideo();
  riegeli::StringReader source(video);
  ABSL_ASSERT_OK(source.status());
  std::string result;
  riegeli::StringWriter destination(&result);
  ABSL_ASSERT_OK(destination.status());
  RiffWriter::OutputMetadata output_metadata;
  ABSL_ASSERT_OK(
      writer_->Write(request, &source, &destination, &output_metadata));
  ASSERT_TRUE(source.Close());
  ASSERT_TRUE(destination.Close());
  EXPECT_TRUE(output_metadata.timestamp_included);

  riegeli::StringReader result_reader(result);
  ABSL_ASSERT_OK(result_reader.status());

  EXPECT_THAT(RiffExtractor().ExtractManifestStoreLocation(result_reader, {}),
              IsOkAndHolds(Optional(AllOf(Field(&ByteRange::offset, Gt(0)),
                                          Field(&ByteRange::length, Gt(0))))));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validation_result =
      validator_->Validate(result_reader, "video/x-msvideo");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
  EXPECT_THAT(
      (*validation_result)
          ->proto()
          .active_manifest()
          .validation()
          .informationals(),
      Not(Contains(PartiallyEqualsProto(ParseTextProtoOrDie<ValidationStatus>(
          R"pb(
            code: "assertion.dataHash.additionalExclusionsPresent"
          )pb")))));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validation_result_no_hint =
      validator_->Validate(result_reader, /*media_type=*/{});
  ABSL_ASSERT_OK(validation_result_no_hint);
  EXPECT_THAT((*validation_result_no_hint)
                  ->proto()
                  .active_manifest()
                  .validation()
                  .failures(),
              IsEmpty());
}

TEST_F(RiffWriterTest, WriteWithOutputMetadataNoTimestamp) {
  auto generator =
      CreateNonProdManifestStoreGenerator({.disable_timestamping = true});
  ABSL_ASSERT_OK(generator);
  auto writer = std::make_unique<RiffWriter>(generator->get());
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");
  std::string image = CreateWebPImage();
  riegeli::StringReader source(image);
  ABSL_ASSERT_OK(source.status());
  std::string result;
  riegeli::StringWriter destination(&result);
  ABSL_ASSERT_OK(destination.status());
  RiffWriter::OutputMetadata output_metadata;
  ABSL_ASSERT_OK(
      writer->Write(request, &source, &destination, &output_metadata));
  ASSERT_TRUE(source.Close());
  ASSERT_TRUE(destination.Close());
  EXPECT_FALSE(output_metadata.timestamp_included);
}

TEST_F(RiffWriterTest, WriteTwiceFails) {
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");

  std::string image = CreateWebPImage();
  riegeli::StringReader source1(image);
  ABSL_ASSERT_OK(source1.status());
  std::string result1;
  riegeli::StringWriter destination1(&result1);
  ABSL_ASSERT_OK(destination1.status());
  ABSL_ASSERT_OK(writer_->Write(request, &source1, &destination1));
  ASSERT_TRUE(source1.Close());
  ASSERT_TRUE(destination1.Close());

  riegeli::StringReader source2(result1);
  ABSL_ASSERT_OK(source2.status());
  std::string result2;
  riegeli::StringWriter destination2(&result2);
  ABSL_ASSERT_OK(destination2.status());
  EXPECT_THAT(
      writer_->Write(request, &source2, &destination2),
      StatusIs(absl::StatusCode::kAlreadyExists, HasSubstr("already exists")));
}

TEST_F(RiffWriterTest, WriteWithHardBindingFails) {
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");
  request.mutable_hard_binding();
  std::string image = CreateWebPImage();
  riegeli::StringReader source(image);
  ABSL_ASSERT_OK(source.status());
  std::string result;
  riegeli::StringWriter destination(&result);
  ABSL_ASSERT_OK(destination.status());
  EXPECT_THAT(writer_->Write(request, &source, &destination),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("must not include `hard_binding`")));
}

TEST_F(RiffWriterTest, DemonstratesUnboundedMemoryBufferingDoS) {
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("demonstrate unbounded buffering DoS");

  constexpr uint32_t kJunkPayloadSize = 10 * 1024 * 1024;  // 10 MiB
  constexpr uint32_t kTotalRiffSize = 4 + 8 + kJunkPayloadSize;

  std::string large_riff_data;
  large_riff_data.reserve(8 + kTotalRiffSize);

  // RIFF Header: "RIFF" <total_riff_size> "WEBP"
  large_riff_data.append("RIFF", 4);
  large_riff_data.push_back(static_cast<char>(kTotalRiffSize & 0xFF));
  large_riff_data.push_back(static_cast<char>((kTotalRiffSize >> 8) & 0xFF));
  large_riff_data.push_back(static_cast<char>((kTotalRiffSize >> 16) & 0xFF));
  large_riff_data.push_back(static_cast<char>((kTotalRiffSize >> 24) & 0xFF));
  large_riff_data.append("WEBP", 4);

  // JUNK Subchunk: "JUNK" <junk_payload_size> <junk_payload>
  large_riff_data.append("JUNK", 4);
  large_riff_data.push_back(static_cast<char>(kJunkPayloadSize & 0xFF));
  large_riff_data.push_back(static_cast<char>((kJunkPayloadSize >> 8) & 0xFF));
  large_riff_data.push_back(static_cast<char>((kJunkPayloadSize >> 16) & 0xFF));
  large_riff_data.push_back(static_cast<char>((kJunkPayloadSize >> 24) & 0xFF));
  large_riff_data.append(kJunkPayloadSize, 'A');

  riegeli::StringReader source(large_riff_data);
  ABSL_ASSERT_OK(source.status());

  std::string result;
  riegeli::StringWriter destination(&result);
  ABSL_ASSERT_OK(destination.status());

  ABSL_ASSERT_OK(writer_->Write(request, &source, &destination));
  ASSERT_TRUE(source.Close());
  ASSERT_TRUE(destination.Close());

  // Verify that the output was successfully generated and contains the 10MB
  // media payload
  EXPECT_GT(result.size(), kJunkPayloadSize);
}

}  // namespace
}  // namespace credentio