// Copyright 2025 The PDFium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "core/fpdfapi/page/jpx_decode_conversion.h"

#include <optional>

#include "core/fpdfapi/page/cpdf_colorspace.h"
#include "core/fxcodec/jpx/cjpx_decoder.h"
#include "core/fxcrt/check_op.h"
#include "core/fxcrt/notreached.h"
#include "core/fxcrt/retain_ptr.h"
#include "core/fxge/dib/fx_dib.h"

namespace {

enum class JpxDecodeAction {
  kDoNothing,
  kUseGray,
  kUseIndexed,
  kUseRgb,
  kUseCmyk,
  kConvertArgbToRgb,
};

// ISO 32000-1:2008 section 7.4.9 says the PDF and JPX colorspaces should have
// the same number of color channels. This helper function checks the
// colorspaces match, but also tolerates unknowns.
bool IsJPXColorSpaceOrUnspecifiedOrUnknown(COLOR_SPACE actual,
                                           COLOR_SPACE expected) {
  return actual == expected || actual == OPJ_CLRSPC_UNSPECIFIED ||
         actual == OPJ_CLRSPC_UNKNOWN;
}

// Decides which JpxDecodeAction to use based on the colorspace information from
// the PDF and the JPX image. Called only when the PDF's image object contains a
// "/ColorSpace" entry.
std::optional<JpxDecodeAction> GetJpxDecodeActionFromColorSpaces(
    const CJPX_Decoder::JpxImageInfo& jpx_info,
    const CPDF_ColorSpace* pdf_colorspace) {
  if (pdf_colorspace ==
      CPDF_ColorSpace::GetStockCS(CPDF_ColorSpace::Family::kDeviceGray)) {
    if (!IsJPXColorSpaceOrUnspecifiedOrUnknown(/*actual=*/jpx_info.colorspace,
                                               /*expected=*/OPJ_CLRSPC_GRAY)) {
      return std::nullopt;
    }
    return JpxDecodeAction::kUseGray;
  }

  if (pdf_colorspace ==
      CPDF_ColorSpace::GetStockCS(CPDF_ColorSpace::Family::kDeviceRGB)) {
    if (!IsJPXColorSpaceOrUnspecifiedOrUnknown(/*actual=*/jpx_info.colorspace,
                                               /*expected=*/OPJ_CLRSPC_SRGB)) {
      return std::nullopt;
    }

    // The channel count of a JPX image can be different from the PDF color
    // space's component count.
    if (jpx_info.channels > 3) {
      return JpxDecodeAction::kConvertArgbToRgb;
    }
    return JpxDecodeAction::kUseRgb;
  }

  if (pdf_colorspace ==
      CPDF_ColorSpace::GetStockCS(CPDF_ColorSpace::Family::kDeviceCMYK)) {
    if (!IsJPXColorSpaceOrUnspecifiedOrUnknown(/*actual=*/jpx_info.colorspace,
                                               /*expected=*/OPJ_CLRSPC_CMYK)) {
      return std::nullopt;
    }
    return JpxDecodeAction::kUseCmyk;
  }

  // Many PDFs generated by iOS meets this condition. Handle the discrepancy.
  // See https://crbug.com/345431077 for example.
  if (pdf_colorspace->ComponentCount() == 3 && jpx_info.channels == 4 &&
      jpx_info.colorspace == OPJ_CLRSPC_SRGB) {
    return JpxDecodeAction::kConvertArgbToRgb;
  }

  if (pdf_colorspace->GetFamily() == CPDF_ColorSpace::Family::kIndexed &&
      pdf_colorspace->ComponentCount() == 1) {
    return JpxDecodeAction::kUseIndexed;
  }

  return JpxDecodeAction::kDoNothing;
}

JpxDecodeAction GetJpxDecodeActionFromImageColorSpace(
    const CJPX_Decoder::JpxImageInfo& jpx_info) {
  switch (jpx_info.colorspace) {
    case OPJ_CLRSPC_UNKNOWN:
    case OPJ_CLRSPC_UNSPECIFIED:
      return jpx_info.channels == 3 ? JpxDecodeAction::kUseRgb
                                    : JpxDecodeAction::kDoNothing;

    case OPJ_CLRSPC_SYCC:
    case OPJ_CLRSPC_EYCC:
      return JpxDecodeAction::kDoNothing;

    case OPJ_CLRSPC_SRGB:
      return jpx_info.channels > 3 ? JpxDecodeAction::kConvertArgbToRgb
                                   : JpxDecodeAction::kUseRgb;

    case OPJ_CLRSPC_GRAY:
      return JpxDecodeAction::kUseGray;

    case OPJ_CLRSPC_CMYK:
      return JpxDecodeAction::kUseCmyk;
  }
  NOTREACHED();
}

int GetComponentCountFromJpxImageInfo(
    const CJPX_Decoder::JpxImageInfo& jpx_info) {
  switch (jpx_info.colorspace) {
    case OPJ_CLRSPC_UNKNOWN:
    case OPJ_CLRSPC_UNSPECIFIED:
      return jpx_info.channels;

    case OPJ_CLRSPC_GRAY:
      return 1;

    case OPJ_CLRSPC_SRGB:
    case OPJ_CLRSPC_SYCC:
    case OPJ_CLRSPC_EYCC:
      return 3;

    case OPJ_CLRSPC_CMYK:
      return 4;
  }
  NOTREACHED();
}

std::optional<FXDIB_Format> GetFormatFromJpxDecodeActionAndImageInfo(
    JpxDecodeAction action,
    uint32_t channels) {
  if (action == JpxDecodeAction::kUseGray ||
      action == JpxDecodeAction::kUseIndexed) {
    return FXDIB_Format::k8bppRgb;
  }
  if (action == JpxDecodeAction::kUseRgb && channels == 3) {
    return FXDIB_Format::kBgr;
  }
  if (action == JpxDecodeAction::kUseRgb && channels == 4) {
    return FXDIB_Format::kBgrx;
  }
  if (action == JpxDecodeAction::kConvertArgbToRgb) {
    CHECK_GE(channels, 4);
    return FXDIB_Format::kBgrx;
  }
  return std::nullopt;
}

}  // namespace

// static
std::optional<JpxDecodeConversion> JpxDecodeConversion::Create(
    const CJPX_Decoder::JpxImageInfo& jpx_info,
    const CPDF_ColorSpace* pdf_colorspace) {
  // When the PDF does not provide a color space, check the image color space.
  std::optional<JpxDecodeAction> maybe_action =
      pdf_colorspace
          ? GetJpxDecodeActionFromColorSpaces(jpx_info, pdf_colorspace)
          : GetJpxDecodeActionFromImageColorSpace(jpx_info);
  if (!maybe_action.has_value()) {
    return std::nullopt;
  }

  JpxDecodeConversion conversion;
  switch (maybe_action.value()) {
    case JpxDecodeAction::kDoNothing:
      break;

    case JpxDecodeAction::kUseGray:
      conversion.override_colorspace_ =
          CPDF_ColorSpace::GetStockCS(CPDF_ColorSpace::Family::kDeviceGray);
      break;

    case JpxDecodeAction::kUseIndexed:
      break;

    case JpxDecodeAction::kUseRgb:
      DCHECK_GE(jpx_info.channels, 3);
      conversion.swap_rgb_ = true;
      conversion.override_colorspace_ = nullptr;
      break;

    case JpxDecodeAction::kUseCmyk:
      conversion.override_colorspace_ =
          CPDF_ColorSpace::GetStockCS(CPDF_ColorSpace::Family::kDeviceCMYK);
      break;

    case JpxDecodeAction::kConvertArgbToRgb:
      conversion.swap_rgb_ = true;
      conversion.convert_argb_to_rgb_ = true;
      conversion.override_colorspace_ = nullptr;
      break;
  }

  // If there exists a PDF colorspace, then CPDF_DIB already has the
  // components count.
  if (!pdf_colorspace) {
    conversion.jpx_components_count_ =
        GetComponentCountFromJpxImageInfo(jpx_info);
  }

  std::optional<FXDIB_Format> maybe_format =
      GetFormatFromJpxDecodeActionAndImageInfo(maybe_action.value(),
                                               jpx_info.channels);
  if (maybe_format.has_value()) {
    conversion.format_ = maybe_format.value();
    conversion.width_ = jpx_info.width;
  } else {
    conversion.format_ = FXDIB_Format::kBgr;
    conversion.width_ = (jpx_info.width * jpx_info.channels + 2) / 3;
  }

  return conversion;
}

JpxDecodeConversion::JpxDecodeConversion() = default;

JpxDecodeConversion::JpxDecodeConversion(JpxDecodeConversion&&) noexcept =
    default;

JpxDecodeConversion& JpxDecodeConversion::operator=(
    JpxDecodeConversion&&) noexcept = default;

JpxDecodeConversion::~JpxDecodeConversion() = default;