#include "Cooker/WorkerRequestsLocal.h"
#include "Cooker/CookDirector.h"
#include "Cooker/PackageTracker.h"
#include "Cooker/CookPlatformManager.h"
namespace UE
{
namespace Cook
{
FWorkerRequestsLocal::FWorkerRequestsLocal(UCookOnTheFlyServer& InCOTFS)
    : IWorkerRequests(InCOTFS)
{
}

ECookMode::Type FWorkerRequestsLocal::GetDirectorCookMode()
{
    return COTFS.GetCookMode();
}

FWorkerRequestsLocal::~FWorkerRequestsLocal()
{
}

void FWorkerRequestsLocal::OnRemoveSessionPlatform(const ITargetPlatform* TargetPlatform)
{
    ExternalRequests.OnRemoveSessionPlatform(TargetPlatform);
}

bool FWorkerRequestsLocal::DequeueCallbacks(TArray<FSchedulerCallback>& OutCallbacks)
{
    return ExternalRequests.DequeueCallbacks(OutCallbacks);
}

bool FWorkerRequestsLocal::HasRequests()
{
    return ExternalRequests.HasRequests();
}

void FWorkerRequestsLocal::AddCallback(FSchedulerCallback&& Callback)
{
    ExternalRequests.AddCallback(std::move(Callback));
}

void FWorkerRequestsLocal::EnqueueUnique(FFilePlatformRequest&& FileRequest, bool bForceFrontOfQueue)
{
    ExternalRequests.EnqueueUnique(std::move(FileRequest), bForceFrontOfQueue);
}

void FWorkerRequestsLocal::DequeueAll(TArray<FSchedulerCallback>& OutCallbacks, TArray<FFilePlatformRequest>& OutCookRequests)
{
    ExternalRequests.DequeueAll(OutCallbacks, OutCookRequests);
}

EExternalRequestType FWorkerRequestsLocal::DequeueRequest(TArray<FSchedulerCallback>& OutCallbacks, FFilePlatformRequest& OutToBuild)
{
    return ExternalRequests.DequeueRequest(OutCallbacks, OutToBuild);
}

int32 FWorkerRequestsLocal::GetNumRequests() const
{
    return ExternalRequests.GetNumRequests();
}

void FWorkerRequestsLocal::RemapTargetPlatforms(const TMap<ITargetPlatform*, ITargetPlatform*>& Remap)
{
    ExternalRequests.RemapTargetPlatforms(Remap);
}

bool FWorkerRequestsLocal::AddToWaitAssign(FPackageData& PackageData, EPackageState TargetState)
{
    EPackageState CurrentState = PackageData.GetState();
    if (CurrentState != EPackageState::Request)
    {
        if (CurrentState == EPackageState::Idle)
        {
            if (TargetState != EPackageState::LoadReady || PackageData.HasAllCookedPlatforms(COTFS.PlatformManager->GetSessionPlatforms(), true))
            {
                UE_LOG(LogCook, Error, TEXT("AddToWaitAssign With Cooked-------- Package %s %d=>%d"), *PackageData.GetPackageName().ToString(), (int32)CurrentState, (int32)TargetState);
                return true;
            }
        }
        else
        {
            if (CurrentState != EPackageState::WaitAssign)
            {
                UE_LOG(LogCook, Error, TEXT("AddToWaitAssign With Cooking++++++++ Package %s %d=>%d"), *PackageData.GetPackageName().ToString(), (int32)CurrentState, (int32)TargetState);
            }
            return true;
        }
    }
    if (TargetState == EPackageState::LoadReady)
    {
        if (COTFS.CookDirector)
        {
            COTFS.CookDirector->AssignRequestTo(&PackageData, COTFS.CookDirector->GetLocalWorkerId());
        }
        PackageData.SendToState(EPackageState::LoadReady, ESendFlags::QueueNone);
        COTFS.PackageDatas->GetLoadReadyQueue().AddFront(&PackageData);
        return false;
    }
    if (COTFS.CookDirector && COTFS.CookDirector->GetWorkerCount() > 0)
    {
        PackageData.SendToState(EPackageState::WaitAssign, ESendFlags::QueueAdd);
        WaitAssignQueue.Add(&PackageData);
        return true;
    }
    PackageData.SendToState(EPackageState::LoadPrepare, ESendFlags::QueueAdd);
    return false;
}

bool FWorkerRequestsLocal::AssignRequests()
{
    if (WaitAssignQueue.Num() == 0)
    {
        return false;
    }

    TArray<FPackageData*> Requests;
    while (WaitAssignQueue.Num() > 0)
    {
        Requests.Add(WaitAssignQueue.PopFrontValue());
    }
    TArray<FWorkerId> OutAssignments;
    if (!COTFS.CookDirector->AssignRequests(Requests, OutAssignments))
    {
        COTFS.CookDirector->AssignRequests(Requests, COTFS.CookDirector->GetLocalWorkerId());
        return true;
    }

    TArray<TArray<FPackageData*>> RemoteBatches;  // Indexed by WorkerId.GetRemoteIndex()
    TArray<FWorkerId> RemoteWorkers;
    int WorkerCount = COTFS.CookDirector->GetWorkerCount();
    RemoteBatches.SetNum(WorkerCount);
    RemoteWorkers.SetNum(WorkerCount);
    for (int RemoteIndex = 0; RemoteIndex < WorkerCount; RemoteIndex++)
    {
        TArray<FPackageData*>& RemoteBatch = RemoteBatches[RemoteIndex];
        RemoteBatch.Reserve(2 * Requests.Num() / WorkerCount);
    }

    for (int32 RequestIndex = 0; RequestIndex < Requests.Num(); ++RequestIndex)
    {
        FWorkerId& WorkerId                = OutAssignments[RequestIndex];
        int32 RemoteIndex                  = WorkerId.GetMultiprocessId();
        RemoteWorkers[RemoteIndex]         = WorkerId;
        TArray<FPackageData*>& RemoteBatch = RemoteBatches[RemoteIndex];
        RemoteBatch.Add(Requests[RequestIndex]);
    }

    for (int RemoteIndex = 0; RemoteIndex < WorkerCount; RemoteIndex++)
    {
        COTFS.CookDirector->AssignRequests(RemoteBatches[RemoteIndex], RemoteWorkers[RemoteIndex]);
    }
    return true;
}

bool FWorkerRequestsLocal::DecideNextCookAction(FTickStackData& StackData, UCookOnTheFlyServer::ECookAction& CookAction)
{
    if (WaitAssignQueue.Num() > static_cast<int32>(DesiredWaitAssignQueueLength))
    {
        CookAction = UCookOnTheFlyServer::ECookAction::Assign;
        return true;
    }
    return false;
}
}  // namespace Cook
}  // namespace UE