// Copyright 2023 The MediaPipe Authors.
//
// 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 "mediapipe/util/frame_buffer/frame_buffer_util.h"

#include <cstdint>
#include <memory>
#include <vector>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "mediapipe/framework/formats/frame_buffer.h"
#include "mediapipe/framework/port/status_macros.h"
#include "mediapipe/util/frame_buffer/gray_buffer.h"
#include "mediapipe/util/frame_buffer/rgb_buffer.h"
#include "mediapipe/util/frame_buffer/yuv_buffer.h"

namespace mediapipe {
namespace frame_buffer {

namespace {

constexpr int kRgbaChannels = 4;
constexpr int kRgbaPixelBytes = 4;
constexpr int kRgbChannels = 3;
constexpr int kRgbPixelBytes = 3;
constexpr int kGrayChannel = 1;
constexpr int kGrayPixelBytes = 1;

// YUV helpers.
//------------------------------------------------------------------------------

// Returns whether the buffer is part of the supported Yuv format.
bool IsSupportedYuvBuffer(const FrameBuffer& buffer) {
  return buffer.format() == FrameBuffer::Format::kNV21 ||
         buffer.format() == FrameBuffer::Format::kNV12 ||
         buffer.format() == FrameBuffer::Format::kYV12 ||
         buffer.format() == FrameBuffer::Format::kYV21;
}

// Shared validation functions.
//------------------------------------------------------------------------------

// Indicates whether the given buffers have the same dimensions.
bool AreBufferDimsEqual(const FrameBuffer& buffer1,
                        const FrameBuffer& buffer2) {
  return buffer1.dimension() == buffer2.dimension();
}

// Indicates whether the given buffers formats are compatible. Same formats are
// compatible and all YUV family formats (e.g. NV21, NV12, YV12, YV21, etc) are
// compatible.
bool AreBufferFormatsCompatible(const FrameBuffer& buffer1,
                                const FrameBuffer& buffer2) {
  switch (buffer1.format()) {
    case FrameBuffer::Format::kRGBA:
    case FrameBuffer::Format::kRGB:
      return (buffer2.format() == FrameBuffer::Format::kRGBA ||
              buffer2.format() == FrameBuffer::Format::kRGB);
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return (buffer2.format() == FrameBuffer::Format::kNV12 ||
              buffer2.format() == FrameBuffer::Format::kNV21 ||
              buffer2.format() == FrameBuffer::Format::kYV12 ||
              buffer2.format() == FrameBuffer::Format::kYV21);
    case FrameBuffer::Format::kGRAY:
    default:
      return buffer1.format() == buffer2.format();
  }
}

absl::Status ValidateBufferFormat(const FrameBuffer& buffer) {
  switch (buffer.format()) {
    case FrameBuffer::Format::kGRAY:
    case FrameBuffer::Format::kRGB:
    case FrameBuffer::Format::kRGBA:
      if (buffer.plane_count() == 1) return absl::OkStatus();
      return absl::InvalidArgumentError(
          "Plane count must be 1 for grayscale and RGB[a] buffers.");
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kYV21:
    case FrameBuffer::Format::kYV12:
      return absl::OkStatus();
    default:
      return absl::InternalError(
          absl::StrFormat("Unsupported buffer format: %i.", buffer.format()));
  }
}

absl::Status ValidateBufferFormats(const FrameBuffer& buffer1,
                                   const FrameBuffer& buffer2) {
  MP_RETURN_IF_ERROR(ValidateBufferFormat(buffer1));
  MP_RETURN_IF_ERROR(ValidateBufferFormat(buffer2));
  return absl::OkStatus();
}

absl::Status ValidateResizeBufferInputs(const FrameBuffer& buffer,
                                        const FrameBuffer& output_buffer) {
  bool valid_format = false;
  switch (buffer.format()) {
    case FrameBuffer::Format::kGRAY:
    case FrameBuffer::Format::kRGB:
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      valid_format = (buffer.format() == output_buffer.format());
      break;
    case FrameBuffer::Format::kRGBA:
      valid_format = (output_buffer.format() == FrameBuffer::Format::kRGBA ||
                      output_buffer.format() == FrameBuffer::Format::kRGB);
      break;
    default:
      return absl::InternalError(
          absl::StrFormat("Unsupported buffer format: %i.", buffer.format()));
  }
  if (!valid_format) {
    return absl::InvalidArgumentError(
        "Input and output buffer formats must match.");
  }
  return ValidateBufferFormats(buffer, output_buffer);
}

absl::Status ValidateRotateBufferInputs(const FrameBuffer& buffer,
                                        const FrameBuffer& output_buffer,
                                        int angle_deg) {
  if (!AreBufferFormatsCompatible(buffer, output_buffer)) {
    return absl::InvalidArgumentError(
        "Input and output buffer formats must match.");
  }

  const bool is_dimension_change = (angle_deg / 90) % 2 == 1;
  const bool are_dimensions_rotated =
      (buffer.dimension().width == output_buffer.dimension().height) &&
      (buffer.dimension().height == output_buffer.dimension().width);
  const bool are_dimensions_equal =
      buffer.dimension() == output_buffer.dimension();

  if (angle_deg >= 360 || angle_deg <= 0 || angle_deg % 90 != 0) {
    return absl::InvalidArgumentError(
        "Rotation angle must be between 0 and 360, in multiples of 90 "
        "degrees.");
  } else if ((is_dimension_change && !are_dimensions_rotated) ||
             (!is_dimension_change && !are_dimensions_equal)) {
    return absl::InvalidArgumentError(
        "Output buffer has invalid dimensions for rotation.");
  }
  return absl::OkStatus();
}

absl::Status ValidateCropBufferInputs(const FrameBuffer& buffer,
                                      const FrameBuffer& output_buffer, int x0,
                                      int y0, int x1, int y1) {
  if (!AreBufferFormatsCompatible(buffer, output_buffer)) {
    return absl::InvalidArgumentError(
        "Input and output buffer formats must match.");
  }

  bool is_buffer_size_valid =
      ((x1 < buffer.dimension().width) && y1 < buffer.dimension().height);
  bool are_points_valid = (x0 >= 0) && (y0 >= 0) && (x1 >= x0) && (y1 >= y0);

  if (!is_buffer_size_valid || !are_points_valid) {
    return absl::InvalidArgumentError("Invalid crop coordinates.");
  }
  return absl::OkStatus();
}

absl::Status ValidateFlipBufferInputs(const FrameBuffer& buffer,
                                      const FrameBuffer& output_buffer) {
  if (!AreBufferFormatsCompatible(buffer, output_buffer)) {
    return absl::InvalidArgumentError(
        "Input and output buffer formats must match.");
  }
  return AreBufferDimsEqual(buffer, output_buffer)
             ? absl::OkStatus()
             : absl::InvalidArgumentError(
                   "Input and output buffers must have the same dimensions.");
}

absl::Status ValidateConvertFormats(FrameBuffer::Format from_format,
                                    FrameBuffer::Format to_format) {
  if (from_format == to_format) {
    return absl::InvalidArgumentError("Formats must be different.");
  }

  switch (from_format) {
    case FrameBuffer::Format::kGRAY:
      return absl::InvalidArgumentError(
          "Grayscale format does not convert to other formats.");
    case FrameBuffer::Format::kRGB:
    case FrameBuffer::Format::kRGBA:
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return absl::OkStatus();
    default:
      return absl::InternalError(
          absl::StrFormat("Unsupported buffer format: %i.", from_format));
  }
}

// Construct buffer helper functions.
//------------------------------------------------------------------------------

// Creates NV12 / NV21 / YV12 / YV21 YuvBuffer from the input `buffer`. The
// output YuvBuffer is agnostic to the YUV format since the YUV buffers are
// managed individually.
absl::StatusOr<YuvBuffer> CreateYuvBuffer(const FrameBuffer& buffer) {
  ASSIGN_OR_RETURN(FrameBuffer::YuvData yuv_data,
                   FrameBuffer::GetYuvDataFromFrameBuffer(buffer));
  return YuvBuffer(const_cast<uint8_t*>(yuv_data.y_buffer),
                   const_cast<uint8_t*>(yuv_data.u_buffer),
                   const_cast<uint8_t*>(yuv_data.v_buffer),
                   buffer.dimension().width, buffer.dimension().height,
                   yuv_data.y_row_stride, yuv_data.uv_row_stride,
                   yuv_data.uv_pixel_stride);
}

absl::StatusOr<GrayBuffer> CreateGrayBuffer(const FrameBuffer& buffer) {
  if (buffer.plane_count() != 1) {
    return absl::InternalError("Unsupported grayscale planar format.");
  }
  return GrayBuffer(const_cast<uint8_t*>(buffer.plane(0).buffer()),
                    buffer.dimension().width, buffer.dimension().height);
}

absl::StatusOr<RgbBuffer> CreateRgbBuffer(const FrameBuffer& buffer) {
  if (buffer.plane_count() != 1) {
    return absl::InternalError("Unsupported rgb[a] planar format.");
  }
  bool alpha = buffer.format() == FrameBuffer::Format::kRGBA ? true : false;
  return RgbBuffer(const_cast<uint8_t*>(buffer.plane(0).buffer()),
                   buffer.dimension().width, buffer.dimension().height,
                   buffer.plane(0).stride().row_stride_bytes, alpha);
}

// Grayscale transformation functions.
//------------------------------------------------------------------------------

absl::Status CropGrayscale(const FrameBuffer& buffer, int x0, int y0, int x1,
                           int y1, FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer));
  bool success_crop = input.Crop(x0, y0, x1, y1);
  if (!success_crop) {
    return absl::UnknownError("Halide grayscale crop operation failed.");
  }
  bool success_resize = input.Resize(&output);
  if (!success_resize) {
    return absl::UnknownError("Halide grayscale resize operation failed.");
  }
  return absl::OkStatus();
}

absl::Status ResizeGrayscale(const FrameBuffer& buffer,
                             FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer));
  return input.Resize(&output)
             ? absl::OkStatus()
             : absl::UnknownError("Halide grayscale resize operation failed.");
}

absl::Status RotateGrayscale(const FrameBuffer& buffer, int angle_deg,
                             FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer));
  return input.Rotate(angle_deg % 360, &output)
             ? absl::OkStatus()
             : absl::UnknownError("Halide grayscale rotate operation failed.");
}

absl::Status FlipHorizontallyGrayscale(const FrameBuffer& buffer,
                                       FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer));
  return input.FlipHorizontally(&output)
             ? absl::OkStatus()
             : absl::UnknownError(
                   "Halide grayscale horizontal flip operation failed.");
}

absl::Status FlipVerticallyGrayscale(const FrameBuffer& buffer,
                                     FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer));
  return input.FlipVertically(&output)
             ? absl::OkStatus()
             : absl::UnknownError(
                   "Halide grayscale vertical flip operation failed.");
}

// Rgb transformation functions.
//------------------------------------------------------------------------------

absl::Status ResizeRgb(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer));
  return input.Resize(&output)
             ? absl::OkStatus()
             : absl::UnknownError("Halide rgb[a] resize operation failed.");
}

absl::Status ConvertRgb(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer));
  bool result = false;
  if (output_buffer->format() == FrameBuffer::Format::kGRAY) {
    ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer));
    result = input.Convert(&output);
  } else if (IsSupportedYuvBuffer(*output_buffer)) {
    ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer));
    result = input.Convert(&output);
  } else if (output_buffer->format() == FrameBuffer::Format::kRGBA ||
             output_buffer->format() == FrameBuffer::Format::kRGB) {
    ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer));
    result = input.Convert(&output);
  }
  return result ? absl::OkStatus()
                : absl::UnknownError("Halide rgb[a] convert operation failed.");
}

absl::Status CropRgb(const FrameBuffer& buffer, int x0, int y0, int x1, int y1,
                     FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer));
  bool success_crop = input.Crop(x0, y0, x1, y1);
  if (!success_crop) {
    return absl::UnknownError("Halide rgb[a] crop operation failed.");
  }
  bool success_resize = input.Resize(&output);
  if (!success_resize) {
    return absl::UnknownError("Halide rgb resize operation failed.");
  }
  return absl::OkStatus();
}

absl::Status FlipHorizontallyRgb(const FrameBuffer& buffer,
                                 FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer));
  return input.FlipHorizontally(&output)
             ? absl::OkStatus()
             : absl::UnknownError(
                   "Halide rgb[a] horizontal flip operation failed.");
}

absl::Status FlipVerticallyRgb(const FrameBuffer& buffer,
                               FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer));
  return input.FlipVertically(&output)
             ? absl::OkStatus()
             : absl::UnknownError(
                   "Halide rgb[a] vertical flip operation failed.");
}

absl::Status RotateRgb(const FrameBuffer& buffer, int angle,
                       FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer));
  return input.Rotate(angle % 360, &output)
             ? absl::OkStatus()
             : absl::UnknownError("Halide rgb[a] rotate operation failed.");
}

// Yuv transformation functions.
//------------------------------------------------------------------------------

absl::Status CropYuv(const FrameBuffer& buffer, int x0, int y0, int x1, int y1,
                     FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer));
  bool success_crop = input.Crop(x0, y0, x1, y1);
  if (!success_crop) {
    return absl::UnknownError("Halide YUV crop operation failed.");
  }
  bool success_resize = input.Resize(&output);
  if (!success_resize) {
    return absl::UnknownError("Halide YUV resize operation failed.");
  }
  return absl::OkStatus();
}

absl::Status ResizeYuv(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer));
  return input.Resize(&output)
             ? absl::OkStatus()
             : absl::UnknownError("Halide YUV resize operation failed.");
}

absl::Status RotateYuv(const FrameBuffer& buffer, int angle_deg,
                       FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer));
  return input.Rotate(angle_deg % 360, &output)
             ? absl::OkStatus()
             : absl::UnknownError("Halide YUV rotate operation failed.");
}

absl::Status FlipHorizontallyYuv(const FrameBuffer& buffer,
                                 FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer));
  return input.FlipHorizontally(&output)
             ? absl::OkStatus()
             : absl::UnknownError(
                   "Halide YUV horizontal flip operation failed.");
}

absl::Status FlipVerticallyYuv(const FrameBuffer& buffer,
                               FrameBuffer* output_buffer) {
  ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer));
  ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer));
  return input.FlipVertically(&output)
             ? absl::OkStatus()
             : absl::UnknownError("Halide YUV vertical flip operation failed.");
}

// Converts input YUV `buffer` into the `output_buffer` in RGB, RGBA or gray
// scale format.
absl::Status ConvertYuv(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
  bool success_convert = false;
  ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer));
  if (output_buffer->format() == FrameBuffer::Format::kRGBA ||
      output_buffer->format() == FrameBuffer::Format::kRGB) {
    ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer));
    bool half_sampling = false;
    if (buffer.dimension().width / 2 == output_buffer->dimension().width &&
        buffer.dimension().height / 2 == output_buffer->dimension().height) {
      half_sampling = true;
    }
    success_convert = input.Convert(half_sampling, &output);
  } else if (output_buffer->format() == FrameBuffer::Format::kGRAY) {
    if (buffer.plane(0).stride().row_stride_bytes == buffer.dimension().width) {
      std::copy(input.y_buffer()->host,
                input.y_buffer()->host + buffer.dimension().Size(),
                const_cast<uint8_t*>(output_buffer->plane(0).buffer()));
    } else {
      // The y_buffer is padded. The conversion removes the padding.
      uint8_t* gray_buffer =
          const_cast<uint8_t*>(output_buffer->plane(0).buffer());
      for (int i = 0; i < buffer.dimension().height; i++) {
        int src_address = i * buffer.plane(0).stride().row_stride_bytes;
        int dest_address = i * buffer.dimension().width;
        std::memcpy(&gray_buffer[dest_address],
                    &buffer.plane(0).buffer()[src_address],
                    buffer.dimension().width);
      }
    }
    success_convert = true;
  } else if (IsSupportedYuvBuffer(*output_buffer)) {
    ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer));
    success_convert = input.Resize(&output);
  }
  return success_convert
             ? absl::OkStatus()
             : absl::UnknownError("Halide YUV convert operation failed.");
}

}  // namespace

// Public methods.
//------------------------------------------------------------------------------

std::shared_ptr<FrameBuffer> CreateFromRgbaRawBuffer(
    uint8_t* input, FrameBuffer::Dimension dimension,
    FrameBuffer::Stride stride) {
  if (stride == kDefaultStride) {
    stride.row_stride_bytes = dimension.width * kRgbaChannels;
    stride.pixel_stride_bytes = kRgbaChannels;
  }
  FrameBuffer::Plane input_plane(/*buffer=*/input,
                                 /*stride=*/stride);
  std::vector<FrameBuffer::Plane> planes{input_plane};
  return std::make_shared<FrameBuffer>(planes, dimension,
                                       FrameBuffer::Format::kRGBA);
}

std::shared_ptr<FrameBuffer> CreateFromRgbRawBuffer(
    uint8_t* input, FrameBuffer::Dimension dimension,
    FrameBuffer::Stride stride) {
  if (stride == kDefaultStride) {
    stride.row_stride_bytes = dimension.width * kRgbChannels;
    stride.pixel_stride_bytes = kRgbChannels;
  }
  FrameBuffer::Plane input_plane(/*buffer=*/input,
                                 /*stride=*/stride);
  std::vector<FrameBuffer::Plane> planes{input_plane};
  return std::make_shared<FrameBuffer>(planes, dimension,
                                       FrameBuffer::Format::kRGB);
}

std::shared_ptr<FrameBuffer> CreateFromGrayRawBuffer(
    uint8_t* input, FrameBuffer::Dimension dimension,
    FrameBuffer::Stride stride) {
  if (stride == kDefaultStride) {
    stride.row_stride_bytes = dimension.width * kGrayChannel;
    stride.pixel_stride_bytes = kGrayChannel;
  }
  FrameBuffer::Plane input_plane(/*buffer=*/input,
                                 /*stride=*/stride);
  std::vector<FrameBuffer::Plane> planes{input_plane};
  return std::make_shared<FrameBuffer>(planes, dimension,
                                       FrameBuffer::Format::kGRAY);
}

absl::StatusOr<std::shared_ptr<FrameBuffer>> CreateFromYuvRawBuffer(
    uint8_t* y_plane, uint8_t* u_plane, uint8_t* v_plane,
    FrameBuffer::Format format, FrameBuffer::Dimension dimension,
    int row_stride_y, int row_stride_uv, int pixel_stride_uv) {
  const int pixel_stride_y = 1;
  std::vector<FrameBuffer::Plane> planes;
  if (format == FrameBuffer::Format::kNV21 ||
      format == FrameBuffer::Format::kYV12) {
    planes = {{y_plane, /*stride=*/{row_stride_y, pixel_stride_y}},
              {v_plane, /*stride=*/{row_stride_uv, pixel_stride_uv}},
              {u_plane, /*stride=*/{row_stride_uv, pixel_stride_uv}}};
  } else if (format == FrameBuffer::Format::kNV12 ||
             format == FrameBuffer::Format::kYV21) {
    planes = {{y_plane, /*stride=*/{row_stride_y, pixel_stride_y}},
              {u_plane, /*stride=*/{row_stride_uv, pixel_stride_uv}},
              {v_plane, /*stride=*/{row_stride_uv, pixel_stride_uv}}};
  } else {
    return absl::InvalidArgumentError(
        absl::StrFormat("Input format is not YUV-like: %i.", format));
  }
  return std::make_shared<FrameBuffer>(planes, dimension, format);
}

absl::StatusOr<std::shared_ptr<FrameBuffer>> CreateFromRawBuffer(
    uint8_t* buffer, FrameBuffer::Dimension dimension,
    const FrameBuffer::Format target_format) {
  switch (target_format) {
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21: {
      FrameBuffer::Plane plane(/*buffer=*/buffer,
                               /*stride=*/{dimension.width, kGrayChannel});
      std::vector<FrameBuffer::Plane> planes{plane};
      return std::make_shared<FrameBuffer>(planes, dimension, target_format);
    }
    case FrameBuffer::Format::kYV12: {
      ASSIGN_OR_RETURN(const FrameBuffer::Dimension uv_dimension,
                       GetUvPlaneDimension(dimension, target_format));
      return CreateFromYuvRawBuffer(
          /*y_plane=*/buffer,
          /*u_plane=*/buffer + dimension.Size() + uv_dimension.Size(),
          /*v_plane=*/buffer + dimension.Size(), target_format, dimension,
          /*row_stride_y=*/dimension.width, uv_dimension.width,
          /*pixel_stride_uv=*/1);
    }
    case FrameBuffer::Format::kYV21: {
      ASSIGN_OR_RETURN(const FrameBuffer::Dimension uv_dimension,
                       GetUvPlaneDimension(dimension, target_format));
      return CreateFromYuvRawBuffer(
          /*y_plane=*/buffer, /*u_plane=*/buffer + dimension.Size(),
          /*v_plane=*/buffer + dimension.Size() + uv_dimension.Size(),
          target_format, dimension, /*row_stride_y=*/dimension.width,
          uv_dimension.width,
          /*pixel_stride_uv=*/1);
    }
    case FrameBuffer::Format::kRGBA:
      return CreateFromRgbaRawBuffer(buffer, dimension);
    case FrameBuffer::Format::kRGB:
      return CreateFromRgbRawBuffer(buffer, dimension);
    case FrameBuffer::Format::kGRAY:
      return CreateFromGrayRawBuffer(buffer, dimension);
    default:
      return absl::InternalError(
          absl::StrFormat("Unsupported buffer format: %i.", target_format));
  }
}

absl::Status Crop(const FrameBuffer& buffer, int x0, int y0, int x1, int y1,
                  FrameBuffer* output_buffer) {
  MP_RETURN_IF_ERROR(
      ValidateCropBufferInputs(buffer, *output_buffer, x0, y0, x1, y1));
  MP_RETURN_IF_ERROR(ValidateBufferFormats(buffer, *output_buffer));

  switch (buffer.format()) {
    case FrameBuffer::Format::kGRAY:
      return CropGrayscale(buffer, x0, y0, x1, y1, output_buffer);
    case FrameBuffer::Format::kRGBA:
    case FrameBuffer::Format::kRGB:
      return CropRgb(buffer, x0, y0, x1, y1, output_buffer);
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return CropYuv(buffer, x0, y0, x1, y1, output_buffer);
    default:
      return absl::InternalError(
          absl::StrFormat("Format %i is not supported.", buffer.format()));
  }
}

absl::Status Resize(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
  MP_RETURN_IF_ERROR(ValidateResizeBufferInputs(buffer, *output_buffer));

  switch (buffer.format()) {
    case FrameBuffer::Format::kGRAY:
      return ResizeGrayscale(buffer, output_buffer);
    case FrameBuffer::Format::kRGBA:
    case FrameBuffer::Format::kRGB:
      return ResizeRgb(buffer, output_buffer);
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return ResizeYuv(buffer, output_buffer);
    default:
      return absl::InternalError(
          absl::StrFormat("Format %i is not supported.", buffer.format()));
  }
}

absl::Status Rotate(const FrameBuffer& buffer, int angle_deg,
                    FrameBuffer* output_buffer) {
  MP_RETURN_IF_ERROR(
      ValidateRotateBufferInputs(buffer, *output_buffer, angle_deg));
  MP_RETURN_IF_ERROR(ValidateBufferFormats(buffer, *output_buffer));

  switch (buffer.format()) {
    case FrameBuffer::Format::kGRAY:
      return RotateGrayscale(buffer, angle_deg, output_buffer);
    case FrameBuffer::Format::kRGBA:
    case FrameBuffer::Format::kRGB:
      return RotateRgb(buffer, angle_deg, output_buffer);
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return RotateYuv(buffer, angle_deg, output_buffer);
    default:
      return absl::InternalError(
          absl::StrFormat("Format %i is not supported.", buffer.format()));
  }
}

absl::Status FlipHorizontally(const FrameBuffer& buffer,
                              FrameBuffer* output_buffer) {
  MP_RETURN_IF_ERROR(ValidateFlipBufferInputs(buffer, *output_buffer));
  MP_RETURN_IF_ERROR(ValidateBufferFormats(buffer, *output_buffer));

  switch (buffer.format()) {
    case FrameBuffer::Format::kGRAY:
      return FlipHorizontallyGrayscale(buffer, output_buffer);
    case FrameBuffer::Format::kRGBA:
    case FrameBuffer::Format::kRGB:
      return FlipHorizontallyRgb(buffer, output_buffer);
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return FlipHorizontallyYuv(buffer, output_buffer);
    default:
      return absl::InternalError(
          absl::StrFormat("Format %i is not supported.", buffer.format()));
  }
}

absl::Status FlipVertically(const FrameBuffer& buffer,
                            FrameBuffer* output_buffer) {
  MP_RETURN_IF_ERROR(ValidateFlipBufferInputs(buffer, *output_buffer));
  MP_RETURN_IF_ERROR(ValidateBufferFormats(buffer, *output_buffer));

  switch (buffer.format()) {
    case FrameBuffer::Format::kGRAY:
      return FlipVerticallyGrayscale(buffer, output_buffer);
    case FrameBuffer::Format::kRGBA:
    case FrameBuffer::Format::kRGB:
      return FlipVerticallyRgb(buffer, output_buffer);
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return FlipVerticallyYuv(buffer, output_buffer);
    default:
      return absl::InternalError(
          absl::StrFormat("Format %i is not supported.", buffer.format()));
  }
}

absl::Status Convert(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
  MP_RETURN_IF_ERROR(
      ValidateConvertFormats(buffer.format(), output_buffer->format()));

  switch (buffer.format()) {
    case FrameBuffer::Format::kRGBA:
    case FrameBuffer::Format::kRGB:
      return ConvertRgb(buffer, output_buffer);
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return ConvertYuv(buffer, output_buffer);
    default:
      return absl::InternalError(
          absl::StrFormat("Format %i is not supported.", buffer.format()));
  }
}

int GetFrameBufferByteSize(FrameBuffer::Dimension dimension,
                           FrameBuffer::Format format) {
  switch (format) {
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return /*y plane*/ dimension.Size() +
             /*uv plane*/ (dimension.width + 1) / 2 * (dimension.height + 1) /
                 2 * 2;
    case FrameBuffer::Format::kRGB:
      return dimension.Size() * kRgbPixelBytes;
    case FrameBuffer::Format::kRGBA:
      return dimension.Size() * kRgbaPixelBytes;
    case FrameBuffer::Format::kGRAY:
      return dimension.Size();
    default:
      return 0;
  }
}

absl::StatusOr<int> GetPixelStrides(FrameBuffer::Format format) {
  switch (format) {
    case FrameBuffer::Format::kGRAY:
      return kGrayPixelBytes;
    case FrameBuffer::Format::kRGB:
      return kRgbPixelBytes;
    case FrameBuffer::Format::kRGBA:
      return kRgbaPixelBytes;
    default:
      return absl::InvalidArgumentError(absl::StrFormat(
          "GetPixelStrides does not support format: %i.", format));
  }
}

absl::StatusOr<const uint8_t*> GetUvRawBuffer(const FrameBuffer& buffer) {
  if (buffer.format() != FrameBuffer::Format::kNV12 &&
      buffer.format() != FrameBuffer::Format::kNV21) {
    return absl::InvalidArgumentError(
        "Only support getting biplanar UV buffer from NV12/NV21 frame buffer.");
  }
  ASSIGN_OR_RETURN(FrameBuffer::YuvData yuv_data,
                   FrameBuffer::GetYuvDataFromFrameBuffer(buffer));
  const uint8_t* uv_buffer = buffer.format() == FrameBuffer::Format::kNV12
                                 ? yuv_data.u_buffer
                                 : yuv_data.v_buffer;
  return uv_buffer;
}

absl::StatusOr<FrameBuffer::Dimension> GetUvPlaneDimension(
    FrameBuffer::Dimension dimension, FrameBuffer::Format format) {
  if (dimension.width <= 0 || dimension.height <= 0) {
    return absl::InvalidArgumentError(
        absl::StrFormat("Invalid input dimension: {%d, %d}.", dimension.width,
                        dimension.height));
  }
  switch (format) {
    case FrameBuffer::Format::kNV12:
    case FrameBuffer::Format::kNV21:
    case FrameBuffer::Format::kYV12:
    case FrameBuffer::Format::kYV21:
      return FrameBuffer::Dimension{(dimension.width + 1) / 2,
                                    (dimension.height + 1) / 2};
    default:
      return absl::InvalidArgumentError(
          absl::StrFormat("Input format is not YUV-like: %i.", format));
  }
}

FrameBuffer::Dimension GetCropDimension(int x0, int x1, int y0, int y1) {
  return {x1 - x0 + 1, y1 - y0 + 1};
}

}  // namespace frame_buffer
}  // namespace mediapipe
