// 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/tiff/writer.h"

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

#include "absl/container/flat_hash_set.h"
#include "absl/log/absl_check.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "formats/byte_range.h"
#include "formats/tiff/constants.h"
#include "formats/tiff/extractor.h"
#include "formats/tiff/reader.h"
#include "generator/generator.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "proto/generator.pb.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/endian/endian_writing.h"
#include "testing/nonprod_generator.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::testing::AllOf;
using ::testing::Field;
using ::testing::Gt;
using ::testing::HasSubstr;
using ::testing::Optional;

std::string CreateMinimalTiff(
    TiffEndianness endianness = TiffEndianness::kBigEndian) {
  std::string buffer;
  riegeli::StringWriter writer(&buffer);
  if (endianness == TiffEndianness::kBigEndian) {
    riegeli::WriteBigEndian<uint32_t>(kTiffBigEndianHeader, writer);
    riegeli::WriteBigEndian<uint32_t>(8, writer);  // Offset to first IFD
    riegeli::WriteBigEndian<uint16_t>(2, writer);  // count
    riegeli::WriteBigEndian<uint16_t>(0x0100, writer);
    riegeli::WriteBigEndian<uint16_t>(kTiffTypeShort, writer);
    riegeli::WriteBigEndian<uint32_t>(1, writer);
    riegeli::WriteBigEndian<uint32_t>(10, writer);
    riegeli::WriteBigEndian<uint16_t>(0xffff, writer);
    riegeli::WriteBigEndian<uint16_t>(kTiffTypeShort, writer);
    riegeli::WriteBigEndian<uint32_t>(1, writer);
    riegeli::WriteBigEndian<uint32_t>(20, writer);
    riegeli::WriteBigEndian<uint32_t>(0, writer);  // next IFD offset
  } else {
    riegeli::WriteBigEndian<uint32_t>(kTiffLittleEndianHeader, writer);
    riegeli::WriteLittleEndian<uint32_t>(8, writer);  // Offset to first IFD
    riegeli::WriteLittleEndian<uint16_t>(2, writer);  // count
    riegeli::WriteLittleEndian<uint16_t>(0x0100, writer);
    riegeli::WriteLittleEndian<uint16_t>(kTiffTypeShort, writer);
    riegeli::WriteLittleEndian<uint32_t>(1, writer);
    riegeli::WriteLittleEndian<uint32_t>(10, writer);
    riegeli::WriteLittleEndian<uint16_t>(0xffff, writer);
    riegeli::WriteLittleEndian<uint16_t>(kTiffTypeShort, writer);
    riegeli::WriteLittleEndian<uint32_t>(1, writer);
    riegeli::WriteLittleEndian<uint32_t>(20, writer);
    riegeli::WriteLittleEndian<uint32_t>(0, writer);  // next IFD offset
  }
  writer.Close();
  return buffer;
}

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

  std::unique_ptr<ManifestStoreGenerator> generator_;
  std::unique_ptr<TiffWriter> writer_;
};

TEST_F(TiffWriterTest, WriteMaintainsSortedTags) {
  std::string buffer = CreateMinimalTiff();

  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");

  riegeli::StringReader source(buffer);
  std::string result;
  riegeli::StringWriter destination(&result);
  ABSL_ASSERT_OK(writer_->Write(request, &source, &destination));
  ASSERT_TRUE(source.Close());
  ASSERT_TRUE(destination.Close());

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

  // Validate that output IFD has the C2PA tag, and all tags are properly
  // sorted.
  ASSERT_TRUE(result_reader.Seek(0));
  absl::flat_hash_set<uint32_t> visited_offsets;
  bool validated = false;
  ABSL_ASSERT_OK(IterateOverImageFileDirectories(
      result_reader, /*known_endianness=*/std::nullopt, visited_offsets,
      [&](const TiffImageFileDirectory& ifd,
          TiffEndianness endianness) -> absl::StatusOr<bool> {
        validated = true;
        EXPECT_EQ(ifd.entries.size(), 3);
        EXPECT_EQ(ifd.entries[0].tag, 0x0100);
        EXPECT_EQ(ifd.entries[1].tag, kTiffTagC2pa);
        EXPECT_EQ(ifd.entries[2].tag, 0xffff);
        return true;
      }));
  EXPECT_TRUE(validated);
}

TEST_F(TiffWriterTest, WriteLittleEndianSucceeds) {
  std::string buffer = CreateMinimalTiff(TiffEndianness::kLittleEndian);

  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");

  riegeli::StringReader source(buffer);
  std::string result;
  riegeli::StringWriter destination(&result);
  ABSL_ASSERT_OK(writer_->Write(request, &source, &destination));
  ASSERT_TRUE(source.Close());
  ASSERT_TRUE(destination.Close());

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

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

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

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

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

TEST_F(TiffWriterTest, WriteInPlaceSucceeds) {
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");
  std::string image = CreateMinimalTiff();

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

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

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

}  // namespace
}  // namespace credentio