#pragma once
#include "CookOnTheSide/CookOnTheFlyServer.h"
#include "CookTypes.h"
#include "Cooker/CookDirector.h"
class FSocket;
class FAssetRegistryGenerator;
namespace UE
{
namespace Cook
{
struct FDirectorConnectionInfo;
struct FAssignPackagesMessage;
struct FCookSideEffectScope;
class FCookWorkerClient: public FCookSocket
{
private:
    enum class EConnectStatus
    {
        Uninitialized,
        WaitInitialConfig,
        Connected,
        ShuttingDown,
        LostConnection,
    };

public:
    FCookWorkerClient(UCookOnTheFlyServer& COTFS);
    ~FCookWorkerClient();
    ECookMode::Type GetDirectorCookMode()
    {
        return DirectorCookMode;
    }

    /** Blocking operation: open the socket to the Director, send the Connected message, receive the setup message. */
    bool TryConnect(FDirectorConnectionInfo&& ConnectInfo);

    /** Reentrant helper for TryConnect which early exits if currently blocked. */
    EPollStatus PollTryConnect(const FDirectorConnectionInfo& ConnectInfo);

    /** Helper for PollTryConnect: create the ServerSocket */
    void CreateServerSocket(const FDirectorConnectionInfo& ConnectInfo);

    /** Send this into the given state. Update any state-dependent variables. */
    void SendToState(EConnectStatus TargetStatus);

    /** Wait for the Config Message, switch state when it succeeds or fails. */
    void WaitInitialConfigMessage();

    /** Mark that initialization is complete and we can free the memory for initialization settings. */
    void DoneWithInitialSettings();

    /** Periodic Tick function to send and receive messages to the Server. */
    void TickFromSchedulerThread(FTickStackData& StackData);

    /** Helper for Tick, pump send messages to the Server. */
    void PumpSendMessages(bool bFlush = false);

    /** Helper for Tick, pump receive messages from the Server. */
    void PumpReceiveMessages();

    /** Helper for Tick, pump Send/Receive and check for whether we are done shutting down. */
    void PumpDisconnect(FTickStackData& StackData);

    /** Helper for PumpReceiveMessages: dispatch the messages received from the socket. */
    void HandleReceiveMessages(FName OptionalPackageName = NAME_None);

    /** Send packages assigned from the server into the request state. */
    void AssignPackages(FAssignPackagesMessage& Message);

    /** Queue a message to the server that a package was discovered as needed in the cook. Will be sent during Tick. */
    void ReportDiscoveredPackage(TArray<FPackageData*>& Requests, TSet<FPackageData*>& LoadedPackages);

    void ReportPackageMessage(const FPackageData& PackageData, EMPPackageResult Reason);

    void LogInvalidMessage(const TCHAR* MessageTypeName, FGuid MessageType);

    void TrySyncFenceMessage();

    void UpdateAssetRegistryPackageData(FAssetRegistryGenerator* Generator, const UPackage& Package, UCookOnTheFlyServer::FCreateOrFindArray& OutputAssets);

    FCookByTheBookOptions&& ConsumeCookByTheBookOptions();
    FCookByTheBookStartupOptions&& ConsumeCookByTheBookStartupOptions();
    /** Reads the OrderedSessionPlatforms received from the Director. */
    const TArray<ITargetPlatform*>& GetTargetPlatforms() const;

    void MarkAssetRegistryDirty() { bAssetRegistryDirty = true; }

private:
    FMPCookSideEffectCollectorData& GetCookCollectorData() { return CookCollectorData; }
    friend class FMPCookSideEffectCollector;

    FString DirectorURI;
    EConnectStatus ConnectStatus         = EConnectStatus::Uninitialized;
    bool bFlushPendingResults            = false;
    bool bSendShutdownAcknowledgeMessage = false;
    bool bAssetRegistryDirty             = false;
    bool bFlushAssetRegistryMessage      = false;
    uint32 SyncFenceVersion = 0, SendSyncFenceVersion = 0;
    double ConnectStartTimeSeconds   = 0.;
    ECookMode::Type DirectorCookMode = ECookMode::CookByTheBook;
    FMPCookSideEffectCollectorData CookCollectorData;
    TUniquePtr<FInitialConfigMessage> InitialConfigMessage;
    TArray<ITargetPlatform*> OrderedSessionPlatforms;
    FDiscoveredPackagesMessage SendDiscoveredPackagesMessage;
    FCookSideEffectMessage SendAssetRegistryMessage;
    FPackageResultsMessage SendPackageResultsMessage;
    FDiscoveredPackagesMessage PendingDiscoveredPackagesMessage;
    FPackageResultsMessage PendingPackageResultsMessage;
    TArray<FMarshalledMessage> PendingMessages;
    UCookOnTheFlyServer& COTFS;
};
}  // namespace Cook
}  // namespace UE