// Copyright (c) the JPEG XL Project Authors. All rights reserved.

//

// Use of this source code is governed by a BSD-style

// license that can be found in the LICENSE file.

#include "lib/jxl/blending.h"

#include <jxl/memory_manager.h>

#include <cstddef>

#include <cstring>

#include <vector>

#include "lib/jxl/alpha.h"

#include "lib/jxl/base/status.h"

#include "lib/jxl/dec_patch_dictionary.h"

#include "lib/jxl/frame_header.h"

#include "lib/jxl/image.h"

#include "lib/jxl/image_metadata.h"

namespace jxl {

bool NeedsBlending(const FrameHeader& frame_header) {

  if (!(frame_header.frame_type == FrameType::kRegularFrame ||

        frame_header.frame_type == FrameType::kSkipProgressive)) {

    return false;

  }

  const auto& info = frame_header.blending_info;

  bool replace_all = (info.mode == BlendMode::kReplace);

  for (const auto& ec_i : frame_header.extra_channel_blending_info) {

    if (ec_i.mode != BlendMode::kReplace) {

      replace_all = false;

    }

  }

  // Replace the full frame: nothing to do.

  if (!frame_header.custom_size_or_origin && replace_all) {

    return false;

  }

  return true;

}

Status PerformBlending(

    JxlMemoryManager* memory_manager, const float* const* bg,

    const float* const* fg, float* const* out, size_t x0, size_t xsize,

    const PatchBlending& color_blending, const PatchBlending* ec_blending,

    const std::vector<ExtraChannelInfo>& extra_channel_info) {

  bool has_alpha = false;

  size_t num_ec = extra_channel_info.size();

  for (size_t i = 0; i < num_ec; i++) {

    if (extra_channel_info[i].type == jxl::ExtraChannel::kAlpha) {

      has_alpha = true;

      break;

    }

  }

  JXL_ASSIGN_OR_RETURN(ImageF tmp,

                       ImageF::Create(memory_manager, xsize, 3 + num_ec));

  // Blend extra channels first so that we use the pre-blending alpha.

  for (size_t i = 0; i < num_ec; i++) {

    switch (ec_blending[i].mode) {

      case PatchBlendMode::kAdd:

        for (size_t x = 0; x < xsize; x++) {

          tmp.Row(3 + i)[x] = bg[3 + i][x + x0] + fg[3 + i][x + x0];

        }

        continue;

      case PatchBlendMode::kBlendAbove: {

        size_t alpha = ec_blending[i].alpha_channel;

        bool is_premultiplied = extra_channel_info[alpha].alpha_associated;

        PerformAlphaBlending(bg[3 + i] + x0, bg[3 + alpha] + x0, fg[3 + i] + x0,

                             fg[3 + alpha] + x0, tmp.Row(3 + i), xsize,

                             is_premultiplied, ec_blending[i].clamp);

        continue;

      }

      case PatchBlendMode::kBlendBelow: {

        size_t alpha = ec_blending[i].alpha_channel;

        bool is_premultiplied = extra_channel_info[alpha].alpha_associated;

        PerformAlphaBlending(fg[3 + i] + x0, fg[3 + alpha] + x0, bg[3 + i] + x0,

                             bg[3 + alpha] + x0, tmp.Row(3 + i), xsize,

                             is_premultiplied, ec_blending[i].clamp);

        continue;

      }

      case PatchBlendMode::kAlphaWeightedAddAbove: {

        size_t alpha = ec_blending[i].alpha_channel;

        PerformAlphaWeightedAdd(bg[3 + i] + x0, fg[3 + i] + x0,

                                fg[3 + alpha] + x0, tmp.Row(3 + i), xsize,

                                ec_blending[i].clamp);

        continue;

      }

      case PatchBlendMode::kAlphaWeightedAddBelow: {

        size_t alpha = ec_blending[i].alpha_channel;

        PerformAlphaWeightedAdd(fg[3 + i] + x0, bg[3 + i] + x0,

                                bg[3 + alpha] + x0, tmp.Row(3 + i), xsize,

                                ec_blending[i].clamp);

        continue;

      }

      case PatchBlendMode::kMul:

        PerformMulBlending(bg[3 + i] + x0, fg[3 + i] + x0, tmp.Row(3 + i),

                           xsize, ec_blending[i].clamp);

        continue;

      case PatchBlendMode::kReplace:

        if (xsize) memcpy(tmp.Row(3 + i), fg[3 + i] + x0, xsize * sizeof(**fg));

        continue;

      case PatchBlendMode::kNone:

        if (xsize) memcpy(tmp.Row(3 + i), bg[3 + i] + x0, xsize * sizeof(**fg));

        continue;

    }

  }

  size_t alpha = color_blending.alpha_channel;

  const auto add = [&]() {

    for (int p = 0; p < 3; p++) {

      float* out = tmp.Row(p);

      for (size_t x = 0; x < xsize; x++) {

        out[x] = bg[p][x + x0] + fg[p][x + x0];

      }

    }

  };

  const auto blend_weighted = [&](const float* const* bottom,

                                  const float* const* top) {

    bool is_premultiplied = extra_channel_info[alpha].alpha_associated;

    PerformAlphaBlending(

        {bottom[0] + x0, bottom[1] + x0, bottom[2] + x0,

         bottom[3 + alpha] + x0},

        {top[0] + x0, top[1] + x0, top[2] + x0, top[3 + alpha] + x0},

        {tmp.Row(0), tmp.Row(1), tmp.Row(2), tmp.Row(3 + alpha)}, xsize,

        is_premultiplied, color_blending.clamp);

  };

  const auto add_weighted = [&](const float* const* bottom,

                                const float* const* top) {

    for (size_t c = 0; c < 3; c++) {

      PerformAlphaWeightedAdd(bottom[c] + x0, top[c] + x0, top[3 + alpha] + x0,

                              tmp.Row(c), xsize, color_blending.clamp);

    }

  };

  const auto copy = [&](const float* const* src) {

    for (size_t p = 0; p < 3; p++) {

      memcpy(tmp.Row(p), src[p] + x0, xsize * sizeof(**src));

    }

  };

  switch (color_blending.mode) {

    case PatchBlendMode::kAdd:

      add();

      break;

    case PatchBlendMode::kAlphaWeightedAddAbove:

      has_alpha ? add_weighted(bg, fg) : add();

      break;

    case PatchBlendMode::kAlphaWeightedAddBelow:

      has_alpha ? add_weighted(fg, bg) : add();

      break;

    case PatchBlendMode::kBlendAbove:

      has_alpha ? blend_weighted(bg, fg) : copy(fg);

      break;

    case PatchBlendMode::kBlendBelow:

      has_alpha ? blend_weighted(fg, bg) : copy(fg);

      break;

    case PatchBlendMode::kMul:

      for (int p = 0; p < 3; p++) {

        PerformMulBlending(bg[p] + x0, fg[p] + x0, tmp.Row(p), xsize,

                           color_blending.clamp);

      }

      break;

    case PatchBlendMode::kReplace:

      copy(fg);

      break;

    case PatchBlendMode::kNone:

      copy(bg);

  }

  for (size_t i = 0; i < 3 + num_ec; i++) {

    if (xsize != 0) memcpy(out[i] + x0, tmp.Row(i), xsize * sizeof(**out));

  }

  return true;

}

}  // namespace jxl