// Copyright 2015 The Cobalt Authors. All Rights Reserved.
//
// 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
//
//     http://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 "cobalt/renderer/rasterizer/skia/skia/src/effects/SkYUV2RGBShader.h"

#include <algorithm>

#include "base/logging.h"
#include "cobalt/base/polymorphic_downcast.h"
#include "third_party/skia/include/core/SkColorPriv.h"
#include "third_party/skia/include/core/SkShader.h"
#include "third_party/skia/include/core/SkString.h"
#include "third_party/skia/src/core/SkReadBuffer.h"
#include "third_party/skia/src/core/SkWriteBuffer.h"
#include "third_party/skia/src/gpu/GrFragmentProcessor.h"
#include "third_party/skia/src/gpu/effects/GrYUVtoRGBEffect.h"
#include "third_party/skia/src/shaders/SkImageShader.h"

SkYUV2RGBShader::SkYUV2RGBShader(SkYUVColorSpace color_space,
                                 const sk_sp<SkImage>& y_image,
                                 const SkMatrix& y_matrix,
                                 const sk_sp<SkImage>& u_image,
                                 const SkMatrix& u_matrix,
                                 const sk_sp<SkImage>& v_image,
                                 const SkMatrix& v_matrix)
    : color_space_(color_space),
      y_image_(y_image),
      y_matrix_(y_matrix),
      u_image_(u_image),
      u_matrix_(u_matrix),
      v_image_(v_image),
      v_matrix_(v_matrix) {
  DCHECK(y_image_);
  DCHECK(u_image_);
  DCHECK(v_image_);
  InitializeShaders();
}

void SkYUV2RGBShader::InitializeShaders() {
  y_shader_.reset(base::polymorphic_downcast<SkImageShader*>(
      SkImageShader::Make(y_image_, SkTileMode::kClamp, SkTileMode::kClamp,
                          &y_matrix_, false)
          .get()));
  DCHECK(y_shader_);

  u_shader_.reset(base::polymorphic_downcast<SkImageShader*>(
      SkImageShader::Make(u_image_, SkTileMode::kClamp, SkTileMode::kClamp,
                          &u_matrix_, false)
          .get()));
  DCHECK(u_shader_);

  v_shader_.reset(base::polymorphic_downcast<SkImageShader*>(
      SkImageShader::Make(v_image_, SkTileMode::kClamp, SkTileMode::kClamp,
                          &v_matrix_, false)
          .get()));
  DCHECK(v_shader_);
}

void SkYUV2RGBShader::flatten(SkWriteBuffer& buffer) const {
  buffer.writeInt(color_space_);
  buffer.writeImage(y_image_.get());
  buffer.writeMatrix(y_matrix_);
  buffer.writeImage(u_image_.get());
  buffer.writeMatrix(u_matrix_);
  buffer.writeImage(v_image_.get());
  buffer.writeMatrix(v_matrix_);
}

uint32_t SkYUV2RGBShader::YUV2RGBShaderContext::getFlags() const {
  return 0;
}

SkShaderBase::Context* SkYUV2RGBShader::onMakeContext(
    const ContextRec& rec, SkArenaAlloc* storage) const {
  return nullptr;
}
SkYUV2RGBShader::YUV2RGBShaderContext::YUV2RGBShaderContext(
    SkYUVColorSpace color_space, const SkYUV2RGBShader& yuv2rgb_shader,
    SkShaderBase::Context* y_shader_context,
    SkShaderBase::Context* u_shader_context,
    SkShaderBase::Context* v_shader_context, const ContextRec& rec)
    : INHERITED(yuv2rgb_shader, rec),
      color_space_(color_space),
      y_shader_context_(y_shader_context),
      u_shader_context_(u_shader_context),
      v_shader_context_(v_shader_context) {}

SkYUV2RGBShader::YUV2RGBShaderContext::~YUV2RGBShaderContext() {
  y_shader_context_->~Context();
  u_shader_context_->~Context();
  v_shader_context_->~Context();
}
void SkYUV2RGBShader::YUV2RGBShaderContext::shadeSpan(
    int x, int y, SkPMColor result[], int count) {
  static const int kPixelCountPerSpan = 64;
  SkPMColor y_values[kPixelCountPerSpan];
  SkPMColor u_values[kPixelCountPerSpan];
  SkPMColor v_values[kPixelCountPerSpan];

  DCHECK_EQ(kRec709_SkYUVColorSpace, color_space_) <<
      "Currently we only support the BT.709 YUV colorspace.";

  do {
    int count_in_chunk =
        count > kPixelCountPerSpan ? kPixelCountPerSpan : count;

    y_shader_context_->shadeSpan(x, y, y_values, count_in_chunk);
    u_shader_context_->shadeSpan(x, y, u_values, count_in_chunk);
    v_shader_context_->shadeSpan(x, y, v_values, count_in_chunk);

    for (int i = 0; i < count_in_chunk; ++i) {
      int32_t y_value = SkColorGetA(y_values[i]) - 16;
      int32_t u_value = SkColorGetA(u_values[i]) - 128;
      int32_t v_value = SkColorGetA(v_values[i]) - 128;

      const float kA = 1.164f;
      const float kB = -0.213f;
      const float kC = 2.112f;
      const float kD = 1.793f;
      const float kE = -0.533f;
      int32_t r_unclamped = static_cast<int32_t>(kA * y_value + kD * v_value);
      int32_t g_unclamped =
          static_cast<int32_t>(kA * y_value + kB * u_value + kE * v_value);
      int32_t b_unclamped = static_cast<int32_t>(kA * y_value + kC * u_value);

      int32_t r_clamped =
          std::min<int32_t>(255, std::max<int32_t>(0, r_unclamped));
      int32_t g_clamped =
          std::min<int32_t>(255, std::max<int32_t>(0, g_unclamped));
      int32_t b_clamped =
          std::min<int32_t>(255, std::max<int32_t>(0, b_unclamped));

      result[i] = SkPackARGB32NoCheck(255, r_clamped, g_clamped, b_clamped);
    }
    result += count_in_chunk;
    x += count_in_chunk;
    count -= count_in_chunk;
  } while (count > 0);
}

#if SK_SUPPORT_GPU

std::unique_ptr<GrFragmentProcessor> SkYUV2RGBShader::asFragmentProcessor(
    const GrFPArgs&) const {
  return nullptr;
}

#endif  // SK_SUPPORT_GPU