// 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.

#include <map>
#include <mutex>
#include <vector>
#include "StableABICompat.h"
#include "ValidationUtils.h"

#include "BetaCudaDeviceInterface.h"

#include "DeviceInterface.h"
#include "FFMPEGCommon.h"
#include "Logging.h"
#include "NVDECCache.h"

#include "NVCUVIDRuntimeLoader.h"
#include "color_conversion.h"
#include "nvcuvid_include/cuviddec.h"
#include "nvcuvid_include/nvcuvid.h"

extern "C" {
#include <libavutil/hwcontext_cuda.h>
#include <libavutil/pixdesc.h>
}

namespace facebook::torchcodec {

namespace {

// Per-device cache for cuvidGetDecoderCaps results.
// The key is a tuple of (device index, codec type, chroma format, bit depth
// minus 8).
struct DecoderCapsCache {
  using Key =
      std::tuple<int, cudaVideoCodec, cudaVideoChromaFormat, unsigned int>;
  std::map<Key, CUVIDDECODECAPS> cache;
  std::mutex mutex;

  std::pair<CUresult, CUVIDDECODECAPS> getDecoderCaps(
      int deviceIndex,
      cudaVideoCodec codecType,
      cudaVideoChromaFormat chromaFormat,
      unsigned int bitDepthMinus8) {
    Key key{deviceIndex, codecType, chromaFormat, bitDepthMinus8};

    std::lock_guard<std::mutex> lock(mutex);
    auto it = cache.find(key);
    if (it != cache.end()) {
      return {CUDA_SUCCESS, it->second};
    }

    CUVIDDECODECAPS caps = {};
    caps.eCodecType = codecType;
    caps.eChromaFormat = chromaFormat;
    caps.nBitDepthMinus8 = bitDepthMinus8;

    CUresult result = cuvidGetDecoderCaps(&caps);
    if (result == CUDA_SUCCESS) {
      cache[key] = caps;
    }
    return {result, caps};
  }
};

static DecoderCapsCache& getDecoderCapsCache() {
  static DecoderCapsCache cache;
  return cache;
}

cudaVideoSurfaceFormat getPreferredSurfaceFormat(OutputDtype outputDtype) {
  return outputDtype == OutputDtype::FLOAT32 ? cudaVideoSurfaceFormat_P016
                                             : cudaVideoSurfaceFormat_NV12;
}

static bool g_cuda_nvdec = registerDeviceInterface(
    DeviceInterfaceKey(kStableCUDA, /*variant=*/"default"),
    [](const StableDevice& device) {
      return new BetaCudaDeviceInterface(device);
    });

static int CUDAAPI
pfnSequenceCallback(void* pUserData, CUVIDEOFORMAT* videoFormat) {
  auto decoder = static_cast<BetaCudaDeviceInterface*>(pUserData);
  return decoder->streamPropertyChange(videoFormat);
}

static int CUDAAPI
pfnDecodePictureCallback(void* pUserData, CUVIDPICPARAMS* picParams) {
  auto decoder = static_cast<BetaCudaDeviceInterface*>(pUserData);
  return decoder->frameReadyForDecoding(picParams);
}

static int CUDAAPI
pfnDisplayPictureCallback(void* pUserData, CUVIDPARSERDISPINFO* dispInfo) {
  auto decoder = static_cast<BetaCudaDeviceInterface*>(pUserData);
  return decoder->frameReadyInDisplayOrder(dispInfo);
}

static UniqueCUvideodecoder createDecoder(
    CUVIDEOFORMAT* videoFormat,
    cudaVideoSurfaceFormat surfaceFormat) {
  // Decoder creation parameters, most are taken from DALI
  CUVIDDECODECREATEINFO decoderParams = {};
  decoderParams.bitDepthMinus8 = videoFormat->bit_depth_luma_minus8;
  decoderParams.ChromaFormat = videoFormat->chroma_format;
  decoderParams.OutputFormat = surfaceFormat;
  decoderParams.ulCreationFlags = cudaVideoCreate_Default;
  decoderParams.CodecType = videoFormat->codec;
  decoderParams.ulHeight = videoFormat->coded_height;
  decoderParams.ulWidth = videoFormat->coded_width;
  decoderParams.ulMaxHeight = videoFormat->coded_height;
  decoderParams.ulMaxWidth = videoFormat->coded_width;
  decoderParams.ulTargetHeight =
      videoFormat->display_area.bottom - videoFormat->display_area.top;
  decoderParams.ulTargetWidth =
      videoFormat->display_area.right - videoFormat->display_area.left;
  decoderParams.ulNumDecodeSurfaces = videoFormat->min_num_decode_surfaces;
  // We should only ever need 1 output surface, since we process frames
  // sequentially, and we always unmap the previous frame before mapping a new
  // one.
  // TODONVDEC P3: set this to 2, allow for 2 frames to be mapped at a time, and
  // benchmark to see if this makes any difference.
  decoderParams.ulNumOutputSurfaces = 1;
  decoderParams.display_area.left = videoFormat->display_area.left;
  decoderParams.display_area.right = videoFormat->display_area.right;
  decoderParams.display_area.top = videoFormat->display_area.top;
  decoderParams.display_area.bottom = videoFormat->display_area.bottom;

  CUvideodecoder* decoder = new CUvideodecoder();
  CUresult result = cuvidCreateDecoder(decoder, &decoderParams);
  STD_TORCH_CHECK(
      result == CUDA_SUCCESS, "Failed to create NVDEC decoder: ", result);
  return UniqueCUvideodecoder(decoder, CUvideoDecoderDeleter{});
}

std::optional<cudaVideoChromaFormat> validateChromaSupport(
    const AVPixFmtDescriptor* desc) {
  // Return the corresponding cudaVideoChromaFormat if supported, std::nullopt
  // otherwise.
  STD_TORCH_CHECK(desc != nullptr, "desc can't be null");

  if (desc->nb_components == 1) {
    return cudaVideoChromaFormat_Monochrome;
  } else if (desc->nb_components >= 3 && !(desc->flags & AV_PIX_FMT_FLAG_RGB)) {
    // Make sure it's YUV: has chroma planes and isn't RGB
    if (desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0) {
      return cudaVideoChromaFormat_444; // 1x1 subsampling = 4:4:4
    } else if (desc->log2_chroma_w == 1 && desc->log2_chroma_h == 1) {
      return cudaVideoChromaFormat_420; // 2x2 subsampling = 4:2:0
    } else if (desc->log2_chroma_w == 1 && desc->log2_chroma_h == 0) {
      return cudaVideoChromaFormat_422; // 2x1 subsampling = 4:2:2
    }
  }

  return std::nullopt;
}

std::optional<cudaVideoCodec> validateCodecSupport(AVCodecID codecId) {
  // Return the corresponding cudaVideoCodec if supported, std::nullopt
  // otherwise
  // Note that we currently return nullopt (and thus fallback to CPU) for some
  // codecs that are technically supported by NVDEC, see comment below.
  switch (codecId) {
    case AV_CODEC_ID_H264:
      return cudaVideoCodec_H264;
    case AV_CODEC_ID_HEVC:
      return cudaVideoCodec_HEVC;
    case AV_CODEC_ID_AV1:
      return cudaVideoCodec_AV1;
    case AV_CODEC_ID_VP9:
      return cudaVideoCodec_VP9;
    case AV_CODEC_ID_VP8:
      return cudaVideoCodec_VP8;
    case AV_CODEC_ID_MPEG4:
      return cudaVideoCodec_MPEG4;
    // Formats below are currently not tested, but they should "mostly" work.
    // MPEG1 was briefly locally tested and it was ok-ish despite duration being
    // off. Since they're far less popular, we keep them disabled by default but
    // we can consider enabling them upon user requests.
    // case AV_CODEC_ID_MPEG1VIDEO:
    //   return cudaVideoCodec_MPEG1;
    // case AV_CODEC_ID_MPEG2VIDEO:
    //   return cudaVideoCodec_MPEG2;
    // case AV_CODEC_ID_MJPEG:
    //   return cudaVideoCodec_JPEG;
    // case AV_CODEC_ID_VC1:
    //   return cudaVideoCodec_VC1;
    default:
      return std::nullopt;
  }
}

std::optional<cudaVideoSurfaceFormat> getNVDECSurfaceFormat(
    const StableDevice& device,
    const SharedAVCodecContext& codecContext,
    OutputDtype outputDtype) {
  // Return the surface format to use for NVDEC decoding if the stream is
  // supported, or nullopt to fall back to CPU.

  auto codecType = validateCodecSupport(codecContext->codec_id);
  if (!codecType.has_value()) {
    return std::nullopt;
  }

  const AVPixFmtDescriptor* desc = av_pix_fmt_desc_get(codecContext->pix_fmt);
  if (!desc) {
    return std::nullopt;
  }

  auto chromaFormat = validateChromaSupport(desc);
  if (!chromaFormat.has_value()) {
    return std::nullopt;
  }

  auto bitDepthMinus8 = static_cast<unsigned int>(desc->comp[0].depth - 8);
  auto [result, caps] = getDecoderCapsCache().getDecoderCaps(
      getDeviceIndex(device),
      codecType.value(),
      chromaFormat.value(),
      bitDepthMinus8);
  if (result != CUDA_SUCCESS) {
    return std::nullopt;
  }

  if (!caps.bIsSupported) {
    return std::nullopt;
  }

  auto coded_width = static_cast<unsigned int>(codecContext->coded_width);
  auto coded_height = static_cast<unsigned int>(codecContext->coded_height);
  if (coded_width < static_cast<unsigned int>(caps.nMinWidth) ||
      coded_height < static_cast<unsigned int>(caps.nMinHeight) ||
      coded_width > caps.nMaxWidth || coded_height > caps.nMaxHeight) {
    return std::nullopt;
  }

  // See nMaxMBCount in cuviddec.h
  constexpr unsigned int macroblockConstant = 256;
  if (coded_width * coded_height / macroblockConstant > caps.nMaxMBCount) {
    return std::nullopt;
  }

  auto preferredFormat = getPreferredSurfaceFormat(outputDtype);
  if ((caps.nOutputFormatMask >> preferredFormat) & 1) {
    return preferredFormat;
  }

  // P016 is typically not supported on 8-bit SDR content. In such cases, we
  // try to fall back to NV12 if supported:
  // NVDEC will decode to NV12, our kernel will do NV12 -> RGB producing
  // uint8, and maybePermuteAndConvertToFloat32 will cast uint8 -> float32.
  // For HDR content, NV12 would lose precision, so we fall back to CPU instead.
  if (preferredFormat == cudaVideoSurfaceFormat_P016 && bitDepthMinus8 == 0 &&
      ((caps.nOutputFormatMask >> cudaVideoSurfaceFormat_NV12) & 1)) {
    return cudaVideoSurfaceFormat_NV12;
  }

  return std::nullopt;
}

// Callback for freeing CUDA memory associated with AVFrame see where it's used
// for more details.
void cudaBufferFreeCallback(void* opaque, [[maybe_unused]] uint8_t* data) {
  cudaFree(opaque);
}

} // namespace

BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device)
    : DeviceInterface(device) {
  STD_TORCH_CHECK(g_cuda_nvdec, "NvdecCudaDeviceInterface was not registered!");
  STD_TORCH_CHECK(
      device_.type() == kStableCUDA, "Unsupported device: must be CUDA");

  initializeCudaContextWithPytorch(device_);

  nvcuvidAvailable_ = loadNVCUVIDLibrary();
}

void BetaCudaDeviceInterface::initializeVideo(
    const AVStream* avStream,
    const UniqueDecodingAVFormatContext& avFormatCtx,
    const VideoStreamOptions& videoStreamOptions,
    const std::vector<std::unique_ptr<Transform>>& transforms,
    const std::optional<FrameDims>& resizedOutputDims) {
  STD_TORCH_CHECK(avStream != nullptr, "AVStream cannot be null");
  rotation_ = rotationFromDegrees(getRotationFromStream(avStream));
  outputDtype_ = videoStreamOptions.outputDtype;

  auto maybeSurfaceFormat = nvcuvidAvailable_
      ? getNVDECSurfaceFormat(device_, codecContext_, outputDtype_)
      : std::nullopt;

  if (!maybeSurfaceFormat.has_value()) {
    if (!nvcuvidAvailable_) {
      TC_LOG("NVCUVID library not available; falling back to CPU decoding.");
    } else {
      TC_LOG(
          "Video stream not supported by NVDEC; falling back to CPU decoding.");
    }
    cpuFallback_ = createDeviceInterface(kStableCPU);
    STD_TORCH_CHECK(
        cpuFallback_ != nullptr, "Failed to create CPU device interface");
    cpuFallback_->initialize(codecContext_);
    cpuFallback_->initializeVideo(
        avStream,
        avFormatCtx,
        videoStreamOptions,
        transforms,
        resizedOutputDims);
    return;
  }

  surfaceFormat_ = maybeSurfaceFormat.value();
  timeBase_ = avStream->time_base;
  frameRateAvgFromFFmpeg_ = avStream->r_frame_rate;

  const AVCodecParameters* codecPar = avStream->codecpar;
  STD_TORCH_CHECK(codecPar != nullptr, "CodecParameters cannot be null");

  initializeBSF(codecPar, avFormatCtx);

  // Create parser. Default values that aren't obvious are taken from DALI.
  CUVIDPARSERPARAMS parserParams = {};
  auto codecType = validateCodecSupport(codecPar->codec_id);
  STD_TORCH_CHECK(
      codecType.has_value(),
      "This should never happen, we should be using the CPU fallback by now. "
      "Please report a bug.");
  parserParams.CodecType = codecType.value();
  parserParams.ulMaxNumDecodeSurfaces = 8;
  parserParams.ulMaxDisplayDelay = 0;
  // Callback setup, all are triggered by the parser within a call
  // to cuvidParseVideoData
  parserParams.pUserData = this;
  parserParams.pfnSequenceCallback = pfnSequenceCallback;
  parserParams.pfnDecodePicture = pfnDecodePictureCallback;
  parserParams.pfnDisplayPicture = pfnDisplayPictureCallback;

  // Some containers (e.g. MP4/MOV) store codec config (H.264 SPS/PPS,
  // MPEG-4 VOS/VOL, etc.) in extradata rather than inline in the
  // bitstream. The NVCUVID parser needs this data to initialize, so we
  // pass it via pExtVideoInfo. Same approach as DALI and FFmpeg cuviddec.
  // DALI does the same thing
  // https://github.com/NVIDIA/DALI/blob/ae79f316ae9b14c464d9cb98465f7f783da9ea89/dali/operators/video/frames_decoder_gpu.cc#L402-L408
  if (codecPar->extradata_size > 0) {
    auto seqhdrSize = std::min(
        static_cast<size_t>(codecPar->extradata_size),
        sizeof(parserExtInfo_.raw_seqhdr_data));
    parserExtInfo_.format.seqhdr_data_length = seqhdrSize;
    memcpy(parserExtInfo_.raw_seqhdr_data, codecPar->extradata, seqhdrSize);
    parserParams.pExtVideoInfo = &parserExtInfo_;
  }

  CUresult result = cuvidCreateVideoParser(&videoParser_, &parserParams);
  STD_TORCH_CHECK(
      result == CUDA_SUCCESS, "Failed to create video parser: ", result);
}

BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
  if (decoder_) {
    // DALI doesn't seem to do any particular cleanup of the decoder before
    // sending it to the cache, so we probably don't need to do anything either.
    // Just to be safe, we flush.
    // What happens to those decode surfaces that haven't yet been mapped is
    // unclear.
    flush();
    unmapPreviousFrame();
    NVDECCache::getCache(device_).returnDecoder(
        &videoFormat_, surfaceFormat_, std::move(decoder_));
  }

  if (videoParser_) {
    cuvidDestroyVideoParser(videoParser_);
    videoParser_ = nullptr;
  }
}

void BetaCudaDeviceInterface::initialize(
    const SharedAVCodecContext& codecContext) {
  codecContext_ = codecContext;
}

void BetaCudaDeviceInterface::initializeBSF(
    const AVCodecParameters* codecPar,
    const UniqueDecodingAVFormatContext& avFormatCtx) {
  // Setup bit stream filters (BSF):
  // https://ffmpeg.org/doxygen/7.0/group__lavc__bsf.html
  // This is only needed for some formats, like H264 or HEVC.

  STD_TORCH_CHECK(codecPar != nullptr, "codecPar cannot be null");
  STD_TORCH_CHECK(avFormatCtx != nullptr, "AVFormatContext cannot be null");
  STD_TORCH_CHECK(
      avFormatCtx->iformat != nullptr,
      "AVFormatContext->iformat cannot be null");
  std::string filterName;

  // Matching logic is taken from DALI
  switch (codecPar->codec_id) {
    case AV_CODEC_ID_H264: {
      const std::string formatName = avFormatCtx->iformat->long_name
          ? avFormatCtx->iformat->long_name
          : "";

      if (formatName == "QuickTime / MOV" ||
          formatName == "FLV (Flash Video)" ||
          formatName == "Matroska / WebM" || formatName == "raw H.264 video") {
        filterName = "h264_mp4toannexb";
      }
      break;
    }

    case AV_CODEC_ID_HEVC: {
      const std::string formatName = avFormatCtx->iformat->long_name
          ? avFormatCtx->iformat->long_name
          : "";

      if (formatName == "QuickTime / MOV" ||
          formatName == "FLV (Flash Video)" ||
          formatName == "Matroska / WebM" || formatName == "raw HEVC video") {
        filterName = "hevc_mp4toannexb";
      }
      break;
    }
    case AV_CODEC_ID_MPEG4: {
      const std::string formatName =
          avFormatCtx->iformat->name ? avFormatCtx->iformat->name : "";
      if (formatName == "avi") {
        filterName = "mpeg4_unpack_bframes";
      }
      break;
    }

    default:
      // No bitstream filter needed for other codecs
      break;
  }

  if (filterName.empty()) {
    // Only initialize BSF if we actually need one
    return;
  }

  const AVBitStreamFilter* avBSF = av_bsf_get_by_name(filterName.c_str());
  STD_TORCH_CHECK(
      avBSF != nullptr, "Failed to find bitstream filter: ", filterName);

  AVBSFContext* avBSFContext = nullptr;
  int retVal = av_bsf_alloc(avBSF, &avBSFContext);
  STD_TORCH_CHECK(
      retVal >= AVSUCCESS,
      "Failed to allocate bitstream filter: ",
      getFFMPEGErrorStringFromErrorCode(retVal));

  bitstreamFilter_.reset(avBSFContext);

  retVal = avcodec_parameters_copy(bitstreamFilter_->par_in, codecPar);
  STD_TORCH_CHECK(
      retVal >= AVSUCCESS,
      "Failed to copy codec parameters: ",
      getFFMPEGErrorStringFromErrorCode(retVal));

  retVal = av_bsf_init(bitstreamFilter_.get());
  STD_TORCH_CHECK(
      retVal == AVSUCCESS,
      "Failed to initialize bitstream filter: ",
      getFFMPEGErrorStringFromErrorCode(retVal));
}

// This callback is called by the parser within cuvidParseVideoData when there
// is a change in the stream's properties (like resolution change), as specified
// by CUVIDEOFORMAT. Particularly (but not just!), this is called at the very
// start of the stream.
// TODONVDEC P1: Code below mostly assume this is called only once at the start,
// we should handle the case of multiple calls. Probably need to flush buffers,
// etc.
int BetaCudaDeviceInterface::streamPropertyChange(CUVIDEOFORMAT* videoFormat) {
  STD_TORCH_CHECK(videoFormat != nullptr, "Invalid video format");

  videoFormat_ = *videoFormat;

  if (videoFormat_.min_num_decode_surfaces == 0) {
    // Same as DALI's fallback
    videoFormat_.min_num_decode_surfaces = 20;
  }

  if (!decoder_) {
    decoder_ =
        NVDECCache::getCache(device_).getDecoder(videoFormat, surfaceFormat_);

    if (!decoder_) {
      // TODONVDEC P2: consider re-configuring an existing decoder instead of
      // re-creating one. See docs, see DALI. Re-configuration doesn't seem to
      // be enabled in DALI by default.
      decoder_ = createDecoder(videoFormat, surfaceFormat_);
    }

    STD_TORCH_CHECK(decoder_, "Failed to get or create decoder");
  }

  // DALI also returns min_num_decode_surfaces from this function. This
  // instructs the parser to reset its ulMaxNumDecodeSurfaces field to this
  // value.
  return static_cast<int>(videoFormat_.min_num_decode_surfaces);
}

// Moral equivalent of avcodec_send_packet(). Here, we pass the AVPacket down to
// the NVCUVID parser.
int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) {
  if (cpuFallback_) {
    return cpuFallback_->sendPacket(packet);
  }

  STD_TORCH_CHECK(
      packet.get() && packet->data && packet->size > 0,
      "sendPacket received an empty packet, this is unexpected, please report.");

  // Apply BSF if needed. We want applyBSF to return a *new* filtered packet, or
  // the original one if no BSF is needed. This new filtered packet must be
  // allocated outside of applyBSF: if it were allocated inside applyBSF, it
  // would be destroyed at the end of the function, leaving us with a dangling
  // reference.
  AutoAVPacket filteredAutoPacket;
  ReferenceAVPacket filteredPacket(filteredAutoPacket);
  ReferenceAVPacket& packetToSend = applyBSF(packet, filteredPacket);

  CUVIDSOURCEDATAPACKET cuvidPacket = {};
  cuvidPacket.payload = packetToSend->data;
  cuvidPacket.payload_size = packetToSend->size;
  cuvidPacket.flags = CUVID_PKT_TIMESTAMP;
  cuvidPacket.timestamp = packetToSend->pts;

  return sendCuvidPacket(cuvidPacket);
}

int BetaCudaDeviceInterface::sendEOFPacket() {
  if (cpuFallback_) {
    return cpuFallback_->sendEOFPacket();
  }

  CUVIDSOURCEDATAPACKET cuvidPacket = {};
  cuvidPacket.flags = CUVID_PKT_ENDOFSTREAM;
  eofSent_ = true;

  return sendCuvidPacket(cuvidPacket);
}

int BetaCudaDeviceInterface::sendCuvidPacket(
    CUVIDSOURCEDATAPACKET& cuvidPacket) {
  CUresult result = cuvidParseVideoData(videoParser_, &cuvidPacket);
  return result == CUDA_SUCCESS ? AVSUCCESS : AVERROR_EXTERNAL;
}

ReferenceAVPacket& BetaCudaDeviceInterface::applyBSF(
    ReferenceAVPacket& packet,
    ReferenceAVPacket& filteredPacket) {
  if (!bitstreamFilter_) {
    return packet;
  }

  int retVal = av_bsf_send_packet(bitstreamFilter_.get(), packet.get());
  STD_TORCH_CHECK(
      retVal >= AVSUCCESS,
      "Failed to send packet to bitstream filter: ",
      getFFMPEGErrorStringFromErrorCode(retVal));

  // TODO P1: the docs mention there can theoretically be multiple output
  // packets for a single input, i.e. we may need to call av_bsf_receive_packet
  // more than once. We should figure out whether that applies to the BSF we're
  // using.
  retVal = av_bsf_receive_packet(bitstreamFilter_.get(), filteredPacket.get());
  STD_TORCH_CHECK(
      retVal >= AVSUCCESS,
      "Failed to receive packet from bitstream filter: ",
      getFFMPEGErrorStringFromErrorCode(retVal));

  return filteredPacket;
}

// Parser triggers this callback within cuvidParseVideoData when a frame is
// ready to be decoded, i.e. the parser received all the necessary packets for a
// given frame. It means we can send that frame to be decoded by the hardware
// NVDEC decoder by calling cuvidDecodePicture.
int BetaCudaDeviceInterface::frameReadyForDecoding(CUVIDPICPARAMS* picParams) {
  STD_TORCH_CHECK(picParams != nullptr, "Invalid picture parameters");
  STD_TORCH_CHECK(decoder_, "Decoder not initialized before picture decode");
  // Send frame to be decoded by NVDEC. This may or may not block, depending on
  // the internal state of the NVDEC. Presumably, when it blocks, it gets
  // automatically unblocked once a frame has been decoded, although how and
  // when it happens is unclear. The docs say:
  // > cuvidDecodePicture() will stall if wait queue on NVDEC inside driver is
  //   full.
  // and cuviddec.h says:
  // > cuvidDecodePicture may block the calling thread if there are too many
  //   pictures pending in the decode queue.
  CUresult result = cuvidDecodePicture(*decoder_.get(), picParams);

  // Yes, you're reading that right, 0 means error, 1 means success
  return (result == CUDA_SUCCESS);
}

int BetaCudaDeviceInterface::frameReadyInDisplayOrder(
    CUVIDPARSERDISPINFO* dispInfo) {
  readyFrames_.push(*dispInfo);
  return 1; // success
}

// Moral equivalent of avcodec_receive_frame().
int BetaCudaDeviceInterface::receiveFrame(UniqueAVFrame& avFrame) {
  if (cpuFallback_) {
    return cpuFallback_->receiveFrame(avFrame);
  }

  if (readyFrames_.empty()) {
    // No frame found, instruct caller to try again later after sending more
    // packets, or to stop if EOF was already sent.
    return eofSent_ ? AVERROR_EOF : AVERROR(EAGAIN);
  }

  CUVIDPARSERDISPINFO dispInfo = readyFrames_.front();
  readyFrames_.pop();

  CUVIDPROCPARAMS procParams = {};
  procParams.progressive_frame = dispInfo.progressive_frame;
  procParams.top_field_first = dispInfo.top_field_first;
  procParams.unpaired_field = dispInfo.repeat_first_field < 0;
  // We set the NVDEC stream to the current stream. It will be waited upon
  // by the color conversion stream before any color conversion.
  // Re types: we get a cudaStream_t from PyTorch but it's interchangeable with
  // CUstream
  procParams.output_stream =
      reinterpret_cast<CUstream>(getCurrentCudaStream(device_.index()));

  CUdeviceptr framePtr = 0;
  unsigned int pitch = 0;

  // We know the frame we want was sent to the hardware decoder, but now we need
  // to "map" it to an "output surface" before we can use its data. This is a
  // blocking calls that waits until the frame is fully decoded and ready to be
  // used.
  // When a frame is mapped to an output surface, it needs to be unmapped
  // eventually, so that the decoder can re-use the output surface. Failing to
  // unmap will cause map to eventually fail. DALI unmaps frames almost
  // immediately  after mapping them: they do the color-conversion in-between,
  // which involves a copy of the data, so that works.
  // We, OTOH, will do the color-conversion later, outside of ReceiveFrame(). So
  // we unmap here: just before mapping a new frame. At that point we know that
  // the previously-mapped frame is no longer needed: it was either
  // color-converted (with a copy), or that's a frame that was discarded in
  // SingleStreamDecoder. Either way, the underlying output surface can be
  // safely re-used.
  unmapPreviousFrame();
  CUresult result = cuvidMapVideoFrame(
      *decoder_.get(), dispInfo.picture_index, &framePtr, &pitch, &procParams);
  if (result != CUDA_SUCCESS) {
    return AVERROR_EXTERNAL;
  }
  previouslyMappedFrame_ = framePtr;

  avFrame = convertCudaFrameToAVFrame(framePtr, pitch, dispInfo);

  return AVSUCCESS;
}

void BetaCudaDeviceInterface::unmapPreviousFrame() {
  if (previouslyMappedFrame_ == 0) {
    return;
  }
  CUresult result =
      cuvidUnmapVideoFrame(*decoder_.get(), previouslyMappedFrame_);
  STD_TORCH_CHECK(
      result == CUDA_SUCCESS, "Failed to unmap previous frame: ", result);
  previouslyMappedFrame_ = 0;
}

UniqueAVFrame BetaCudaDeviceInterface::convertCudaFrameToAVFrame(
    CUdeviceptr framePtr,
    unsigned int pitch,
    const CUVIDPARSERDISPINFO& dispInfo) {
  STD_TORCH_CHECK(framePtr != 0, "Invalid CUDA frame pointer");

  // Get frame dimensions from video format display area (not coded dimensions)
  // This matches DALI's approach and avoids padding issues
  int width = videoFormat_.display_area.right - videoFormat_.display_area.left;
  int height = videoFormat_.display_area.bottom - videoFormat_.display_area.top;

  STD_TORCH_CHECK(width > 0 && height > 0, "Invalid frame dimensions");
  STD_TORCH_CHECK(
      pitch >= static_cast<unsigned int>(width), "Pitch must be >= width");

  UniqueAVFrame avFrame(av_frame_alloc());
  STD_TORCH_CHECK(avFrame.get() != nullptr, "Failed to allocate AVFrame");

  avFrame->width = width;
  avFrame->height = height;
  avFrame->format = (surfaceFormat_ == cudaVideoSurfaceFormat_P016)
      ? AV_PIX_FMT_P016LE
      : AV_PIX_FMT_NV12;
  avFrame->pts = dispInfo.timestamp;

  // TODONVDEC P2: We compute the duration based on average frame rate info, so
  // so if the video has variable frame rate, the durations may be off. We
  // should try to see if we can set the duration more accurately. Unfortunately
  // it's not given by dispInfo. One option would be to set it based on the pts
  // difference between consecutive frames, if the next frame is already
  // available.
  // Note that we used to rely on videoFormat_.frame_rate for this, but that
  // proved less accurate than FFmpeg.
  setDuration(avFrame, computeSafeDuration(frameRateAvgFromFFmpeg_, timeBase_));

  // We need to assign the frame colorspace. This is crucial for proper color
  // conversion. NVCUVID stores that in the matrix_coefficients field, but
  // doesn't document the semantics of the values. Claude code generated this,
  // which seems to work. Reassuringly, the values seem to match the
  // corresponding indices in the FFmpeg enum for colorspace conversion
  // (ff_yuv2rgb_coeffs):
  // https://ffmpeg.org/doxygen/trunk/yuv2rgb_8c_source.html#l00047
  switch (videoFormat_.video_signal_description.matrix_coefficients) {
    case 1:
      avFrame->colorspace = AVCOL_SPC_BT709;
      break;
    case 6:
      avFrame->colorspace = AVCOL_SPC_SMPTE170M; // BT.601
      break;
    case 9:
      avFrame->colorspace = AVCOL_SPC_BT2020_NCL;
      break;
    case 10:
      avFrame->colorspace = AVCOL_SPC_BT2020_CL;
      break;
    default:
      // Default to BT.601
      avFrame->colorspace = AVCOL_SPC_SMPTE170M;
      break;
  }

  avFrame->color_range =
      videoFormat_.video_signal_description.video_full_range_flag
      ? AVCOL_RANGE_JPEG
      : AVCOL_RANGE_MPEG;

  // NVDEC's surface layout places the UV plane after the Y plane. For
  // NV12/P016 the Y plane has an even number of rows (NVDEC rounds up
  // internally), so we must use the rounded-up height for the UV offset.
  unsigned int evenHeight = roundUpToEven(height);
  avFrame->data[0] = reinterpret_cast<uint8_t*>(framePtr);
  avFrame->data[1] =
      reinterpret_cast<uint8_t*>(framePtr + (pitch * evenHeight));
  avFrame->data[2] = nullptr;
  avFrame->data[3] = nullptr;
  avFrame->linesize[0] = pitch;
  avFrame->linesize[1] = pitch;
  avFrame->linesize[2] = 0;
  avFrame->linesize[3] = 0;

  return avFrame;
}

void BetaCudaDeviceInterface::flush() {
  if (cpuFallback_) {
    cpuFallback_->flush();
    return;
  }

  // The NVCUVID docs mention that after seeking, i.e. when flush() is called,
  // we should send a packet with the CUVID_PKT_DISCONTINUITY flag. The docs
  // don't say whether this should be an empty packet, or whether it should be a
  // flag on the next non-empty packet. It doesn't matter: neither work :)
  // Sending an EOF packet, however, does work. So we do that. And we re-set the
  // eofSent_ flag to false because that's not a true EOF notification.
  sendEOFPacket();
  eofSent_ = false;

  std::queue<CUVIDPARSERDISPINFO> emptyQueue;
  std::swap(readyFrames_, emptyQueue);
}

UniqueAVFrame BetaCudaDeviceInterface::transferCpuFrameToGpu(
    UniqueAVFrame& cpuFrame,
    AVPixelFormat targetPixFmt) {
  // This is called in the context of the CPU fallback: the frame was decoded on
  // the CPU, and in this function we convert that frame into NV12 or P016
  // format and send it to the GPU.
  // We do that in 2 steps:
  // - First we convert the input CPU frame into an intermediate NV12/P016 CPU
  //   frame using sws_scale.
  // - Then we allocate GPU memory and copy the CPU frame to the GPU. This
  //   is what we return.
  // Since NV12/P016 require even dimensions, the returned frame will have even
  // (rounded up) width and height, even if the original CPU frame had odd
  // dimensions.

  STD_TORCH_CHECK(cpuFrame != nullptr, "CPU frame cannot be null");
  // NV12 = 1 byte per sample, P016 = 2 bytes per sample
  STD_TORCH_CHECK(
      targetPixFmt == AV_PIX_FMT_NV12 || targetPixFmt == AV_PIX_FMT_P016LE,
      "targetPixFmt must be NV12 or P016LE");
  int bytesPerSample = (targetPixFmt == AV_PIX_FMT_P016LE) ? 2 : 1;

  int width = cpuFrame->width;
  int height = cpuFrame->height;
  int evenWidth = roundUpToEven(width);
  int evenHeight = roundUpToEven(height);

  UniqueAVFrame intermediateCpuFrame(av_frame_alloc());
  STD_TORCH_CHECK(
      intermediateCpuFrame != nullptr,
      "Failed to allocate intermediate CPU frame");

  intermediateCpuFrame->format = targetPixFmt;
  intermediateCpuFrame->width = evenWidth;
  intermediateCpuFrame->height = evenHeight;

  int ret = av_frame_get_buffer(intermediateCpuFrame.get(), 0);
  STD_TORCH_CHECK(
      ret >= 0,
      "Failed to allocate intermediate CPU frame buffer: ",
      getFFMPEGErrorStringFromErrorCode(ret));

  SwsConfig swsConfig(
      width,
      height,
      static_cast<AVPixelFormat>(cpuFrame->format),
      cpuFrame->colorspace,
      evenWidth,
      evenHeight,
      targetPixFmt);

  if (!swsContext_ || prevSwsConfig_ != swsConfig) {
    swsContext_ = createSwsContext(swsConfig, SWS_BILINEAR);
    prevSwsConfig_ = swsConfig;
  }

  int convertedHeight = sws_scale(
      swsContext_.get(),
      cpuFrame->data,
      cpuFrame->linesize,
      0,
      height,
      intermediateCpuFrame->data,
      intermediateCpuFrame->linesize);
  STD_TORCH_CHECK(
      convertedHeight == evenHeight,
      "sws_scale failed for CPU->NV12/P016 conversion");

  int rowBytes = evenWidth * bytesPerSample;
  int ySize = rowBytes * evenHeight;
  int uvSize = ySize / 2;
  size_t totalSize = static_cast<size_t>(ySize + uvSize);

  uint8_t* cudaBuffer = nullptr;
  cudaError_t err =
      cudaMalloc(reinterpret_cast<void**>(&cudaBuffer), totalSize);
  STD_TORCH_CHECK(
      err == cudaSuccess,
      "Failed to allocate CUDA memory: ",
      cudaGetErrorString(err));

  UniqueAVFrame gpuFrame(av_frame_alloc());
  STD_TORCH_CHECK(gpuFrame != nullptr, "Failed to allocate GPU AVFrame");

  gpuFrame->format = targetPixFmt;
  gpuFrame->width = evenWidth;
  gpuFrame->height = evenHeight;
  gpuFrame->data[0] = cudaBuffer;
  gpuFrame->data[1] = cudaBuffer + ySize;
  gpuFrame->linesize[0] = rowBytes;
  gpuFrame->linesize[1] = rowBytes;

  // Note that we use cudaMemcpy2D here instead of cudaMemcpy because the
  // linesizes (strides) may be different than the widths for the input CPU
  // frame. That's precisely what cudaMemcpy2D is for.
  err = cudaMemcpy2D(
      gpuFrame->data[0],
      gpuFrame->linesize[0],
      intermediateCpuFrame->data[0],
      intermediateCpuFrame->linesize[0],
      rowBytes,
      evenHeight,
      cudaMemcpyHostToDevice);
  STD_TORCH_CHECK(
      err == cudaSuccess,
      "Failed to copy Y plane to GPU: ",
      cudaGetErrorString(err));

  err = cudaMemcpy2D(
      gpuFrame->data[1],
      gpuFrame->linesize[1],
      intermediateCpuFrame->data[1],
      intermediateCpuFrame->linesize[1],
      rowBytes,
      evenHeight / 2,
      cudaMemcpyHostToDevice);
  STD_TORCH_CHECK(
      err == cudaSuccess,
      "Failed to copy UV plane to GPU: ",
      cudaGetErrorString(err));

  ret = av_frame_copy_props(gpuFrame.get(), cpuFrame.get());
  STD_TORCH_CHECK(
      ret >= 0,
      "Failed to copy frame properties: ",
      getFFMPEGErrorStringFromErrorCode(ret));

  // We need to make sure the CUDA memory is freed properly. Since we allocated
  // it ourselves, FFmpeg doesn't know how to free it. We associate a `free`
  // callback via opaque_ref that will be called by av_frame_free().
  gpuFrame->opaque_ref = av_buffer_create(
      nullptr, // data - we don't need any
      0, // data size
      cudaBufferFreeCallback, // callback triggered by av_frame_free()
      cudaBuffer, // parameter to callback
      0); // flags
  STD_TORCH_CHECK(
      gpuFrame->opaque_ref != nullptr,
      "Failed to create GPU memory cleanup reference");

  return gpuFrame;
}

void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
    UniqueAVFrame& avFrame,
    FrameOutput& frameOutput,
    std::optional<torch::stable::Tensor> preAllocatedOutputTensor) {
  if (cpuFallback_) {
    // When the CPU fallback happens, we'll try to run the color-conversion on
    // GPU by sending those CPU frames to the GPU as NV12 or P016 (See
    // transferCpuFrameToGpu() below). However, it's not always
    // possible: NV12/P016 would downsample 4:4:4 frames and lose chroma
    // resolution, resulting in poorly decoded frames. So for those, we still
    // do the color conversion on the CPU and then send the full RGB frame to
    // the GPU.
    const AVPixFmtDescriptor* desc =
        av_pix_fmt_desc_get(static_cast<AVPixelFormat>(avFrame->format));
    bool is444 = desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0;
    if (is444) {
      FrameOutput cpuFrameOutput;
      cpuFallback_->convertAVFrameToFrameOutput(avFrame, cpuFrameOutput);
      if (preAllocatedOutputTensor.has_value()) {
        torch::stable::copy_(
            preAllocatedOutputTensor.value(), cpuFrameOutput.data);
        frameOutput.data = preAllocatedOutputTensor.value();
      } else {
        frameOutput.data = torch::stable::to(cpuFrameOutput.data, device_);
      }
      if (rotation_ != Rotation::NONE) {
        applyRotation(frameOutput, preAllocatedOutputTensor);
      }
      return;
    }
  }

  // Capture original dimensions before transferCpuFrameToGpu()
  // may round them up to even.
  FrameDims originalDims(avFrame->height, avFrame->width);

  UniqueAVFrame gpuFrame;
  if (cpuFallback_) {
    AVPixelFormat targetPixFmt = (outputDtype_ == OutputDtype::FLOAT32)
        ? AV_PIX_FMT_P016LE
        : AV_PIX_FMT_NV12;
    gpuFrame = transferCpuFrameToGpu(avFrame, targetPixFmt);
  } else {
    gpuFrame = std::move(avFrame);
  }

  STD_TORCH_CHECK(
      gpuFrame->format == AV_PIX_FMT_NV12 ||
          gpuFrame->format == AV_PIX_FMT_P016LE,
      "Expected NV12 or P016LE format frame");

  cudaStream_t nvdecStream = getCurrentCudaStream(device_.index());

  auto convertFrame = [&](std::optional<torch::stable::Tensor> preAlloc)
      -> torch::stable::Tensor {
    bool isP016 = (gpuFrame->format == AV_PIX_FMT_P016LE);
    int bitDepth = 8;
    if (isP016) {
      bitDepth = cpuFallback_
          ? codecContext_->bits_per_raw_sample
          : static_cast<int>(videoFormat_.bit_depth_luma_minus8) + 8;
    }
    return convertYUVFrameToRGB(
        gpuFrame,
        device_,
        nvdecStream,
        preAlloc,
        originalDims,
        isP016,
        bitDepth,
        cachedColorMatrix_);
  };

  if (rotation_ == Rotation::NONE) {
    validatePreAllocatedTensorShape(preAllocatedOutputTensor, originalDims);
    frameOutput.data = convertFrame(preAllocatedOutputTensor);
  } else {
    // preAllocatedOutputTensor has post-rotation dimensions, but the
    // conversion outputs pre-rotation dimensions, so we can't use it as the
    // conversion destination or validate it against the frame shape.
    // Once we support native transforms on the NVDEC CUDA interface,
    // rotation should be handled as part of the transform pipeline instead.
    frameOutput.data = convertFrame(/*preAlloc=*/std::nullopt);
    applyRotation(frameOutput, preAllocatedOutputTensor);
  }
}

void BetaCudaDeviceInterface::applyRotation(
    FrameOutput& frameOutput,
    std::optional<torch::stable::Tensor> preAllocatedOutputTensor) {
  int k = 0;
  switch (rotation_) {
    case Rotation::CCW90:
      k = 1;
      break;
    case Rotation::ROTATE180:
      k = 2;
      break;
    case Rotation::CW90:
      k = 3;
      break;
    default:
      STD_TORCH_CHECK(false, "Unexpected rotation value");
      break;
  }
  // Apply rotation using rot90 on the H and W dims of our HWC tensor.
  // stableRot90 returns a view, so we need to make it contiguous.
  frameOutput.data =
      torch::stable::contiguous(stableRot90(frameOutput.data, k, 0, 1));

  if (preAllocatedOutputTensor.has_value()) {
    torch::stable::copy_(preAllocatedOutputTensor.value(), frameOutput.data);
    frameOutput.data = preAllocatedOutputTensor.value();
  }
}

OutputDtype BetaCudaDeviceInterface::getPreAllocationDtype(
    OutputDtype requestedDtype) const {
  if (requestedDtype == OutputDtype::FLOAT32 &&
      surfaceFormat_ == cudaVideoSurfaceFormat_NV12) {
    return OutputDtype::UINT8;
  }
  return requestedDtype;
}

std::string BetaCudaDeviceInterface::getDetails() {
  std::string details = "NVDEC CUDA Device Interface.";
  if (cpuFallback_) {
    details += " Using CPU fallback.";
    if (!nvcuvidAvailable_) {
      details += " NVCUVID not available!";
    }
  } else {
    details += " Using NVDEC.";
  }
  return details;
}

} // namespace facebook::torchcodec
