#include "stdafx.h"
#include "MPScheduler.h"
#include "../Scheduler.h"
#include <vector>
#include "../../BWAPI_Addon.h"

#ifdef NON_FINAL_VERSION
#include "LogServer.h"
#include "StatsRecorder.h"
#endif

namespace MP
{
  class TaskScope : boost::noncopyable
  {
    MPScheduler& m_scheduler;
    Task* const m_task;
    const bool m_owning;
  public:
    TaskScope(MPScheduler& scheduler, Task * task, bool owning = true) : m_scheduler(scheduler), m_task(task), m_owning(owning)
    {
#ifdef NON_FINAL_VERSION
      TaskEvent ev(m_task);
      m_scheduler.onStartEvent(ev);
#endif
    }
    TaskScope(MPScheduler& scheduler, TaskPtr task) : m_scheduler(scheduler), m_task(task.get()), m_owning(false)
    {
#ifdef NON_FINAL_VERSION
      TaskEvent ev(m_task);
      m_scheduler.onStartEvent(ev);
#endif
    }
    ~TaskScope()
    {
#ifdef NON_FINAL_VERSION
      TaskEvent ev(m_task);
      m_scheduler.onEndEvent(ev);
#endif
      if(m_owning)
        delete m_task;
    }

    Task * operator->() const
    {
      return m_task;
    }
    Task & operator*() const
    {
      return *m_task;
    }
  };

#ifdef USE_MULTITHREADED_SCHEDULER

  class MainThreadTaskHandler : public SchedulerTask, boost::noncopyable
  {
    typedef std::vector<TaskPtr> TaskVector;
    TaskVector m_tasks;
    CriticalSection m_mutex;
    MP::MPScheduler& m_scheduler;
    ::Scheduler& m_owner;
  public:
    MainThreadTaskHandler(::Scheduler &owner, MP::MPScheduler& scheduler)
      : SchedulerTask(owner), m_scheduler(scheduler), m_owner(owner)
    {
      owner.activate(this, 1, 1, 1);
    }

    ~MainThreadTaskHandler()
    {
      m_owner.deactivate(this);
    }

    virtual void run()
    {
      TaskVector tasks;
      {
        CriticalSection::scoped_lock lock(m_mutex);
        m_tasks.swap(tasks);
      }
      for(TaskVector::iterator it = tasks.begin(); it != tasks.end(); ++it)
      {
        TaskScope task(m_scheduler, *it);
        task->execute();
      }
    }
    void schedule(TaskPtr task)
    {
      CriticalSection::scoped_lock lock(m_mutex);
      m_tasks.push_back(task);
    }
  protected:
    virtual void _output(std::ostream& out) const
    {
      out << "Main thread task";
    }
  };

  void MPScheduler::runWorker(boost::shared_ptr<volatile bool> isClosed)
  {
    // Main function for each worker thread

    bool lastWasSuccessful = true;
    while(lastWasSuccessful || !*isClosed)
    {
      Task * task;
      lastWasSuccessful = m_workerTasks.try_pop(task);
      if(lastWasSuccessful)
      {
        {
          CriticalSection::scoped_lock lock(m_mutex);
          ++m_nbActiveThreads;
        }

        TaskScope scope(*this, task);

        scope->execute();

        {
          CriticalSection::scoped_lock lock(m_mutex);
          --m_nbActiveThreads;
        }
      }
      Sleep(0);
    }
  }
   
  void MPScheduler::scheduleOnMainThread(TaskPtr task)
  {
    assert(m_mainThreadTasks);
    m_mainThreadTasks->schedule(task);
  }
#endif

  void MPScheduler::schedule(Task * task)
  {
#ifdef USE_MULTITHREADED_SCHEDULER
    if(m_workerThreads.size() == 0) // always empty if USE_MULTITHREADED_SCHEDULER is not defined
    {
      // No worker threads, run synchronously
      TaskScope scope(*this, task);
      scope->execute();
    }
    else
    {
      m_workerTasks.push(task);
    }
#else
    m_workerTasks.push(task);
#endif
  }

  MPScheduler::MPScheduler(::Scheduler* scheduler)
#ifdef USE_MULTITHREADED_SCHEDULER
    : m_scheduler(scheduler), m_mainThreadTasks(new MainThreadTaskHandler(*scheduler, *this)), m_isClosed(new bool(false)), m_nbActiveThreads(0)
#endif
  {
    std::vector<DWORD> ids;
#ifdef USE_MULTITHREADED_SCHEDULER
    const unsigned int cpuCount = Config::getInstance().getWorkerThreadCount();
    for(unsigned int i = 0; i < cpuCount; ++i)
    {
      boost::thread * th = m_workerThreads.create_thread(boost::bind(&MPScheduler::runWorker, this, m_isClosed));
      DWORD id = GetThreadId(th->native_handle());
      SetThreadName(id, "MP MPScheduler Worker");
      ids.push_back(id);
    }
#else
    unused(scheduler);
#endif

#if defined(_DEBUG) && defined(_BWAPI_MPSC)
    // See BWAPI-MPSC project
    BWAPI::Broodwar->setSCAiThreadsId(ids);
#endif

#ifdef NON_FINAL_VERSION
    if (Config::getInstance().startLogServerThread())
    {
      LogServer * srv = new LogServer(Config::getInstance().getLogServerPort());
      srv->open(BWAPI::getMainThreadId(), ids);
      addListener(srv);
    }
    if (Config::getInstance().startStatsThread())
    {
      StatsRecorder * rcrd = new StatsRecorder(ids.size() + 1); // +1 for main thread
      rcrd->start();
      addListener(rcrd);
    }
#endif
  }

  MPScheduler::~MPScheduler()
  { 
#ifdef USE_MULTITHREADED_SCHEDULER
    *m_isClosed = true;
    m_workerThreads.join_all();
    delete m_mainThreadTasks;
#endif
  }

#ifdef NON_FINAL_VERSION
   void MPScheduler::onStartEvent(const TaskEvent &ev)
   {
     std::for_each(m_listeners.begin(), m_listeners.end(), boost::bind(&ITaskListener::onTaskStarted, _1, ev));
   }
   void MPScheduler::onEndEvent(const TaskEvent &ev)
   {
     std::for_each(m_listeners.begin(), m_listeners.end(), boost::bind(&ITaskListener::onTaskEnded, _1, ev));
   }
#endif
}