// Copyright 2023 DeepMind Technologies Limited
//
// 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.

#import <algorithm>
#import <atomic>
#import <cstdint>
#import <cstdlib>
#import <cstring>
#import <iostream>
#import <vector>

#import <dlfcn.h>
#import <mach-o/dyld.h>
#import <pthread.h>
#import <sys/resource.h>
#import <unistd.h>

#import <Cocoa/Cocoa.h>
#import <Python.h>

extern "C" {
extern char **environ;  // for execve

// Wrap Objective-C Cocoa calls into C-style functions with default visibility,
// so that we can dlsym and call them from Python via ctypes.
__attribute__((used)) void mjpython_hide_dock_icon() {
  [NSApp setActivationPolicy:NSApplicationActivationPolicyAccessory];
}
__attribute__((used)) void mjpython_show_dock_icon() {
  [NSApp setActivationPolicy:NSApplicationActivationPolicyRegular];
}
}  // extern "C"

// TODO(b/273744079): Remove Python 3.7 code after end-of-life (27 Jun 2023).
namespace {
// 16MiB is the default Python thread stack size on macOS as of Python 3.11
// https://bugs.python.org/issue18075
constexpr rlim_t kThreadStackSize = 0x1000000;

struct {
#define CPYTHON_FN(fname) decltype(&::fname) fname

#if PY_MINOR_VERSION >= 8
  CPYTHON_FN(Py_InitializeFromConfig);
  CPYTHON_FN(Py_RunMain);
  CPYTHON_FN(PyConfig_Clear);
  CPYTHON_FN(PyConfig_InitPythonConfig);
  CPYTHON_FN(PyConfig_SetBytesArgv);
#else
  CPYTHON_FN(Py_DecodeLocale);
  CPYTHON_FN(Py_Initialize);
  CPYTHON_FN(Py_Main);
  CPYTHON_FN(PyMem_RawFree);
  CPYTHON_FN(Py_SetProgramName);
#endif

  // go/keep-sorted start
  CPYTHON_FN(Py_FinalizeEx);
  CPYTHON_FN(PyGILState_Ensure);
  CPYTHON_FN(PyGILState_Release);
  CPYTHON_FN(PyRun_SimpleStringFlags);
  // go/keep-sorted end

#undef CPYTHON_FN
} cpython;

std::atomic_bool py_initialized = false;

struct Args {
  int argc;
  char** argv;
};

// The Python main thread (distinct from the macOS main thread, which executes the main function).
void* mjpython_pymain(void* vargs) {
  Args* args = static_cast<Args*>(vargs);
  PyGILState_STATE gil;

  // Initialize the Python interpreter.
#if PY_MINOR_VERSION >= 8
  PyConfig config;
  cpython.PyConfig_InitPythonConfig(&config);
  cpython.PyConfig_SetBytesArgv(&config, args->argc, args->argv);
  cpython.Py_InitializeFromConfig(&config);
  cpython.PyConfig_Clear(&config);
#else
  // Convert each argv to wchar_t* (needed for Py_Main).
  wchar_t** wargv = static_cast<wchar_t**>(std::calloc(args->argc, sizeof(wchar_t*)));
  for (int i = 0; i < args->argc; ++i) {
    wargv[i] = cpython.Py_DecodeLocale(args->argv[i], nullptr);
  }
  cpython.Py_SetProgramName(wargv[0]);
  cpython.Py_Initialize();
#endif

  // Set up the condition variable to pass control back to the macOS main thread.
  gil = cpython.PyGILState_Ensure();
  cpython.PyRun_SimpleStringFlags("import threading; cond = threading.Condition()", nullptr);
  py_initialized.store(true);

  // Wait until GLFW is initialized on macOS main thread, set up the queue and an atexit hook
  // to enqueue a termination flag upon exit.
  cpython.PyRun_SimpleStringFlags(R"(
import atexit

# The mujoco.viewer module should only be imported here after glfw.init() in the macOS main thread.
with cond:
  cond.wait()
import mujoco.viewer

# Similar to a queue.Queue(maxsize=1), but where only one active task is allowed at a time.
# With queue.Queue(1), another item is allowed to be enqueued before task_done is called.
class _MjPythonImpl(mujoco.viewer._MjPythonBase):

  # Termination statuses
  NOT_TERMINATED = 0
  TERMINATION_REQUESTED = 1
  TERMINATION_ACCEPTED = 2
  TERMINATED = 3

  def __init__(self):
    self._cond = threading.Condition()
    self._model_data = None
    self._termination = self.__class__.NOT_TERMINATED
    self._busy = False

  def launch_on_ui_thread(self, model, data):
    with self._cond:
      if self._busy or self._model_data is not None:
        raise RuntimeError('another MuJoCo viewer is already open')
      else:
        self._model_data = (model, data)
        self._cond.notify()

  def terminate(self):
    with self._cond:
      self._termination = self.__class__.TERMINATION_REQUESTED
      self._cond.notify()
      self._cond.wait_for(
          lambda: self._termination == self.__class__.TERMINATED)

  def get(self):
    with self._cond:
      self._cond.wait_for(
          lambda: self._model_data is not None or self._termination)

      if self._termination:
        if self._termination == self.__class__.TERMINATION_REQUESTED:
          self._termination = self.__class__.TERMINATION_ACCEPTED
        return None

      model_data = self._model_data
      self._busy = True
      self._model_data = None
      return model_data

  def done(self):
    with self._cond:
      self._busy = False
      if self._termination == self.__class__.TERMINATION_ACCEPTED:
        self._termination = self.__class__.TERMINATED
      self._cond.notify()


mujoco.viewer._MJPYTHON = _MjPythonImpl()
atexit.register(mujoco.viewer._MJPYTHON.terminate)
del _MjPythonImpl  # Don't pollute globals for user script.

with cond:
  cond.notify()
del cond  # Don't pollute globals for user script.
)", nullptr);

  // Run the Python interpreter main loop.
#if PY_MINOR_VERSION >= 8
  cpython.Py_RunMain();
#else
  cpython.Py_Main(args->argc, wargv);
#endif

  // Tear down the interpreter.
  cpython.Py_FinalizeEx();
#if PY_MINOR_VERSION < 8
  for (int i = 0; i < args->argc; ++i) {
    cpython.PyMem_RawFree(wargv[i]);
    wargv[i] = nullptr;
  }
  std::free(wargv);
  wargv = nullptr;
#endif
  return nullptr;
}
}  // namespace

int main(int argc, char** argv) {
  const char* libpython_path = getenv("MJPYTHON_LIBPYTHON");
  if (!libpython_path || !libpython_path[0]) {
    std::cerr << "This binary must be launched via the mjpython.py script.\n";
    return EXIT_FAILURE;
  }

  // Enlarge the stack if necessary to match what Python normally expects to have when launching
  // a new thread.
  rlimit stack;
  if (getrlimit(RLIMIT_STACK, &stack)) {
    std::cerr << "getrlimit failed to query stack size with error code " << errno << " ("
              << std::strerror(errno) << ")\n";
    std::cerr << "continuing anyway but crashes may occur if the stack is too small\n";
  } else if (stack.rlim_cur < kThreadStackSize && stack.rlim_cur < stack.rlim_max) {
    auto rlim_old = stack.rlim_cur;
    stack.rlim_cur = std::min(kThreadStackSize, stack.rlim_max);
    if (setrlimit(RLIMIT_STACK, &stack)) {
      std::cerr << "setrlimit failed to increase stack size with error code " << errno << " ("
                << std::strerror(errno) << ")\n";
      std::cerr << "continuing anyway with stack size " << rlim_old << " but crashes may occur\n";
    } else {
      // re-exec the binary so that the new stack size takes effect
      std::uint32_t path_size = 0;
      _NSGetExecutablePath(nullptr, &path_size);
      std::vector<char> path(path_size);
      if (_NSGetExecutablePath(path.data(), &path_size)) {
        std::cerr << "unexpected error from _NSGetExecutablePath, continuing anyway\n";
      } else {
        execve(path.data(), argv, environ);
      }
    }
  }

  // Resolve libpython at runtime to prevent linking against the wrong dylib. The correct libpython
  // path is passed from a Python trampoline script, which ran inside the desired interpreter and
  // exec'd this binary.
  void* libpython = dlopen(libpython_path, RTLD_NOW | RTLD_GLOBAL);
  if (!libpython) {
    std::cerr << "failed to dlopen path '" << libpython_path << "': " << dlerror() << "\n";
    return EXIT_FAILURE;
  }

  // Look up required CPython API functions from table of symbols already loaded into the process.
#define CPYTHON_INITFN(fname)                                                            \
  {                                                                                      \
    cpython.fname = reinterpret_cast<decltype(cpython.fname)>(dlsym(libpython, #fname)); \
    if (!cpython.fname) {                                                                \
      std::cerr << "failed to dlsym '" << #fname << "': " << dlerror() << "\n";          \
      return EXIT_FAILURE;                                                               \
    }                                                                                    \
  }

#if PY_MINOR_VERSION >= 8
  CPYTHON_INITFN(Py_InitializeFromConfig);
  CPYTHON_INITFN(Py_RunMain);
  CPYTHON_INITFN(PyConfig_Clear);
  CPYTHON_INITFN(PyConfig_InitPythonConfig);
  CPYTHON_INITFN(PyConfig_SetBytesArgv);
#else
  CPYTHON_INITFN(Py_DecodeLocale);
  CPYTHON_INITFN(Py_Initialize);
  CPYTHON_INITFN(Py_Main);
  CPYTHON_INITFN(PyMem_RawFree);
  CPYTHON_INITFN(Py_SetProgramName);
#endif

  // go/keep-sorted start
  CPYTHON_INITFN(Py_FinalizeEx);
  CPYTHON_INITFN(PyGILState_Ensure);
  CPYTHON_INITFN(PyGILState_Release);
  CPYTHON_INITFN(PyRun_SimpleStringFlags);
  // go/keep-sorted end

#undef CPYTHON_INITFN

  // Package up argc and argv together to pass to pthread_create.
  Args args{argc, argv};

#define PTHREAD_CHECKED(func, ...)                                                              \
  {                                                                                             \
    int result = func(__VA_ARGS__);                                                             \
    if (result) {                                                                               \
      std::cerr << #func << " failed with " << result << "(" << std::strerror(result) << ")\n"; \
      return EXIT_FAILURE;                                                                      \
    }                                                                                           \
  }

  // Configure the new thread with the correct stack size;
  pthread_attr_t pthread_attr;
  PTHREAD_CHECKED(pthread_attr_init, &pthread_attr);
  PTHREAD_CHECKED(pthread_attr_setstacksize, &pthread_attr, stack.rlim_cur);

  // Create a thread to be used as the "Python main thread".
  pthread_t pymain_thread;
  PTHREAD_CHECKED(pthread_create, &pymain_thread, &pthread_attr, &mjpython_pymain, &args);
  pthread_attr_destroy(&pthread_attr);
#undef PTHREAD_CHECKED

  // Busy-wait until Python interpreter is initialized.
  while (!py_initialized.load()) {}

  // Initialize GLFW on the macOS main thread, yield control to Python main thread and wait for it
  // to finish setting up _MJPYTHON, then serve incoming viewer launch requests.
  PyGILState_STATE gil = cpython.PyGILState_Ensure();
  cpython.PyRun_SimpleStringFlags(R"(
import ctypes

# GLFW must be initialized on the OS main thread (i.e. here).
import glfw
import mujoco.viewer

glfw.init()
glfw.poll_events()
ctypes.CDLL(None).mjpython_hide_dock_icon()

# Wait for Python main thread to finish setting up _MJPYTHON
with cond:
  cond.notify()
  cond.wait()

while True:
  try:
    # Wait for an incoming payload.
    payload = mujoco.viewer._MJPYTHON.get()

    # None means that we are exiting.
    if payload is None:
      glfw.terminate()
      break

    # Otherwise, launch the viewer.
    model, data = payload
    ctypes.CDLL(None).mjpython_show_dock_icon()
    mujoco.viewer._launch_internal(model, data, run_physics_thread=False)
    ctypes.CDLL(None).mjpython_hide_dock_icon()

  finally:
    mujoco.viewer._MJPYTHON.done()
)", nullptr);
  cpython.PyGILState_Release(gil);

  // Tear everything down.
  pthread_join(pymain_thread, nullptr);
  dlclose(libpython);
}