// Copyright 2023 Google LLC
// SPDX-License-Identifier: BSD-2-Clause

#include "extractgainmap_command.h"

#include "avif/avif_cxx.h"
#include "avifutil.h"
#include "imageio.h"

namespace avif {

ExtractGainMapCommand::ExtractGainMapCommand()
    : ProgramCommand("extractgainmap",
                     "Save the gain map of an AVIF file as an image") {
  argparse_.add_argument(arg_input_filename_, "input_filename");
  argparse_.add_argument(arg_output_filename_, "output_filename");
  arg_image_encode_.Init(argparse_, /*can_have_alpha=*/false);
  arg_jobs_.Init(argparse_);
}

avifResult ExtractGainMapCommand::Run() {
  DecoderPtr decoder(avifDecoderCreate());
  if (decoder == nullptr) {
    return AVIF_RESULT_OUT_OF_MEMORY;
  }
  decoder->maxThreads = arg_jobs_.jobs.value();
  decoder->imageContentToDecode = AVIF_IMAGE_CONTENT_GAIN_MAP;
  decoder->ignoreICC = true;

  avifResult result = ReadAvif(decoder.get(), arg_input_filename_);
  if (result != AVIF_RESULT_OK) {
    return result;
  }

  ImagePtr image(avifImageCreateEmpty());
  if (!image) {
    return AVIF_RESULT_OUT_OF_MEMORY;
  }
  result = avifImageCreateView(image.get(), decoder->image,
                               /*ignoreAlpha=*/false);
  if (result != AVIF_RESULT_OK) {
    return result;
  }

  if (image->gainMap == nullptr || image->gainMap->image == nullptr) {
    std::cerr << "Input image " << arg_input_filename_
              << " does not contain a gain map\n";
    return AVIF_RESULT_INVALID_ARGUMENT;
  }

  return WriteImage(image->gainMap->image,
                    arg_image_encode_.grid.value().grid_cols,
                    arg_image_encode_.grid.value().grid_rows,
                    arg_output_filename_, arg_image_encode_.quality,
                    arg_image_encode_.speed, arg_jobs_.jobs.value());
}

}  // namespace avif