//----------------------------------------------------------------------
// Copyright (C) 2006-2007 Electronic Arts    
//----------------------------------------------------------------------

#include "EABase/eabase.h"

#if ! defined EA_PLATFORM_PS3_SPU

#include "job_util/parallel_for_each.h"
#include "job_manager/atomic.h"
#include "job_util/batch_desc.h"
#include "job_util/image_desc.h"
#include "job_manager/allocator.h"
#include "job_manager/atomic_allocator.h"
#include "job_manager/job_scheduler.h"
#include "job_manager/version.h"
#include <string.h>
#include <new>
#include "eathread/eathread_futex.h"

namespace EA {
    namespace Jobs {

namespace ForEachDescManager {
    EA_COMPILETIME_ASSERT( sizeof(BatchDesc) <= ForEachDesc::sLargeEnoughToHoldAllKnownDescTypes );
    EA_COMPILETIME_ASSERT( sizeof(ImageDesc) <= ForEachDesc::sLargeEnoughToHoldAllKnownDescTypes );

    static AtomicAllocator* aa = 0;
    static EA::Thread::Futex& GetAAFEDFutex()
    {
        static EA::Thread::Futex aaFutex;
        return aaFutex;
    }

    static void Setup()
    {
        EA::Thread::AutoFutex af(GetAAFEDFutex());
        if (!aa)
        {
            aa = (AtomicAllocator*)EA::Jobs::GetAllocator().Alloc(sizeof(AtomicAllocator), "JUFD", EA::Allocator::MEM_PERM, EA_ALIGN_OF(AtomicAllocator));
            new (aa) AtomicAllocator(
                &EA::Jobs::GetAllocator(),
#if JM_VERSION >= JM_CREATE_VERSION_NUMBER(3,2,1)
                4,
                4,
#else
                4096,
                4096,
#endif
                ForEachDesc::sLargeEnoughToHoldAllKnownDescTypes,
                16,
				JM_OFFSET_OF(ForEachDesc, mNext),
				"EA::Jobs::AtomicAllocator::(ForEachDescManager)"
            );
        }
    }
    static void Teardown()
    {
        EA::Thread::AutoFutex af(GetAAFEDFutex());
        if (aa)
        {
            aa->~AtomicAllocator();
            EA::Jobs::GetAllocator().Free(aa);
            aa = 0;
        }
    }

    static ForEachDesc * Get()
    {
        if (!aa) Setup();

        ForEachDesc * d = aa->Alloc<ForEachDesc>();
        d->mAtomicAllocator = (uintptr_t)aa;
        return d;
    }

    static uint32_t Used()
    {
        if (!aa) return 0;
        return aa->Used();
    }
}

size_t ParallelForEach::UsedDescriptors()
{
    return ForEachDescManager::Used();
}

void ParallelForEach::Setup() { ForEachDescManager::Setup(); }
void ParallelForEach::Teardown() { ForEachDescManager::Teardown(); }

//----------------------------------------------------------------------
ParallelForEach::ParallelForEach(JobScheduler& jobScheduler, const EntryPoint& entryPoint, const ForEachDesc& forEachDesc, int numJobs, ForEachDivisionEnum forEachDivisionEnum)
  : mJobScheduler(&jobScheduler),
    mEntryPoint(entryPoint),
    mDesc(0),
    mForEachDivisionEnum(forEachDivisionEnum),
    mNumJobs(numJobs),
    mParam2(0),
    mParam3(0),
    mAllowSleepOn(false)
{
    SetDesc( forEachDesc );
}

//----------------------------------------------------------------------
ParallelForEach::ParallelForEach()
  : mJobScheduler(0),
    mDesc(0),
    mForEachDivisionEnum(FOR_EACH_DIVISION_DYNAMIC),
    mNumJobs(0),
    mParam2(0),
    mParam3(0),
    mAllowSleepOn(false)
{
}

//----------------------------------------------------------------------
void ParallelForEach::SetJobScheduler(JobScheduler & jobScheduler)
{
    mJobScheduler = &jobScheduler;
}

//----------------------------------------------------------------------
JobScheduler & ParallelForEach::GetJobScheduler() const
{
    EA_ASSERT(mJobScheduler);
    return *mJobScheduler;
}

//----------------------------------------------------------------------
void ParallelForEach::SetEntryPoint(EntryPoint const& entryPoint)
{
    mEntryPoint = entryPoint;
}

//----------------------------------------------------------------------
EntryPoint ParallelForEach::GetEntryPoint() const
{
    return mEntryPoint;
}

//----------------------------------------------------------------------
void ParallelForEach::SetDesc(ForEachDesc const& forEachDesc)
{
    EA_ASSERT(!mDesc);
    mDesc = ForEachDescManager::Get();
    mDesc->Copy(forEachDesc);
}

//----------------------------------------------------------------------
void ParallelForEach::SetDescUnmanaged(ForEachDesc * forEachDesc)
{
    EA_ASSERT(!mDesc);
    forEachDesc->mRefCounter = 0;
    forEachDesc->mAtomicAllocator = 0;
    forEachDesc->mNext = ~uintptr_t(0);
    mDesc = forEachDesc;
}

//----------------------------------------------------------------------
void ParallelForEach::SetNumJobs(int numJobs)
{
    mNumJobs = numJobs;
}

//----------------------------------------------------------------------
int ParallelForEach::GetNumJobs() const
{
    return mNumJobs;
}

//----------------------------------------------------------------------
void ParallelForEach::Run()
{
    AddJobInstances();
    mBeginJobInstanceHandle.Enable();
}

//----------------------------------------------------------------------
void ParallelForEach::AddJobInstances()
{
    EA_ASSERT(mDesc);
    EA_ASSERT(mJobScheduler);
    EA_ASSERT(mNumJobs>0);
    EA_ASSERT(!mEntryPoint.IsNull());

    mDesc->InitDynamicCounter();
    mDesc->InitRefCounter(static_cast<uint32_t>(mNumJobs));

    // Highjack the descriptor to contain the division params so we only do one DMA
    mDesc->SetDivision(mForEachDivisionEnum, mNumJobs);

    if (mNumJobs == 1)
    {
        // avoid adding the extra gather job to the queue
        mEntryPoint.SetAllowSleepOn(mAllowSleepOn);

        mBeginJobInstanceHandle = mJobScheduler->AddNotReady(
            mEntryPoint,
            reinterpret_cast<uintptr_t>(mDesc),
            static_cast<uintptr_t>(0),
            mParam2,
            mParam3
        );
        mEndJobInstanceHandle = mBeginJobInstanceHandle;
        return;
    }

    // Force user supplied EP to not use SleepOn since we don't SleepOn those jobs.
    // This is just in case the caller thinks they need to set SleepOn for their EP.
    mEntryPoint.SetAllowSleepOn(false);

    EntryPoint gatherEp = mEntryPoint;

    gatherEp.SetCode(0);
    gatherEp.SetCode(mEntryPoint.GetEnvironment(), 0, 0);

    gatherEp.SetAllowSleepOn(mAllowSleepOn);

    mEndJobInstanceHandle = mJobScheduler->AddNotReady(gatherEp);

    JobInstanceHandle prevJih;

    for (int i = 0; i < mNumJobs; ++i)
    {
        JobInstanceHandle jih = mJobScheduler->AddNotReady(mEntryPoint, reinterpret_cast<uintptr_t>(mDesc), static_cast<uintptr_t>(i), mParam2, mParam3);

        jih.AddEvent(mEndJobInstanceHandle.AddBarrier(), Event::EVENT_WHEN_JOB_END);

        if (i == 0)
            mBeginJobInstanceHandle = jih;
        else
            prevJih.AddEvent(jih.AddBarrier(), Event::EVENT_WHEN_JOB_BEGIN);

        prevJih = jih;
    }
}


//----------------------------------------------------------------------
void ParallelForEach::WaitOn(WaitOnCallback* waitOnCallback, void* callbackContext, int sleepInMs) const
{
    mEndJobInstanceHandle.WaitOn(waitOnCallback, callbackContext, sleepInMs);
}

//----------------------------------------------------------------------
void ParallelForEach::SleepOn() const
{
    mEndJobInstanceHandle.SleepOn();
}

//----------------------------------------------------------------------
bool ParallelForEach::IsDone() const
{
    return mEndJobInstanceHandle.IsDone();
}

//----------------------------------------------------------------------
void ParallelForEach::SetAllowSleepOn(bool b)
{
    mAllowSleepOn = b;
}

//----------------------------------------------------------------------
bool ParallelForEach::GetAllowSleepOn() const
{
    return mAllowSleepOn;
}

//----------------------------------------------------------------------
void ParallelForEach::SetDivision(ForEachDivisionEnum forEachDivisionEnum)
{
    mForEachDivisionEnum = forEachDivisionEnum;
}

//----------------------------------------------------------------------
ForEachDivisionEnum ParallelForEach::GetDivision() const
{
    return mForEachDivisionEnum;
}

//----------------------------------------------------------------------
void ParallelForEach::SetParam2(Param param2)
{
    mParam2 = param2;
}

//----------------------------------------------------------------------
Param ParallelForEach::GetParam2() const
{
    return mParam2;
}

//----------------------------------------------------------------------
void ParallelForEach::SetParam3(Param param3)
{
    mParam3 = param3;
}

//----------------------------------------------------------------------
Param ParallelForEach::GetParam3() const
{
    return mParam3;
}

//----------------------------------------------------------------------
JobInstanceHandle ParallelForEach::GetBeginJobInstanceHandle() const
{
    return mBeginJobInstanceHandle;
}

//----------------------------------------------------------------------
JobInstanceHandle ParallelForEach::GetEndJobInstanceHandle() const
{
    return mEndJobInstanceHandle;
}

}}

#endif
