// 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/jpeg/hasher.h"

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

#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/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "bindings/input_hasher.h"
#include "crypto/algorithms.h"
#include "formats/byte_range.h"
#include "formats/jpeg/assessor.h"
#include "formats/jpeg/c2pa_segment_handler.h"
#include "formats/jpeg/reader.h"
#include "proto/data_hash_assertion.pb.h"
#include "proto/generator.pb.h"
#include "proto/multi_asset_hash_assertion.pb.h"
#include "riegeli/bytes/reader.h"

namespace credentio {

namespace {

constexpr uint64_t kMaxChunkSize = 1024 * 1024;  // 1 MiB

// Updates the given hashers with the indicated content, leaving the input
// position at the end of the content.
absl::Status UpdateTwoHashers(riegeli::Reader& input, int64_t offset,
                              int64_t length, InputHasher& hasher1,
                              InputHasher& hasher2) {
  if (!input.Seek(offset)) {
    return absl::InvalidArgumentError(
        absl::StrCat("Failed to seek to offset ", offset));
  }

  int64_t chunk_size = kMaxChunkSize;
  for (int64_t i = 0; i < length; i += chunk_size) {
    int64_t l = std::min(length - i, chunk_size);

    std::string data;
    if (!input.Read(l, data)) {
      return absl::InvalidArgumentError("Failed to read data");
    }

    hasher1.Update(data);
    hasher2.Update(data);
  }

  return absl::OkStatus();
}

// Returns true if the current box hash range should end based on the next
// segment header and whether a C2PA segment needs to be injected.
bool InsertC2paBeforeSegment(absl::string_view next_segment_label) {
  if (next_segment_label == "SOI") {
    return false;
  }
  if (!absl::StartsWith(next_segment_label, "APP")) {
    // Inject before any non-APP segments.
    return true;
  }
  if (next_segment_label == "APP2") {
    // Inject before the APP2 segment.
    return true;
  }
  if (next_segment_label == "APP11") {
    // Inject before any existing APP11 segments.
    return true;
  }
  return false;
}

absl::StatusOr<uint16_t> GetUnusedSegmentId(
    const absl::flat_hash_set<uint16_t>& used_segment_ids) {
  if (used_segment_ids.size() == UINT16_MAX) {
    return absl::InvalidArgumentError("No Available Segment ID");
  }
  uint16_t next_segment_id = 1;
  while (used_segment_ids.contains(next_segment_id)) {
    ++next_segment_id;
  }
  return next_segment_id;
}

// Computes data hash(es) for a JPEG asset (which may be a part of a multi-part
// asset). If the starting input position is 0, then
//  * a placeholder exclusion will be added for the C2PA segment
//  * the overall file hash will be computed (result.hard_binding.data_hash)
// In all cases, a part hash will be computed, covering the asset starting at
// the current input position and ending at the next EOI marker.
//
// If `exclusion_processor` is provided, it is invoked for each segment to
// determine byte ranges to exclude from the hashes.
absl::Status HashJpeg(riegeli::Reader& input, HashAlgorithm alg,
                      std::optional<JpegExclusionProcessor> exclusion_processor,
                      InputHasher& asset_hasher, DigestJpegResult& result) {
  int64_t asset_start_offset = input.pos();
  bool is_first_asset = asset_start_offset == 0;

  C2paSegmentHandler c2pa_segment_handler;
  absl::flat_hash_set<uint16_t> used_segment_ids;

  std::optional<JpegInsertionPoint> insertion_point;

  ABSL_ASSIGN_OR_RETURN(auto part_hasher, InputHasher::Create(alg));

  PartBinding* part_binding = result.hard_binding.add_part_bindings();
  part_binding->mutable_location()->set_byte_offset(asset_start_offset);
  part_binding->set_is_optional(!is_first_asset);
  part_binding->mutable_data_hash()->set_alg(absl::StrCat(alg));

  std::optional<JpegSegment> data_past_asset = std::nullopt;

  ABSL_RETURN_IF_ERROR(IterateOverJpegSegments(
      input, [&](const JpegSegment& segment) -> absl::StatusOr<bool> {
        if (segment.marker_label == "c2pa.after") {
          data_past_asset = std::move(segment);
          return false;  // stop iteration
        }

        if (segment.app11_info.has_value()) {
          used_segment_ids.insert(segment.app11_info->segment_id);
        }

        if (is_first_asset && !insertion_point.has_value() &&
            InsertC2paBeforeSegment(segment.marker_label)) {
          // Add an exclusion placeholder for the C2PA segment. Length will be
          // filled in during encoding.
          auto* asset_exclusion =
              result.hard_binding.mutable_data_hash()->add_exclusions();
          asset_exclusion->set_start(segment.offset);
          asset_exclusion->set_length(0);

          auto* part_exclusion =
              part_binding->mutable_data_hash()->add_exclusions();
          part_exclusion->set_start(segment.offset - asset_start_offset);
          part_exclusion->set_length(0);

          insertion_point = JpegInsertionPoint{
              .segment_id = 1,  // to be updated with the next available ID
              .offset = segment.offset,
          };
        }

        std::vector<ByteRange> exclusions;
        if (exclusion_processor.has_value()) {
          ABSL_ASSIGN_OR_RETURN(exclusions,
                                (*exclusion_processor)(input, segment));
        }
        int64_t offset = segment.offset;
        for (ByteRange exclusion : exclusions) {
          uint64_t segment_end = segment.offset + segment.length;
          if (exclusion.offset < offset || exclusion.offset > segment_end ||
              exclusion.length > segment_end - exclusion.offset) {
            return absl::InvalidArgumentError(absl::StrCat(
                "Invalid or overlapping exclusion range [", exclusion.offset,
                ", ", exclusion.offset + exclusion.length, ") for segment at ",
                segment.offset));
          }

          auto* asset_exclusion =
              result.hard_binding.mutable_data_hash()->add_exclusions();
          asset_exclusion->set_start(exclusion.offset);
          asset_exclusion->set_length(exclusion.length);

          auto* part_exclusion =
              part_binding->mutable_data_hash()->add_exclusions();
          part_exclusion->set_start(exclusion.offset - asset_start_offset);
          part_exclusion->set_length(exclusion.length);

          ABSL_RETURN_IF_ERROR(UpdateTwoHashers(input, offset,
                                                exclusion.offset - offset,
                                                asset_hasher, *part_hasher));
          offset = exclusion.offset + exclusion.length;
        }
        if (offset < (segment.offset + segment.length)) {
          ABSL_RETURN_IF_ERROR(UpdateTwoHashers(
              input, offset, segment.offset + segment.length - offset,
              asset_hasher, *part_hasher));
        }

        c2pa_segment_handler.RecordSegment(segment);
        return true;
      }));

  if (is_first_asset) {
    absl::Status c2pa_segment_status =
        c2pa_segment_handler.GetC2paSegment(input).status();
    if (c2pa_segment_status.ok() ||
        absl::StrContains(c2pa_segment_status.message(), "multiple")) {
      return absl::AlreadyExistsError("C2PA chunk already exists");
    }

    if (!insertion_point.has_value()) {
      return absl::InvalidArgumentError(
          "No valid C2PA insertion point found in JPEG");
    }
    ABSL_ASSIGN_OR_RETURN(uint16_t next_segment_id,
                          GetUnusedSegmentId(used_segment_ids));
    insertion_point->segment_id = next_segment_id;

    result.insertion_point = std::move(*insertion_point);
  }

  if (data_past_asset.has_value()) {
    if (!input.Seek(data_past_asset->offset)) {
      return absl::InvalidArgumentError("Failed to seek to data past asset");
    }
  }

  part_binding->mutable_location()->set_length(input.pos() -
                                               asset_start_offset);
  part_binding->mutable_data_hash()->set_hash(part_hasher->Digest());
  return absl::OkStatus();
}

// Adds a part hash covering the remaining bytes of the input, marked as
// optional.
absl::Status HashRemainingBytes(riegeli::Reader& input, HashAlgorithm alg,
                                InputHasher& asset_hasher,
                                DigestJpegResult& result) {
  if (!input.SupportsSize() || !input.Size().has_value()) {
    return absl::InvalidArgumentError("Failed to get size of input");
  }

  int64_t asset_start_offset = input.pos();
  int64_t asset_length = *input.Size() - asset_start_offset;

  ABSL_ASSIGN_OR_RETURN(auto part_hasher, InputHasher::Create(alg));

  ABSL_RETURN_IF_ERROR(UpdateTwoHashers(input, asset_start_offset, asset_length,
                                        asset_hasher, *part_hasher));

  auto* part_binding = result.hard_binding.add_part_bindings();
  part_binding->mutable_location()->set_byte_offset(asset_start_offset);
  part_binding->mutable_location()->set_length(asset_length);
  part_binding->set_is_optional(true);
  part_binding->mutable_data_hash()->set_alg(absl::StrCat(alg));
  part_binding->mutable_data_hash()->set_hash(part_hasher->Digest());

  return absl::OkStatus();
}

}  // namespace

absl::StatusOr<DigestJpegResult> DigestJpeg(
    riegeli::Reader& input, HashAlgorithm alg, bool generate_multi_asset_hash,
    std::optional<JpegExclusionProcessor> exclusion_processor) {
  ABSL_ASSIGN_OR_RETURN(auto correct_format, JpegAssessor().IsSupported(input));
  if (!correct_format) {
    return absl::InvalidArgumentError("File does not start with a valid JPEG");
  }
  if (!input.Seek(0)) {
    return absl::InvalidArgumentError("Failed to seek to start of input");
  }
  if (!input.SupportsSize() || !input.Size().has_value()) {
    return absl::InvalidArgumentError("Failed to get size of input");
  }

  DigestJpegResult result = {};
  result.hard_binding.mutable_data_hash()->set_alg(absl::StrCat(alg));

  uint64_t input_size = *input.Size();

  ABSL_ASSIGN_OR_RETURN(auto asset_hasher, InputHasher::Create(alg));

  int64_t pos;
  do {
    pos = input.pos();

    ABSL_ASSIGN_OR_RETURN(bool is_jpeg, JpegAssessor().IsSupported(input));
    if (is_jpeg) {
      ABSL_RETURN_IF_ERROR(
          HashJpeg(input, alg, exclusion_processor, *asset_hasher, result));
    } else {
      ABSL_RETURN_IF_ERROR(
          HashRemainingBytes(input, alg, *asset_hasher, result));
    }

    if (!generate_multi_asset_hash) {
      break;
    }
    int64_t newpos = input.pos();
    if (pos == newpos && pos < input_size) {
      // Avoid infinite loop in case of bugs.
      return absl::InternalError("Hashing made no progress");
    }
    pos = newpos;
  } while (pos < input_size);

  if (!generate_multi_asset_hash && input.pos() < input_size) {
    ABSL_RETURN_IF_ERROR(
        asset_hasher->Update(input, input.pos(), input_size - input.pos()));
  }

  result.hard_binding.mutable_data_hash()->set_hash(asset_hasher->Digest());

  if (!generate_multi_asset_hash ||
      result.hard_binding.part_bindings_size() == 1) {
    // Only a single asset was found, no need for multi-asset hash assertion.
    result.hard_binding.clear_part_bindings();
  }

  return result;
}

}  // namespace credentio