// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

#pragma once

#include <map>
#include <optional>
#include <string>
#include <string_view>
#include "StableABICompat.h"

namespace facebook::torchcodec {

enum ColorConversionLibrary {
  // Use the libavfilter library for color conversion.
  FILTERGRAPH,
  // Use the libswscale library for color conversion.
  SWSCALE
};

// The resolved output dtype. Only UINT8 or FLOAT32 — no AUTO.
// All code downstream of addVideoStream() should use this.
enum class OutputDtype { UINT8, FLOAT32 };

// The user-facing output dtype config, which may include AUTO.
// AUTO is resolved in addVideoStream() into an OutputDtype.
// UINT8: Always output uint8 tensors (default, backward compatible). Uses an
//        8-bit / RGB24 intermediate.
// FLOAT32: Always output float32 tensors normalized to [0, 1]. Uses a 16-bit /
//          RGB48 intermediate so the YUV->RGB matrix output is preserved at
//          full precision through the float cast, regardless of source bit
//          depth.
// AUTO: Output uint8 for SDR (<=8-bit) sources, float32 for HDR (>8-bit).
enum class OutputDtypeConfig { UINT8, FLOAT32, AUTO };

struct VideoStreamOptions {
  VideoStreamOptions() {}

  // Number of threads we pass to FFMPEG for decoding.
  // 0 means FFMPEG will choose the number of threads automatically to fully
  // utilize all cores. If not set, it will be the default FFMPEG behavior for
  // the given codec.
  std::optional<int> ffmpeg_thread_count;

  // Currently the dimension order can be either NHWC or NCHW.
  // H=height, W=width, C=channel.
  std::string dimension_order = "NCHW";

  // By default we have to use filtergraph, as it is more general. We can only
  // use swscale when we have met strict requirements. See
  // CpuDeviceInterface::initialze() for the logic.
  ColorConversionLibrary color_conversion_library =
      ColorConversionLibrary::FILTERGRAPH;

  // By default we use CPU for decoding for both C++ and python users.
  // Note: This is not used for video encoding, because device is determined by
  // the device of the input frame tensor.
  StableDevice device = StableDevice(kStableCPU);
  // Device variant (e.g., "nvdec", "ffmpeg")
  std::string_view device_variant = "default";

  // The user-specified output dtype config. May be AUTO, which gets resolved
  // in addVideoStream() into outputDtype below.
  OutputDtypeConfig output_dtype_config = OutputDtypeConfig::UINT8;

  // Set by addVideoStream() after resolving AUTO. All downstream code should
  // read this field.
  OutputDtype output_dtype = OutputDtype::UINT8;

  // Encoding options
  std::optional<std::string> codec;
  // Optional pixel format for video encoding (e.g., "yuv420p", "yuv444p")
  // If not specified, uses codec's default format.
  std::optional<std::string> pixel_format;
  std::optional<double> crf;
  std::optional<std::string> preset;
  std::optional<std::map<std::string, std::string>> extra_options;
};

struct AudioStreamOptions {
  AudioStreamOptions() {}

  // Encoding only
  std::optional<int> bit_rate;
  // Decoding and encoding:
  std::optional<int> num_channels;
  std::optional<int> sample_rate;
};

} // namespace facebook::torchcodec