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

#include "job_util/batch_processor.h"
#include "job_manager/atomic.h"
#include "job_manager/atomic_allocator.h"

#include "job_util/align.h"

#ifdef EA_PLATFORM_PS3_SPU
    #include <cell/dma.h>
    #include <spu_printf.h>
#else
    #include <string.h>
    #include <stdio.h>
#endif

namespace
{
    //----------------------------------------------------------------------
    int LeastNumElementsForAlignment(int elementSize)
    {
        static int const sPrecomputed[16] = { 1, 16, 8, 16, 4, 16, 8, 16, 2, 16, 8, 16, 4, 16, 8, 16 };
        return sPrecomputed[elementSize & 0xf];
    }
}

namespace EA {
    namespace Jobs {

using namespace EA::Jobs::UtilDetail;

//----------------------------------------------------------------------
void BatchProcessor::GetNextReadCount()
{
    switch (mBatchDesc.GetDivision())
    {
    case FOR_EACH_DIVISION_EVEN:
        mReadBatchNumber += mBatchDesc.JobCount();
        break;
    case FOR_EACH_DIVISION_DYNAMIC:
        mReadBatchNumber = static_cast<int>(AtomicIncrement(mDynamicCounterEa)) - 1;
        break;
    default:
        break;
    }
}

//----------------------------------------------------------------------
BatchProcessor::BatchProcessor(BatchDesc * batchDesc, char * buffer, size_t bufferSize) :
    mBatchDescEa((uintptr_t)batchDesc),
    mDynamicCounterEa(0),
    mRefCounterEa(0),
    mBuffer(buffer),
    mBufferSize(static_cast<int>(bufferSize)),
    mDone(false),
    mJobIndex(0),
    mBatchDesc(*batchDesc)
{
    Intialize();
}


//----------------------------------------------------------------------
BatchProcessor::BatchProcessor(BatchDesc * batchDesc, size_t jobIndex, char * buffer, size_t bufferSize) :
    mBatchDescEa((uintptr_t)batchDesc),
    mDynamicCounterEa(0),
    mRefCounterEa(0),
    mBuffer(buffer),
    mBufferSize(static_cast<int>(bufferSize)),
    mDone(false),
    mJobIndex(static_cast<int>(jobIndex)),
    mBatchDesc(*batchDesc)
{
    Intialize();
}

//----------------------------------------------------------------------
BatchProcessor::BatchProcessor(uintptr_t batchDescEa, char * buffer, size_t bufferSize) :
    mBatchDescEa(batchDescEa),
    mDynamicCounterEa(0),
    mRefCounterEa(0),
    mBuffer(buffer),
    mBufferSize(static_cast<int>(bufferSize)),
    mDone(false),
    mJobIndex(0)
    #if ! defined EA_PLATFORM_PS3_SPU
        , mBatchDesc(*(BatchDesc *)batchDescEa)
    #endif
{
    EA_ASSERT(batchDescEa);

    #if defined EA_PLATFORM_PS3_SPU
        int const batchDescDmaTag = 31;
        EA_ASSERT(batchDescEa % 16 == 0);
        cellDmaGet((void*)&mBatchDesc, (uint64_t)batchDescEa, sizeof(BatchDesc), batchDescDmaTag, 0, 0);
        cellDmaWaitTagStatusAll(1 << batchDescDmaTag);
    #endif

    mDynamicCounterEa = mBatchDesc.GetDynamicCounterEa(batchDescEa);
    mRefCounterEa = mBatchDesc.GetRefCounterEa(batchDescEa);

    Intialize();
}

//----------------------------------------------------------------------
BatchProcessor::BatchProcessor(uintptr_t batchDescEa, size_t jobIndex, char * buffer, size_t bufferSize) :
    mBatchDescEa(batchDescEa),
    mDynamicCounterEa(0),
    mRefCounterEa(0),
    mBuffer(buffer),
    mBufferSize(static_cast<int>(bufferSize)),
    mDone(false),
    mJobIndex(static_cast<int>(jobIndex))
    #if ! defined EA_PLATFORM_PS3_SPU
        , mBatchDesc(*reinterpret_cast<BatchDesc *>(batchDescEa))
    #endif
{
    EA_ASSERT(batchDescEa);

    #if defined EA_PLATFORM_PS3_SPU
        int const batchDescDmaTag = 31;
        EA_ASSERT(batchDescEa % 16 == 0);
        cellDmaGet((void*)&mBatchDesc, (uint64_t)batchDescEa, sizeof(BatchDesc), batchDescDmaTag, 0, 0);
        cellDmaWaitTagStatusAll(1 << batchDescDmaTag);
    #endif

    mDynamicCounterEa = mBatchDesc.GetDynamicCounterEa(batchDescEa);
    mRefCounterEa = mBatchDesc.GetRefCounterEa(batchDescEa);

    Intialize();
}

void BatchProcessor::Intialize()
{
#ifdef EA_ASSERT_ENABLED
    static uint32_t const sizeOfSpuLocalStore = 256*1024;
    (void)sizeOfSpuLocalStore;
#endif

    mBufferSize &= ~(15); // round down to a multiple of 16

    EA_ASSERT((mBufferSize == 0 && mBuffer == 0) || mBufferSize >= 16);
    EA_ASSERT(mBufferSize % 16 == 0);
    EA_ASSERT((uintptr_t)mBuffer % 16 == 0);

    #if defined EA_PLATFORM_PS3_SPU
        EA_ASSERT((uintptr_t)mBuffer < (sizeOfSpuLocalStore - mBufferSize));
    #else
        EA_ASSERT((mBuffer == 0) || (uintptr_t)mBuffer > sizeOfSpuLocalStore); // assume's VM won't map ppu allocations to < sizeOfSpuLocalStore
    #endif

    EA_ASSERT(mBatchDesc.Verify());

    EA_ASSERT(mBatchDesc.JobCount() > 0);
    EA_ASSERT(mJobIndex < static_cast<int>(mBatchDesc.JobCount()));

    mElementsPerBatch = mBatchDesc.GetElementsPerBatch();

    int elementMultiplePerBatch = mBatchDesc.GetElementMultiplePerBatch();

    int minBatchSize = 1;
    int bytesPerElementSlice = 0;
    int largestElementSize = 0;

    for (int i = 0; i < BatchDesc::sMaxArrays; ++i)
    {
        BatchDesc::ArrayDesc & arrayDesc = mBatchDesc.GetArrayDesc(i);
        if ( arrayDesc.mInPtr || arrayDesc.mOutPtr )
        {
            EA_ASSERT(arrayDesc.mElementSize > 0);

            bytesPerElementSlice += arrayDesc.mElementSize;
            int temp = LeastNumElementsForAlignment(arrayDesc.mElementSize);
            if ( minBatchSize < temp )
                minBatchSize = temp;
            if ( largestElementSize < arrayDesc.mElementSize)
                largestElementSize = arrayDesc.mElementSize;
            if (!mBuffer && arrayDesc.mInPtr && arrayDesc.mOutPtr && (arrayDesc.mInPtr != arrayDesc.mOutPtr))
            {
                EA_FAIL_MSG("Input and Output pointers must be the same when using pass-through mode. (IE: when no buffer is used.)");
            }
        }
    }

    EA_ASSERT(bytesPerElementSlice > 0);

    if (mElementsPerBatch == 0)
    {

        EA_ASSERT(minBatchSize != 0);

        if (minBatchSize % elementMultiplePerBatch)
        {
            if (elementMultiplePerBatch % minBatchSize == 0)
                minBatchSize = elementMultiplePerBatch;
            else
                minBatchSize *= elementMultiplePerBatch; // cheat, should really be lcm(minBatchSize, minElementsPerBatch)
        }

        EA_ASSERT(mBatchDesc.NumElements() % minBatchSize == 0);

        int minBatchBytes = minBatchSize * bytesPerElementSlice;
        int batchCount = mBufferSize / (minBatchBytes * sBufferingCount);

        if (batchCount == 0)
        {
            #if ! defined EA_PLATFORM_PS3_SPU
                if (mBuffer || mBufferSize)
                {
                    EA_FAIL_MSG("Unable to compute elements per batch. Buffer to small");
                }

                switch (mBatchDesc.GetDivision())
                {
                case FOR_EACH_DIVISION_EVEN:
                    batchCount = mBatchDesc.NumElements() / mBatchDesc.JobCount();
                    break;
                case FOR_EACH_DIVISION_DYNAMIC:
                    EA_FAIL_MSG("Unable to compute elements per batch. Use BatchDesc::SetElementsPerBatch or specify a bufferSize. (You can specify a bufferSize and still pass in null for the buffer.  It will be used to compute the number of elements in each Batch.)");
                    break;
                default:
                    break;
                }

                if (batchCount <= 0)
                    batchCount = 1;

            #else
                EA_FAIL_MSG("Unable to compute elements per batch. Buffer to small");
            #endif
        }

        mElementsPerBatch = batchCount * minBatchSize;
    }
    else
    {
        // user overrode the mElementsPerBatch so they must have an multiple of that many elements in total
        EA_ASSERT((mBatchDesc.NumElements() % mElementsPerBatch) == 0);
    }

    /*
    int const dmaLimit = 16 * 1024; // cell hardware limitation
    if (mElementsPerBatch * largestElementSize > dmaLimit)
    {
        EA_ASSERT(mElementsPerBatch * largestElementSize <= dmaLimit);

        mElementsPerBatch = dmaLimit / largestElementSize;
        mElementsPerBatch = minBatchSize * (mElementsPerBatch / minBatchSize);
    }
    */

    EA_ASSERT(mElementsPerBatch);

    mLastBatchNumber = mBatchDesc.NumElements() / mElementsPerBatch;
    mElementsInLastBatch = mBatchDesc.NumElements() % mElementsPerBatch;

    if (mElementsInLastBatch)
        mLastBatchNumber += 1;
    else
        mElementsInLastBatch = mElementsPerBatch;

    EA_ASSERT(0 == (mElementsInLastBatch % minBatchSize));

    mLastBatchNumber -= 1;

    EA_ASSERT(mLastBatchNumber >= 0);
    EA_ASSERT(mElementsInLastBatch != 0);

    if (mBuffer)
    {
        char* b = mBuffer;
        for (int i = 0; i < BatchDesc::sMaxArrays; ++i)
        {
            BatchDesc::ArrayDesc & arrayDesc = mBatchDesc.GetArrayDesc(i);
            if ( arrayDesc.mInPtr || arrayDesc.mOutPtr )
            {
                int bufferSize = mElementsPerBatch * arrayDesc.mElementSize;

                for (int j = 0; j < sBufferingCount; ++j)
                {
                    mArrayOffsetInBuffer[i*sBufferingCount + j] = Align(b);
                    b += bufferSize;
                    EA_ASSERT(b <= (mBuffer + mBufferSize));
                }
            }
            else
            {
                for (int j = 0; j < sBufferingCount; ++j)
                    mArrayOffsetInBuffer[i * sBufferingCount + j] = 0;
            }
        }
    }
}

BatchProcessor::~BatchProcessor()
{
    if (mBatchDesc.mAtomicAllocator && mRefCounterEa && AtomicDecrement(mRefCounterEa)==0)
    {
        #ifdef EA_PLATFORM_PS3_SPU
            //asm("stop 0x3fff\n");
            AtomicAllocator aa;
            aa.Dma(mBatchDesc.mAtomicAllocator);
            //spu_printf("SPU FEDF 0x%08x\n", (uint32_t)mBatchDescEa);
            aa.Free((ForEachDesc*)mBatchDescEa);
        #else
            //printf("PPU FEDF 0x%08x\n", (uint32_t)mBatchDescEa);
            AtomicAllocator* aa = (AtomicAllocator*)mBatchDesc.mAtomicAllocator;
            aa->Free<ForEachDesc>((ForEachDesc*)mBatchDescEa);
        #endif
    }
}

//----------------------------------------------------------------------
void BatchProcessor::First()
{
    mReadBufferIndex = 0;
    mWriteBufferIndex = 0;
    mCurrentBatchNumber = 0;
    mReadBatchNumber = 0;
    mBufferBatchWritesPending = 0;

    switch (mBatchDesc.GetDivision())
    {
    case FOR_EACH_DIVISION_EVEN:
        if (mJobIndex > mLastBatchNumber)
        {
            mDone = true;
            return;
        }
        mReadBatchNumber = mJobIndex - mBatchDesc.JobCount();
        break;
    case FOR_EACH_DIVISION_DYNAMIC:
        break;
    default:
        break;
    }

    for (int i = 0; i < sBufferingCount - 1; ++i)
        Read();

    mCurrentBufferIndex = 0;
    mCurrentBatchNumber = mBufferBatchNumber[mCurrentBufferIndex];
    if (mCurrentBatchNumber > mLastBatchNumber)
    {
        mDone = true;
        return;
    }

    Wait(1<<mCurrentBufferIndex);
}

//----------------------------------------------------------------------
bool BatchProcessor::Done() const
{
    return mDone;
}

//----------------------------------------------------------------------
void BatchProcessor::Next()
{
    Write();
    Read();
    NextBuffer(mCurrentBufferIndex);
    mCurrentBatchNumber = mBufferBatchNumber[mCurrentBufferIndex];

    if (mCurrentBatchNumber > mLastBatchNumber)
    {
        while (mBufferBatchWritesPending > 0)
            Write();

        Wait(sDmaWaitAllMask);
        mDone = true;
    }
    else
    {
        Wait(1<<mCurrentBufferIndex);
    }
}

//----------------------------------------------------------------------
int BatchProcessor::GetNumElements() const {
    EA_ASSERT(mCurrentBatchNumber <= mLastBatchNumber);
    return mBufferBatchNumber[mCurrentBufferIndex] == mLastBatchNumber ? mElementsInLastBatch : mElementsPerBatch;
}

//----------------------------------------------------------------------
int BatchProcessor::GetElementOffset() const {
    EA_ASSERT(mCurrentBatchNumber <= mLastBatchNumber);
    return mCurrentBatchNumber * mElementsPerBatch;
}

//----------------------------------------------------------------------
void BatchProcessor::NextBuffer(int & buffer)
{
    ++buffer;
    if (buffer == sBufferingCount)
        buffer = 0;
}

//----------------------------------------------------------------------
void BatchProcessor::Read()
{
    GetNextReadCount();

    mBufferBatchNumber[mReadBufferIndex] = mReadBatchNumber;

    if (mReadBatchNumber > mLastBatchNumber)
        return;

    int numElements = mReadBatchNumber == mLastBatchNumber ? mElementsInLastBatch : mElementsPerBatch;
    int startingElement = mReadBatchNumber * mElementsPerBatch;

    for (int i = 0; i < BatchDesc::sMaxArrays; ++i)
    {
        BatchDesc::ArrayDesc & arrayDesc = mBatchDesc.GetArrayDesc(i);
        if (arrayDesc.mInPtr)
        {
            int bufferIndex = i * sBufferingCount + mReadBufferIndex;

            EA_ASSERT(0 == ((numElements * arrayDesc.mElementSize) % 16));
            size_t bytes = static_cast<size_t>( numElements * arrayDesc.mElementSize );
            const char * startingEa = arrayDesc.mInPtr + startingElement * arrayDesc.mElementSize;

            #if defined EA_PLATFORM_PS3_SPU
                //EA_ASSERT(bytes <= 16*1024);
                cellDmaLargeGetf(mArrayOffsetInBuffer[bufferIndex], reinterpret_cast<uint64_t>(startingEa), bytes, mReadBufferIndex, 0, 0);
            #else
                if (mBuffer)
                    memcpy(mArrayOffsetInBuffer[bufferIndex], startingEa, bytes);
                else
                    mArrayOffsetInBuffer[bufferIndex] = const_cast<char *>(startingEa);
            #endif
        }
        #if ! defined EA_PLATFORM_PS3_SPU
            if (arrayDesc.mOutPtr && !mBuffer)
            {
                int bufferIndex = i * sBufferingCount + mReadBufferIndex;

                EA_ASSERT(0 == ((numElements * arrayDesc.mElementSize) % 16));
                const char * startingEa = arrayDesc.mOutPtr + startingElement * arrayDesc.mElementSize;
                mArrayOffsetInBuffer[bufferIndex] = const_cast<char *>(startingEa);
            }
        #endif
    }
    NextBuffer(mReadBufferIndex);
    ++mBufferBatchWritesPending;
}

//----------------------------------------------------------------------
void BatchProcessor::Write()
{
    EA_ASSERT(mBufferBatchWritesPending);

    #if ! defined EA_PLATFORM_PS3_SPU
        if (mBuffer)
    #endif
    {
        int numElements = mBufferBatchNumber[mWriteBufferIndex] == mLastBatchNumber ? mElementsInLastBatch : mElementsPerBatch;
        int startingElement = mBufferBatchNumber[mWriteBufferIndex] * mElementsPerBatch;

        for (int i = 0; i < BatchDesc::sMaxArrays; ++i)
        {
            BatchDesc::ArrayDesc & arrayDesc = mBatchDesc.GetArrayDesc(i);
            if (arrayDesc.mOutPtr)
            {
                int bufferIndex = i * sBufferingCount + mWriteBufferIndex;

                EA_ASSERT(0 == ((numElements * arrayDesc.mElementSize) % 16));
                size_t bytes = static_cast<size_t>( numElements * arrayDesc.mElementSize );
                char * startingEa = arrayDesc.mOutPtr + startingElement * arrayDesc.mElementSize;

                #if defined EA_PLATFORM_PS3_SPU
                    //EA_ASSERT(bytes <= 16*1024);
                    cellDmaLargePutf(mArrayOffsetInBuffer[bufferIndex], reinterpret_cast<uint64_t>(startingEa), bytes, mWriteBufferIndex, 0, 0);
                #else
                    memcpy(startingEa, mArrayOffsetInBuffer[bufferIndex], bytes);
                #endif
            }
        }
        NextBuffer(mWriteBufferIndex);
    }

    --mBufferBatchWritesPending;
}

//----------------------------------------------------------------------
void BatchProcessor::Wait(int dmaTagMask)
{
    #if defined EA_PLATFORM_PS3_SPU
        cellDmaWaitTagStatusAll(dmaTagMask);
    #else
        (void)dmaTagMask;
    #endif
}

//----------------------------------------------------------------------
char * BatchProcessor::GetArrayTypeless(int index)
{
    EA_ASSERT(index >= 0 && index < BatchDesc::sMaxArrays);
    return static_cast<char *>(mArrayOffsetInBuffer[index * sBufferingCount + mCurrentBufferIndex]);
}

}}
