// 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 <optional>
#include <string>
#include <vector>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/pixfmt.h>
#include <libavutil/rational.h>
}

namespace facebook::torchcodec {

enum class SeekMode { exact, approximate, custom_frame_mappings };

struct StreamMetadata {
  // Common (video and audio) fields derived from the AVStream.
  int streamIndex = -1;

  // See this link for what various values are available:
  // https://ffmpeg.org/doxygen/trunk/group__lavu__misc.html#ga9a84bba4713dfced21a1a56163be1f48
  AVMediaType mediaType = AVMEDIA_TYPE_UNKNOWN;

  std::optional<AVCodecID> codecId;
  std::optional<std::string> codecName;
  std::optional<double> durationSecondsFromHeader;
  std::optional<double> beginStreamSecondsFromHeader;
  std::optional<int64_t> numFramesFromHeader;
  std::optional<int64_t> numKeyFrames;
  std::optional<double> averageFpsFromHeader;
  std::optional<double> bitRate;

  // Used as fallback in approximate mode when stream duration is unavailable.
  std::optional<double> durationSecondsFromContainer;

  // More accurate duration, obtained by scanning the file.
  // These presentation timestamps are in time base.
  std::optional<int64_t> beginStreamPtsFromContent;
  std::optional<int64_t> endStreamPtsFromContent;

  // These presentation timestamps are in seconds.
  std::optional<double> beginStreamPtsSecondsFromContent;
  std::optional<double> endStreamPtsSecondsFromContent;

  // This can be useful for index-based seeking.
  std::optional<int64_t> numFramesFromContent;

  // Video-only fields
  // Post-rotation dimensions
  std::optional<int> postRotationWidth;
  std::optional<int> postRotationHeight;
  std::optional<AVRational> sampleAspectRatio;
  // Rotation angle in degrees from display matrix, in the range [-180, 180].
  std::optional<double> rotation;
  std::optional<AVColorPrimaries> colorPrimaries;
  std::optional<AVColorSpace> colorSpace;
  std::optional<AVColorTransferCharacteristic> colorTransferCharacteristic;
  // The pixel format of the encoded video, e.g. "yuv420p".
  std::optional<std::string> pixelFormat;

  // Audio-only fields
  std::optional<int64_t> sampleRate;
  std::optional<int64_t> numChannels;
  std::optional<std::string> sampleFormat;

  // Computed methods with fallback logic
  std::optional<double> getDurationSeconds(SeekMode seekMode) const;
  double getBeginStreamSeconds(SeekMode seekMode) const;
  std::optional<double> getEndStreamSeconds(SeekMode seekMode) const;
  std::optional<int64_t> getNumFrames(SeekMode seekMode) const;
  std::optional<double> getAverageFps(SeekMode seekMode) const;

  // Color metadata name accessors. These return nullopt if the field is unset
  // or if FFmpeg returns NULL for the name.
  std::optional<std::string> getColorPrimariesName() const;
  std::optional<std::string> getColorSpaceName() const;
  std::optional<std::string> getColorTransferCharacteristicName() const;
};

struct ContainerMetadata {
  std::vector<StreamMetadata> allStreamMetadata;
  int numAudioStreams = 0;
  int numVideoStreams = 0;

  // Note that this is the container-level duration, which is usually the max
  // of all stream durations available in the container.
  std::optional<double> durationSecondsFromHeader;

  // Total BitRate level information at the container level in bit/s
  std::optional<double> bitRate;

  // If set, this is the index to the default audio stream.
  std::optional<int> bestAudioStreamIndex;

  // If set, this is the index to the default video stream.
  std::optional<int> bestVideoStreamIndex;
};

} // namespace facebook::torchcodec
