#include "rar.hpp"


#ifdef RAR_SMP
#include "threadmisc.cpp"

#ifdef _WIN_ALL
int ThreadPool::ThreadPriority=THREAD_PRIORITY_NORMAL;
#endif

ThreadPool::ThreadPool(uint MaxThreads)
{
  MaxAllowedThreads = MaxThreads;
  if (MaxAllowedThreads>MaxPoolThreads)
    MaxAllowedThreads=MaxPoolThreads;
  if (MaxAllowedThreads==0)
    MaxAllowedThreads=1;

  ThreadsCreatedCount=0;

  // If we have more threads than queue size, we'll hang on pool destroying,
  // not releasing all waiting threads.
  if (MaxAllowedThreads>ASIZE(TaskQueue))
    MaxAllowedThreads=ASIZE(TaskQueue);

  Closing=false;

  bool Success = CriticalSectionCreate(&CritSection);
#ifdef _WIN_ALL
  QueuedTasksCnt=CreateSemaphore(NULL,0,ASIZE(TaskQueue),NULL);
  NoneActive=CreateEvent(NULL,TRUE,TRUE,NULL);
  Success=Success && QueuedTasksCnt!=NULL && NoneActive!=NULL;
#elif defined(_UNIX)
  AnyActive = false;
  QueuedTasksCnt = 0;
  Success=Success && pthread_cond_init(&AnyActiveCond,NULL)==0 &&
          pthread_mutex_init(&AnyActiveMutex,NULL)==0 &&
          pthread_cond_init(&QueuedTasksCntCond,NULL)==0 &&
          pthread_mutex_init(&QueuedTasksCntMutex,NULL)==0;
#endif
  if (!Success)
  {
    ErrHandler.GeneralErrMsg(L"\nThread pool initialization failed.");
    ErrHandler.Exit(RARX_FATAL);
  }

  QueueTop = 0;
  QueueBottom = 0;
  ActiveThreads = 0;
}


ThreadPool::~ThreadPool()
{
  WaitDone();
  Closing=true;

#ifdef _WIN_ALL
  ReleaseSemaphore(QueuedTasksCnt,ASIZE(TaskQueue),NULL);
#elif defined(_UNIX)
  // Threads still can access QueuedTasksCnt for a short time after WaitDone(),
  // so lock is required. We would occassionally hang without it.
  pthread_mutex_lock(&QueuedTasksCntMutex);
  QueuedTasksCnt+=ASIZE(TaskQueue);
  pthread_mutex_unlock(&QueuedTasksCntMutex);

  pthread_cond_broadcast(&QueuedTasksCntCond);
#endif

  for(uint I=0;I<ThreadsCreatedCount;I++)
  {
#ifdef _WIN_ALL
    // Waiting until the thread terminates.
    CWaitForSingleObject(ThreadHandles[I]);
#endif
    // Close the thread handle. In Unix it results in pthread_join call,
    // which also waits for thread termination.
    ThreadClose(ThreadHandles[I]);
  }

  CriticalSectionDelete(&CritSection);
#ifdef _WIN_ALL
  CloseHandle(QueuedTasksCnt);
  CloseHandle(NoneActive);
#elif defined(_UNIX)
  pthread_cond_destroy(&AnyActiveCond);
  pthread_mutex_destroy(&AnyActiveMutex);
  pthread_cond_destroy(&QueuedTasksCntCond);
  pthread_mutex_destroy(&QueuedTasksCntMutex);
#endif
}


void ThreadPool::CreateThreads()
{
#ifdef WIN32_CPU_GROUPS
  // 2024.12.28: Implement processor group support for pre-Windows 11 systems
  // with number of CPUs exceeding 64. For example, for 72 CPU the single
  // processor group size would be 36 and this is what RAR would use without
  // processor group support.

  uint GroupCount=0;
  uint CurGroupNumber=(uint)-1; // We'll increment it to 0 later.
  uint CurGroupSize=0,CumulativeGroupSize=0;

  typedef DWORD (WINAPI *GETACTIVEPROCESSORCOUNT)(WORD GroupNumber);
  GETACTIVEPROCESSORCOUNT pGetActiveProcessorCount=nullptr;
  typedef BOOL (WINAPI *GETTHREADGROUPAFFINITY)(HANDLE hThread,PGROUP_AFFINITY GroupAffinity);
  GETTHREADGROUPAFFINITY pGetThreadGroupAffinity=nullptr;
  typedef BOOL (WINAPI *SETTHREADGROUPAFFINITY)(HANDLE hThread,const GROUP_AFFINITY *GroupAffinity,PGROUP_AFFINITY PreviousGroupAffinity);
  SETTHREADGROUPAFFINITY pSetThreadGroupAffinity=nullptr;

  // Doesn't seem to be needed in Windows 11 and newer.
  // https://learn.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-getprocessaffinitymask
  // "Starting with Windows 11 and Windows Server 2022, on a system with
  //  more than 64 processors, process and thread affinities span all
  //  processors in the system, across all processor groups, by default."
  if (!IsWindows11OrGreater())
  {
    HMODULE hKernel=GetModuleHandle(L"kernel32.dll");
    if (hKernel!=nullptr)
    {
      typedef WORD (WINAPI *GETACTIVEPROCESSORGROUPCOUNT)();
      GETACTIVEPROCESSORGROUPCOUNT pGetActiveProcessorGroupCount=(GETACTIVEPROCESSORGROUPCOUNT)GetProcAddress(hKernel,"GetActiveProcessorGroupCount");

      pGetActiveProcessorCount=(GETACTIVEPROCESSORCOUNT)GetProcAddress(hKernel,"GetActiveProcessorCount");
      pGetThreadGroupAffinity=(GETTHREADGROUPAFFINITY)GetProcAddress(hKernel,"GetThreadGroupAffinity");
      pSetThreadGroupAffinity=(SETTHREADGROUPAFFINITY)GetProcAddress(hKernel,"SetThreadGroupAffinity");

      if (pGetActiveProcessorCount!=nullptr && pGetActiveProcessorGroupCount!=nullptr &&
          pGetThreadGroupAffinity!=nullptr && pSetThreadGroupAffinity!=nullptr)
        GroupCount=pGetActiveProcessorGroupCount();
    }
  }
#endif

  for (uint I=0;I<MaxAllowedThreads;I++)
  {
    THREAD_HANDLE hThread=ThreadCreate(PoolThread, this);
    ThreadHandles[I] = hThread;
    ThreadsCreatedCount++;
#ifdef _WIN_ALL
#ifdef WIN32_CPU_GROUPS
    if (GroupCount>1) // If we have multiple processor groups.
    {
      if (I>=CumulativeGroupSize) // Filled the processor group, go to next.
      {
        if (++CurGroupNumber>=GroupCount)
        {
          // If we exceeded the group number, such as when user specified
          // -mt64 for lower core count, start assigning from beginning.
          CurGroupNumber=0;
          CumulativeGroupSize=0;
        }
        // Current group size.
        CurGroupSize=pGetActiveProcessorCount(CurGroupNumber);
        // Size of all preceding groups including the current.
        CumulativeGroupSize+=CurGroupSize;
      }
      GROUP_AFFINITY GroupAffinity;
      pGetThreadGroupAffinity(hThread,&GroupAffinity);

      // Since normally before Windows 11 all threads belong to same source
      // group, we could set this value only once. But we set it every time
      // in case we'll decide for some reason to use it to rearrange threads
      // from different source groups in Windows 11+.
      uint SrcGroupSize=pGetActiveProcessorCount(GroupAffinity.Group);

      // Shifting by 64 would be the undefined behavior, so we treat 64 separately.
      KAFFINITY SrcGroupMask=(KAFFINITY)(SrcGroupSize==64 ? (uint64)0xffffffffffffffff:(uint64(1)<<SrcGroupSize)-1);

      // Here we check that process affinity for existing thread group
      // matches the entire group size. If user limited the process
      // affinity, we prefer to not extend the process to other groups,
      // because user might want to restrict the resource usage.
      // Also if source processor group is larger than required number
      // of threads, we do not need to move threads between groups.
      if (SrcGroupSize!=0 && GroupAffinity.Mask==SrcGroupMask &&
          GroupAffinity.Group!=CurGroupNumber && SrcGroupSize<MaxAllowedThreads)
      {
        // Shifting by 64 would be the undefined behavior, so we treat 64 separately.
        KAFFINITY CurGroupMask=(KAFFINITY)(CurGroupSize==64 ? (uint64)0xffffffffffffffff:(uint64(1)<<CurGroupSize)-1);
        GroupAffinity.Mask=CurGroupMask; // Use the entire group.
        GroupAffinity.Group=CurGroupNumber;

        // Assign the thread to a new group.
        pSetThreadGroupAffinity(hThread,&GroupAffinity,NULL);
      }
    }
#endif
    
    // Set the thread priority if needed.
    if (ThreadPool::ThreadPriority!=THREAD_PRIORITY_NORMAL)
      SetThreadPriority(ThreadHandles[I],ThreadPool::ThreadPriority);
#endif
  }
}


NATIVE_THREAD_TYPE ThreadPool::PoolThread(void *Param)
{
  ((ThreadPool*)Param)->PoolThreadLoop();
  return 0;
}


void ThreadPool::PoolThreadLoop()
{
  QueueEntry Task;
  while (GetQueuedTask(&Task))
  {
    Task.Proc(Task.Param);
    
    CriticalSectionStart(&CritSection); 
    if (--ActiveThreads == 0)
    {
#ifdef _WIN_ALL
      SetEvent(NoneActive);
#elif defined(_UNIX)
      pthread_mutex_lock(&AnyActiveMutex);
      AnyActive=false;
      pthread_cond_signal(&AnyActiveCond);
      pthread_mutex_unlock(&AnyActiveMutex);
#endif
    }
    CriticalSectionEnd(&CritSection); 
  }
}


bool ThreadPool::GetQueuedTask(QueueEntry *Task)
{
#ifdef _WIN_ALL
  CWaitForSingleObject(QueuedTasksCnt);
#elif defined(_UNIX)
  pthread_mutex_lock(&QueuedTasksCntMutex);
  while (QueuedTasksCnt==0)
    cpthread_cond_wait(&QueuedTasksCntCond,&QueuedTasksCntMutex);
  QueuedTasksCnt--;
  pthread_mutex_unlock(&QueuedTasksCntMutex);
#endif

  if (Closing)
    return false;

  CriticalSectionStart(&CritSection); 

  *Task = TaskQueue[QueueBottom];
  QueueBottom = (QueueBottom + 1) % ASIZE(TaskQueue);

  CriticalSectionEnd(&CritSection); 

  return true;
}


// Add task to queue. We assume that it is always called from main thread,
// it allows to avoid any locks here. We process collected tasks only
// when WaitDone is called.
void ThreadPool::AddTask(PTHREAD_PROC Proc,void *Data)
{
  if (ThreadsCreatedCount == 0)
    CreateThreads();
  
  // If queue is full, wait until it is empty.
  if (ActiveThreads>=ASIZE(TaskQueue))
    WaitDone();

  TaskQueue[QueueTop].Proc = Proc;
  TaskQueue[QueueTop].Param = Data;
  QueueTop = (QueueTop + 1) % ASIZE(TaskQueue);
  ActiveThreads++;
}


// Start queued tasks and wait until all threads are inactive.
// We assume that it is always called from main thread, when pool threads
// are sleeping yet.
void ThreadPool::WaitDone()
{
  if (ActiveThreads==0)
    return;
#ifdef _WIN_ALL
  ResetEvent(NoneActive);
  ReleaseSemaphore(QueuedTasksCnt,ActiveThreads,NULL);
  CWaitForSingleObject(NoneActive);
#elif defined(_UNIX)
  AnyActive=true;

  // Threads reset AnyActive before accessing QueuedTasksCnt and even
  // preceding WaitDone() call does not guarantee that some slow thread
  // is not accessing QueuedTasksCnt now. So lock is necessary.
  pthread_mutex_lock(&QueuedTasksCntMutex);
  QueuedTasksCnt+=ActiveThreads;
  pthread_mutex_unlock(&QueuedTasksCntMutex);

  pthread_cond_broadcast(&QueuedTasksCntCond);

  pthread_mutex_lock(&AnyActiveMutex);
  while (AnyActive)
    cpthread_cond_wait(&AnyActiveCond,&AnyActiveMutex);
  pthread_mutex_unlock(&AnyActiveMutex);
#endif
}
#endif // RAR_SMP