/*=========================================================================

  Program:   ParaView
  Module:    vtkPythonRunLoop.cxx

  Copyright (c) Kitware, Inc.
  All rights reserved.
  See Copyright.txt or http://www.paraview.org/HTML/Copyright.html for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/

#include "vtkPythonRunLoop.h"
#include "vtkLogger.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPythonInterpreter.h"
#include "vtkSmartPyObject.h"

#include <chrono>
#include <thread>

namespace
{
static struct PyMethodDef schedule_callback_descr = { "Paraview_on_event_scheduled",
  (PyCFunction)vtkPythonRunLoop::ScheduleCallback, METH_NOARGS, nullptr };

bool EnsureThread(std::thread::id owner)
{
  if (owner != std::this_thread::get_id())
  {
    vtkLogF(ERROR, "Runloop accessed by wrong thread !");
    return false;
  }
  return true;
}

bool CheckAndClearError()
{
  if (PyObject* exception = PyErr_Occurred())
  {
    PyErr_Print();
    PyErr_Clear();
    return true;
  }
  return false;
}
}

class vtkPythonRunLoop::vtkInternals
{
public:
  const std::thread::id OwnerTID{ std::this_thread::get_id() };

  // the event loop of python
  vtkSmartPyObject AsyncioLoop;

  // Python object that corresponds to vtkPythonRunLoop::ScheduleCallback
  vtkSmartPyObject PyScheduleCallback;
  // Convenience method to call "asyncio.call_later"
  vtkSmartPyObject PyCallLater;
  // Convenience method to call "asyncio.call_soon_threadsafe"
  vtkSmartPyObject PyCallSoonThreadsafe;

  rxcpp::schedulers::run_loop rxcpp_run_loop;

  void Initialize(PyObject* asyncioLoop)
  {
    vtkPythonScopeGilEnsurer gilEnsurer;
    this->AsyncioLoop = asyncioLoop;
    rxcpp_run_loop.set_notify_earlier_wakeup(
      [this](auto const& when) { this->on_earlier_wakeup(this->seconds_until(when)); });

    this->PyScheduleCallback = PyCFunction_New(&schedule_callback_descr, NULL);
    this->PyCallLater = PyObject_GetAttrString(this->AsyncioLoop, "call_later");
    this->PyCallSoonThreadsafe = PyObject_GetAttrString(this->AsyncioLoop, "call_soon_threadsafe");
  }

  // to be called after an event is pushed into the eventloop
  void on_earlier_wakeup(double delayInSeconds)
  {

    vtkPythonScopeGilEnsurer gilEnsurer;
    PyObject* pyDelayInSeconds = PyFloat_FromDouble(delayInSeconds);
    if (OwnerTID == std::this_thread::get_id())
    {
      vtkSmartPyObject result = PyObject_CallFunctionObjArgs(
        PyCallLater.GetPointer(), pyDelayInSeconds, PyScheduleCallback.GetPointer(), nullptr);
      if (result == nullptr)
      {
        CheckAndClearError();
        throw std::runtime_error("Encountered error calling asyncio.call_later");
      }
    }
    else
    {
      vtkSmartPyObject result =
        PyObject_CallFunctionObjArgs(this->PyCallSoonThreadsafe.GetPointer(),
          PyCallLater.GetPointer(), pyDelayInSeconds, PyScheduleCallback.GetPointer(), nullptr);
      if (result == nullptr)
      {
        CheckAndClearError();
        throw std::runtime_error("Encountered error calling asyncio.call_soon_threadsafe");
      }
    }
    // TODO by holding to `result` we can cancel tasks. Do we need this ?
  }

  // Flush the RxCpp run loop
  void on_event_scheduled()
  {
    // Dispatch outstanding RxCpp events
    while (!rxcpp_run_loop.empty() && rxcpp_run_loop.peek().when < rxcpp_run_loop.now())
    {
      rxcpp_run_loop.dispatch();
    }
    // If there are outstanding events, set the timer to wakeup for the first one
    if (!rxcpp_run_loop.empty())
    {
      const double seconds_till_next_event = this->seconds_until(rxcpp_run_loop.peek().when);
      this->on_earlier_wakeup(seconds_till_next_event);
    }
  }
  // Calculate seconds from now until `when`
  double seconds_until(rxcpp::schedulers::run_loop::clock_type::time_point const& when) const
  {
    const auto duration = std::chrono::duration<double>(when - this->rxcpp_run_loop.now());
    return std::max(duration.count(), 0.0);
  }
};

vtkPythonRunLoop* vtkPythonRunLoop::Singleton = nullptr;

vtkStandardNewMacro(vtkPythonRunLoop);

//----------------------------------------------------------------------------
vtkPythonRunLoop::vtkPythonRunLoop()
  : Internals(new vtkPythonRunLoop::vtkInternals)
{
  vtkPythonRunLoop::Singleton = this;
}

//----------------------------------------------------------------------------
vtkPythonRunLoop* vtkPythonRunLoop::GetInstance()
{
  return vtkPythonRunLoop::Singleton;
}

//----------------------------------------------------------------------------
vtkPythonRunLoop::~vtkPythonRunLoop()
{
  vtkPythonRunLoop::Singleton = nullptr;
}

//----------------------------------------------------------------------------
bool vtkPythonRunLoop::IsAsyncInitialized()
{
  if (!vtkPythonInterpreter::IsInitialized())
  {
    vtkLogF(WARNING, "Python interpreter is not initialized !");
    return false;
  }

  vtkPythonScopeGilEnsurer gilEnsurer;
  vtkSmartPyObject moduleName = PyUnicode_FromString("asyncio");

  vtkSmartPyObject asyncioModule = PyImport_GetModule(moduleName);
  if (!asyncioModule)
  {
    vtkLogF(WARNING, "asyncio is not loaded !");
    return false;
  }
  // Get current loop
  vtkSmartPyObject asyncioLoop = PyObject_CallMethod(asyncioModule, "get_running_loop", nullptr);
  if (!asyncioLoop)
  {
    vtkLogF(WARNING, "No asyncio running loop !\n Please run inside an async def function");
    return false;
  }

  return true;
}

//----------------------------------------------------------------------------
PyObject* vtkPythonRunLoop::CreateFuture()
{
  auto& internals = *(this->Internals);
  if (!EnsureThread(internals.OwnerTID))
  {
    return nullptr;
  }
  vtkPythonScopeGilEnsurer gilEnsurer;
  PyObject* pythonFuture = PyObject_CallMethod(internals.AsyncioLoop, "create_future", nullptr);
  if (pythonFuture == nullptr)
  {
    CheckAndClearError();
    return nullptr;
  }

  return pythonFuture;
}

//----------------------------------------------------------------------------
PyObject* vtkPythonRunLoop::CreateAsyncIterator()
{
  auto& internals = *(this->Internals);
  if (!EnsureThread(internals.OwnerTID))
  {
    return nullptr;
  }
  // TODO cache modules to avoid re-importing
  vtkPythonScopeGilEnsurer gilEnsurer;
  vtkSmartPyObject moduleName = PyUnicode_FromString("asyncio");
  vtkSmartPyObject asyncioModule = PyImport_GetModule(moduleName);
  vtkSmartPyObject queue = PyObject_CallMethod(asyncioModule, "Queue", nullptr);

  vtkSmartPyObject future = this->CreateFuture();
  vtkSmartPyObject paratModule = PyImport_ImportModule("parat.utils");
  if (paratModule == nullptr)
  {
    CheckAndClearError();
    return nullptr;
  }

  PyObject* iterator = PyObject_CallMethod(
    paratModule, "QueueIterator", "OO", queue.GetPointer(), future.GetPointer());

  if (iterator == nullptr)
  {
    CheckAndClearError();
  }
  return iterator;
}

//----------------------------------------------------------------------------
void vtkPythonRunLoop::AcquireRunningLoopFromPython()
{
  if (!this->IsAsyncInitialized())
  {
    throw std::runtime_error("Async module is not initialized");
  }
  vtkPythonScopeGilEnsurer gilEnsurer;
  auto& internals = *(this->Internals);
  vtkSmartPyObject moduleName = PyUnicode_FromString("asyncio");

  vtkSmartPyObject asyncioModule = PyImport_GetModule(moduleName);
  // Get current loop
  vtkSmartPyObject asyncioLoop = PyObject_CallMethod(asyncioModule, "get_running_loop", nullptr);
  internals.Initialize(asyncioLoop);
}

//----------------------------------------------------------------------------
bool vtkPythonRunLoop::HasRunningLoop()
{
  return this->Internals->AsyncioLoop != nullptr;
}

//----------------------------------------------------------------------------
PyObject* vtkPythonRunLoop::ScheduleCallback()
{

  // There are (at least) two locks involved during scheduling: One from rxcpp
  // (lets call it R) which is held during `on_earlier_wakeup` when called from
  // `rxcpp::run_loop_scheduler::schedule` and the GIL of python.  During the
  // execution of `ScheduleCallback` the GIL is hold because we call it through
  // Python (via `call_later` and `call_soon_threadsafe`).  Note that
  // `rxcpp::run_loop::peek()/empty()` require holding R.
  // So, having a thread (M) run `on_event_scheduled` while an other (O) running
  // on_earlier_wakeup will cause a deadlock since thread M has GIL and needs
  // R and thread `O`  has R  and needs GIL.
  // To resolve it we release the GIL here before calling on_event_scheduled.

  Py_BEGIN_ALLOW_THREADS;

  vtkPythonRunLoop::GetInstance()->Internals->on_event_scheduled();

  Py_END_ALLOW_THREADS;

  vtkPythonScopeGilEnsurer gilEnsurer;
  Py_RETURN_NONE;
}

//----------------------------------------------------------------------------
rxcpp::schedulers::scheduler vtkPythonRunLoop::get_scheduler() const
{
  auto& internals = *(this->Internals);
  return internals.rxcpp_run_loop.get_scheduler();
}

//----------------------------------------------------------------------------
rxcpp::observe_on_one_worker vtkPythonRunLoop::observe_on_run_loop() const
{
  auto& internals = *(this->Internals);
  return rxcpp::observe_on_run_loop(internals.rxcpp_run_loop);
}

//----------------------------------------------------------------------------
bool vtkPythonRunLoop::empty() const
{
  auto& internals = *(this->Internals);
  return internals.rxcpp_run_loop.empty();
}

//----------------------------------------------------------------------------
void vtkPythonRunLoop::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
}
