// 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 <algorithm>
#include <cstdint>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/base/nullability.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "bindings/input_hasher.h"
#include "crypto/algorithms.h"
#include "formats/byte_range.h"
#include "formats/tiff/constants.h"
#include "formats/tiff/hasher.h"
#include "formats/tiff/reader.h"
#include "google/protobuf/repeated_ptr_field.h"
#include "proto/data_hash_assertion.pb.h"
#include "proto/generator.pb.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/bytes/writer.h"
#include "riegeli/endian/endian_writing.h"
#include "utils/byte_instruction.h"

namespace credentio {
namespace {

// The maximum number of entries per IFD in a TIFF file.
constexpr uint16_t kTiffMaxEntriesPerIfd = 60000;

template <typename T>
absl::Status WriteEndian(T value, TiffEndianness endianness,
                         riegeli::Writer& writer) {
  const bool success = endianness == TiffEndianness::kBigEndian
                           ? riegeli::WriteBigEndian<T>(value, writer)
                           : riegeli::WriteLittleEndian<T>(value, writer);

  return success ? absl::OkStatus()
                 : writer.StatusOrAnnotate(
                       absl::InternalError("Failed to write value"));
}

struct ArtifactInformation {
  TiffImageFileDirectory ifd_to_add_to;
  TiffEndianness endianness;
};

// Scans the TIFF file for the first IFD that can be used to store the C2PA
// data. If no suitable IFD is found, a new IFD will be created at the end of
// the file. Returns an error if a C2PA entry already exists.
absl::StatusOr<ArtifactInformation> ScanForArtifactInformation(
    riegeli::Reader& source, uint32_t file_size,
    const GenerateManifestStoreRequest& request,
    const TiffWriter::ExclusionProcessor& exclusion_processor,
    DataHashAssertion* absl_nonnull assertion) {
  ArtifactInformation result;

  absl::flat_hash_set<uint32_t> visited_offsets;
  ABSL_RETURN_IF_ERROR(IterateOverImageFileDirectories(
      source, /*known_endianness=*/std::nullopt, visited_offsets,
      [&result, &request, &exclusion_processor, assertion, &source,
       &visited_offsets](const TiffImageFileDirectory& ifd,
                         TiffEndianness endianness) -> absl::StatusOr<bool> {
        for (const TiffImageFileDirectoryEntry& entry : ifd.entries) {
          if (entry.tag == kTiffTagC2pa) {
            return absl::AlreadyExistsError("C2PA entry already exists");
          }
          if (exclusion_processor != nullptr) {
            ABSL_ASSIGN_OR_RETURN(
                std::vector<ByteRange> exclusions,
                exclusion_processor(source, entry, endianness, visited_offsets,
                                    request));
            for (const ByteRange& exclusion_range : exclusions) {
              auto* exclusion = assertion->add_exclusions();
              exclusion->set_start(exclusion_range.offset);
              exclusion->set_length(exclusion_range.length);
            }
          }
        }

        result.endianness = endianness;
        result.ifd_to_add_to = ifd;
        return true;
      }));

  if (result.ifd_to_add_to.offset_of_pointer != 4 ||
      result.ifd_to_add_to.entries.size() >= kTiffMaxEntriesPerIfd) {
    // Need to insert a new IFD at the end of the file
    result.ifd_to_add_to = TiffImageFileDirectory{
        .offset_of_pointer = static_cast<uint32_t>(source.pos()),
        .offset = file_size,
    };
  }

  return result;
}

template <typename T>
absl::Status AppendNumberToString(T value, TiffEndianness endianness,
                                  std::string& output) {
  riegeli::StringWriter writer(
      &output, riegeli::StringWriterBase::Options().set_append(true));

  const bool success = endianness == TiffEndianness::kBigEndian
                           ? riegeli::WriteBigEndian<T>(value, writer)
                           : riegeli::WriteLittleEndian<T>(value, writer);
  if (!success) {
    return writer.StatusOrAnnotate(
        absl::InternalError("Failed to write value"));
  }
  if (!writer.Close()) {
    return writer.StatusOrAnnotate(
        absl::InternalError("Failed to close writer"));
  }
  return absl::OkStatus();
}

// Adjusts the pointer to the IFD to point to the new location. If the IFD is
// being moved, zero out the original location.
absl::StatusOr<std::vector<ByteInstruction>>
CreateByteInstructionsForIfdInsertion(
    const TiffImageFileDirectory& ifd_to_add_to, TiffEndianness endianness,
    uint32_t file_size) {
  std::vector<ByteInstruction> byte_instructions;

  // Add a byte instruction to update pointer to the IFD.
  byte_instructions.push_back(ByteInstruction{
      .operation = ByteInstruction::Operation::kReplace,
      .offset = ifd_to_add_to.offset_of_pointer,
  });
  ABSL_RETURN_IF_ERROR(AppendNumberToString<uint32_t>(
      file_size, endianness, byte_instructions.back().bytes));

  if (ifd_to_add_to.offset_of_pointer == 4) {
    // C2PA data will be written to the first (and only) IFD, it will be moved
    // to the end of the file, zero out the original location.
    byte_instructions.push_back(ByteInstruction{
        .operation = ByteInstruction::Operation::kReplace,
        .offset = ifd_to_add_to.offset,
        .bytes = std::string(
            ByteSizeOfImageFileDirectory(ifd_to_add_to.entries.size()), '\0'),
    });
  }
  return std::move(byte_instructions);
}

void SanitizeExclusions(
    google::protobuf::RepeatedPtrField<DataRange>* exclusions) {
  if (exclusions->empty()) return;
  std::sort(exclusions->pointer_begin(), exclusions->pointer_end(),
            [](const DataRange* a, const DataRange* b) {
              if (a->start() != b->start()) return a->start() < b->start();
              return a->length() < b->length();
            });

  google::protobuf::RepeatedPtrField<DataRange> merged;
  DataRange current = (*exclusions)[0];
  for (int i = 1; i < exclusions->size(); ++i) {
    const DataRange& next = (*exclusions)[i];
    uint64_t current_end = current.start() + current.length();
    uint64_t next_end = next.start() + next.length();
    if (next.start() <= current_end) {
      if (next_end > current_end) {
        current.set_length(next_end - current.start());
      }
    } else {
      *merged.Add() = current;
      current = next;
    }
  }
  *merged.Add() = current;
  exclusions->Swap(&merged);
}

absl::Status HashAroundExclusion(const DataRange& exclusion,
                                 riegeli::Reader& source, InputHasher& hasher) {
  if (exclusion.start() < source.pos()) {
    return absl::InvalidArgumentError(
        "Exclusion start is before current source position");
  }
  // Hash up to the start of the exclusion.
  uint64_t hash_length = exclusion.start() - source.pos();
  ABSL_RETURN_IF_ERROR(ReadHash(source, hash_length, hasher));

  // Skip the excluded data
  if (!source.Skip(exclusion.length())) {
    return absl::InternalError("failed to skip data");
  }

  return absl::OkStatus();
}

absl::Status HashInstructionAroundExclusion(const DataRange& exclusion,
                                            uint64_t instruction_offset,
                                            absl::string_view instruction_bytes,
                                            InputHasher& hasher) {
  if (exclusion.start() < instruction_offset ||
      exclusion.start() - instruction_offset > instruction_bytes.size()) {
    return absl::InvalidArgumentError(
        "Exclusion range out of bounds of byte instruction");
  }

  // hash until the exclusion offset
  absl::string_view data_to_hash =
      instruction_bytes.substr(0, exclusion.start() - instruction_offset);
  hasher.Update(data_to_hash);

  if (exclusion.length() >
      instruction_bytes.size() - (exclusion.start() - instruction_offset)) {
    return absl::InvalidArgumentError(
        "Exclusion data out of bounds of byte instruction");
  }

  return absl::OkStatus();
}

absl::Status Hash(riegeli::Reader& source, uint64_t length_with_ifd,
                  absl::Span<const ByteInstruction> byte_instructions,
                  const DataHashAssertion& assertion, InputHasher& hasher) {
  auto exclusion_iter = assertion.exclusions().begin();

  for (const ByteInstruction& ins : byte_instructions) {
    while (exclusion_iter->start() < ins.offset) {
      ABSL_RETURN_IF_ERROR(
          HashAroundExclusion(*exclusion_iter, source, hasher));
      ++exclusion_iter;
      if (exclusion_iter == assertion.exclusions().end()) {
        return absl::InternalError("ran out of exclusions");
      }
    }
    // Next exclusion is after the start of the instruction, hash up to the
    // instruction offset.
    if (ins.offset < source.pos()) {
      return absl::InvalidArgumentError(
          "Byte instruction offset is before current source position");
    }
    ABSL_RETURN_IF_ERROR(ReadHash(source, ins.offset - source.pos(), hasher));

    uint64_t ins_processed = 0;
    while (exclusion_iter->start() < ins.offset + ins.bytes.size()) {
      ABSL_RETURN_IF_ERROR(HashInstructionAroundExclusion(
          *exclusion_iter, ins.offset + ins_processed,
          absl::string_view(ins.bytes).substr(ins_processed), hasher));

      ins_processed =
          (exclusion_iter->start() + exclusion_iter->length()) - ins.offset;

      ++exclusion_iter;
      if (exclusion_iter == assertion.exclusions().end()) {
        return absl::InternalError("ran out of exclusions");
      }
    }

    // Process remaining instruction bytes
    if (ins_processed < ins.bytes.size()) {
      absl::string_view data_to_hash =
          absl::string_view(ins.bytes).substr(ins_processed);
      hasher.Update(data_to_hash);
    }

    if (ins.operation == ByteInstruction::Operation::kReplace) {
      if (!source.Skip(ins.bytes.size())) {
        return absl::InternalError("failed to skip data");
      }
    }
  }

  return absl::OkStatus();
}

absl::Status WriteImageFileDirectory(
    absl::Span<const TiffImageFileDirectoryEntry> entries,
    TiffEndianness endianness, riegeli::Writer& writer) {
  // Entry Count
  ABSL_RETURN_IF_ERROR(
      WriteEndian<uint16_t>(entries.size(), endianness, writer));
  // Entries
  for (const auto& entry : entries) {
    ABSL_RETURN_IF_ERROR(WriteEndian<uint16_t>(entry.tag, endianness, writer));
    ABSL_RETURN_IF_ERROR(WriteEndian<uint16_t>(entry.type, endianness, writer));
    ABSL_RETURN_IF_ERROR(
        WriteEndian<uint32_t>(entry.count, endianness, writer));
    ABSL_RETURN_IF_ERROR(
        WriteEndian<uint32_t>(entry.value_or_offset, endianness, writer));
  }
  // Next IFD Offset
  ABSL_RETURN_IF_ERROR(WriteEndian<uint32_t>(0, endianness, writer));
  return absl::OkStatus();
}

absl::Status WriteImageFileDirectory(
    absl::Span<const TiffImageFileDirectoryEntry> entries,
    TiffEndianness endianness, std::string& destination) {
  riegeli::StringWriter<> writer(
      &destination, riegeli::StringWriterBase::Options().set_append(true));
  ABSL_RETURN_IF_ERROR(WriteImageFileDirectory(entries, endianness, writer));
  if (!writer.Close()) {
    return writer.status();
  }
  return absl::OkStatus();
}

}  // namespace

absl::StatusOr<std::vector<ByteInstruction>> TiffWriter::WriteInternal(
    GenerateManifestStoreRequest request, riegeli::Reader* absl_nonnull source,
    OutputMetadata* absl_nullable output_metadata) const {
  if (request.has_hard_binding()) {
    return absl::InvalidArgumentError(
        "request must not include `hard_binding` (binding is computed by "
        "TiffWriter)");
  }
  if (!source->SupportsSize()) {
    return absl::InvalidArgumentError("source must support Size()");
  }
  if (!source->Size().has_value()) {
    return absl::InternalError("source->Size() returned no value");
  }

  uint64_t source_size = *source->Size();
  uint64_t padding_size = source_size % 2;
  uint64_t ifd_offset = source_size + padding_size;

  auto hash_algorithm = HashAlgorithm::kSha256;

  auto* assertion = request.mutable_hard_binding()->mutable_data_hash();
  assertion->set_alg(absl::StrCat(hash_algorithm));
  assertion->set_hash(std::string(32, '\0'));

  ABSL_ASSIGN_OR_RETURN(
      ArtifactInformation artifact_information,
      ScanForArtifactInformation(*source, ifd_offset, request,
                                 options_.exclusion_processor, assertion));

  ABSL_ASSIGN_OR_RETURN(std::vector<ByteInstruction> byte_instructions,
                        CreateByteInstructionsForIfdInsertion(
                            artifact_information.ifd_to_add_to,
                            artifact_information.endianness, ifd_offset));

  uint32_t file_size_after_ifd =
      ifd_offset + ByteSizeOfImageFileDirectory(
                       artifact_information.ifd_to_add_to.entries.size() + 1);

  artifact_information.ifd_to_add_to.entries.push_back(
      TiffImageFileDirectoryEntry{
          .tag = kTiffTagC2pa,
          .type = kTiffTypeUndefined,
          .count = 1,
          .value_or_offset = file_size_after_ifd,
      });

  std::sort(artifact_information.ifd_to_add_to.entries.begin(),
            artifact_information.ifd_to_add_to.entries.end(),
            [](const TiffImageFileDirectoryEntry& a,
               const TiffImageFileDirectoryEntry& b) { return a.tag < b.tag; });

  {
    // Holder for the C2PA IFD entry count.
    auto* count_exclusion = assertion->add_exclusions();
    // IFD count, IFD Entry tag, and IFD Entry type are 3 uint16_t's

    uint16_t c2pa_entry_index = 0;
    for (const auto& entry : artifact_information.ifd_to_add_to.entries) {
      if (entry.tag == kTiffTagC2pa) {
        break;
      }
      ++c2pa_entry_index;
    }

    count_exclusion->set_start(ifd_offset +        // Offset of relocated IFD
                               sizeof(uint16_t) +  // final IFD count
                               // Number of entries before C2PA entry
                               (c2pa_entry_index * kTiffIfdEntrySize) +
                               sizeof(uint16_t) +   // C2PA entry tag
                               sizeof(uint16_t));   // C2PA entry type
    count_exclusion->set_length(sizeof(uint32_t));  // C2PA entry count

    // Holder for the Manifest Store.
    auto* length_exclusion = assertion->add_exclusions();
    // Will be located after the IFD containing the C2PA data.
    length_exclusion->set_start(file_size_after_ifd);
    length_exclusion->set_length(1);
  }

  SanitizeExclusions(assertion->mutable_exclusions());

  ABSL_ASSIGN_OR_RETURN(auto template_response,
                        generator_.GenerateManifestStoreTemplate(request));

  if ((std::numeric_limits<uint32_t>::max() - file_size_after_ifd -
       template_response.manifest_store_size()) < 0) {
    return absl::InvalidArgumentError(
        "Manifest store size is too large (not enough space for TIFF "
        "structure)");
  }

  // Update the length exclusion to the actual length of the manifest store.
  assertion->mutable_exclusions()->rbegin()->set_length(
      template_response.manifest_store_size());
  for (auto& entry : artifact_information.ifd_to_add_to.entries) {
    if (entry.tag == kTiffTagC2pa) {
      entry.count = template_response.manifest_store_size();
      break;
    }
  }

  {
    ByteInstruction last_ifd{
        .operation = ByteInstruction::Operation::kInsert,
        .offset = source_size,
        .bytes = std::string(padding_size, '\0'),
    };
    ABSL_RETURN_IF_ERROR(WriteImageFileDirectory(
        artifact_information.ifd_to_add_to.entries,
        artifact_information.endianness, last_ifd.bytes));
    byte_instructions.push_back(last_ifd);
  }

  std::sort(byte_instructions.begin(), byte_instructions.end(),
            [](const ByteInstruction& a, const ByteInstruction& b) {
              return a.offset < b.offset;
            });

  source->Seek(0);

  std::unique_ptr<credentio::InputHasher> hasher;
  ABSL_ASSIGN_OR_RETURN(hasher, InputHasher::Create(hash_algorithm));

  ABSL_RETURN_IF_ERROR(Hash(*source, file_size_after_ifd, byte_instructions,
                            *assertion, *hasher));

  FinalizeManifestStoreRequest finalize_request;
  *finalize_request.mutable_manifest_store_template() =
      std::move(template_response);
  *finalize_request.mutable_hard_binding()->mutable_data_hash() =
      request.hard_binding().data_hash();
  finalize_request.mutable_hard_binding()->mutable_data_hash()->set_hash(
      hasher->Digest());

  ABSL_ASSIGN_OR_RETURN(auto finalize_response,
                        generator_.FinalizeManifestStore(finalize_request));

  byte_instructions.push_back(
      ByteInstruction{.operation = ByteInstruction::Operation::kInsert,
                      .offset = static_cast<uint64_t>(file_size_after_ifd),
                      .bytes = finalize_response.manifest_store()});

  std::sort(byte_instructions.begin(), byte_instructions.end(),
            [](const ByteInstruction& a, const ByteInstruction& b) {
              return a.offset < b.offset;
            });

  if (output_metadata != nullptr) {
    output_metadata->timestamp_included = finalize_response.is_timestamped();
  }

  return byte_instructions;
}

absl::Status TiffWriter::Write(
    GenerateManifestStoreRequest request, riegeli::Reader* absl_nonnull source,
    riegeli::Writer* absl_nonnull destination,
    OutputMetadata* absl_nullable output_metadata) const {
  ABSL_ASSIGN_OR_RETURN(
      auto byte_instructions,
      WriteInternal(std::move(request), source, output_metadata));

  return ApplyByteInstructions(source, byte_instructions, destination);
}

absl::Status TiffWriter::Write(
    GenerateManifestStoreRequest request, riegeli::Writer* absl_nonnull asset,
    OutputMetadata* absl_nullable output_metadata) const {
  if (!asset->SupportsRandomAccess()) {
    return absl::InvalidArgumentError("asset must support RandomAccess()");
  }
  if (!asset->SupportsReadMode()) {
    return absl::InvalidArgumentError("asset must support ReadMode()");
  }

  riegeli::Reader* source = asset->ReadMode(0);
  if (source == nullptr) {
    return absl::InternalError("Failed to switch asset to ReadMode");
  }

  ABSL_ASSIGN_OR_RETURN(
      auto byte_instructions,
      WriteInternal(std::move(request), source, output_metadata));

  return ApplyByteInstructions(byte_instructions, asset);
}

}  // namespace credentio