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

// pybind11 headers from torch error out if TORCH_TARGET_VERSION is defined,
// so we temporarily undefine it.
// See https://github.com/pytorch/pytorch/pull/174372 for context
#pragma push_macro("TORCH_TARGET_VERSION")
#undef TORCH_TARGET_VERSION
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#pragma pop_macro("TORCH_TARGET_VERSION")

#include "AVIOContextHolder.h"

namespace py = pybind11;

namespace facebook::torchcodec {

// Enables uers to pass in a Python file-like object. We then forward all read
// and seek calls back up to the methods on the Python object.
class AVIOFileLikeContext : public AVIOContextHolder {
 public:
  explicit AVIOFileLikeContext(const py::object& fileLike, bool isForWriting);

  int read(uint8_t* buf, int size) override;
  int write(const uint8_t* buf, int size) override;
  int64_t seek(int64_t offset, int whence) override;
  int64_t getSize() override;

 private:
  // Note that we dynamically allocate the Python object because we need to
  // strictly control when its destructor is called. We must hold the GIL
  // when its destructor gets called, as it needs to update the reference
  // count. It's easiest to control that when it's dynamic memory. Otherwise,
  // we'd have to ensure whatever enclosing scope holds the object has the GIL,
  // and that's, at least, hard. For all of the common pitfalls, see:
  //
  //   https://pybind11.readthedocs.io/en/stable/advanced/misc.html#common-sources-of-global-interpreter-lock-errors
  //
  // We maintain a reference to the file-like object because the file-like
  // object that was created on the Python side must live as long as our
  // potential use. That is, even if there are no more references to the object
  // on the Python side, we require that the object is still live.
  struct PyObjectDeleter {
    inline void operator()(py::object* obj) const {
      if (obj) {
        py::gil_scoped_acquire gil;
        delete obj;
      }
    }
  };

  using UniquePyObject = std::unique_ptr<py::object, PyObjectDeleter>;
  UniquePyObject fileLike_;
};

} // namespace facebook::torchcodec
