#pragma once

#include "Cooker/MPCookMessage.h"
#include "Cooker/CookTypes.h"
#include "Cooker/CookDirector.h"
class FSocket;
namespace UE
{
namespace Cook
{
enum class ECookDirectorThread : uint8;
class FCookDirector;
enum class ENotifyRemote
{
    NotifyRemote,
    LocalOnly,
};
class FCookWorkerServer: public FThreadSafeRefCountedObject
    , public FCookSocket
{
public:
    FCookWorkerServer(FCookDirector& InDirector, FWorkerId InWorkerId);
    ~FCookWorkerServer();

    /** Is this either shutting down or completed shutdown of its remote Client? */
    bool IsShuttingDown() const;

    /** Is this CookWorker fully disconnected? */
    bool IsConnectionLost() const;

    /** Periodic Tick function to send and receive messages to the Client. */
    void TickCommunication(ECookDirectorThread TickThread);

    /** Helper for PumpConnect, launch the remote Client process. */
    void LaunchProcess();
    /** Helper for Tick, pump send messages to a connected Client. */
    void PumpSendMessages();
    /** Helper for Tick, pump receive messages from a connected Client. */
    void PumpReceiveMessages();

    bool TryHandleConnectMessage(FWorkerConnectMessage& Message, FSocket* InSocket);

    /** Add the given assignments for the CookWorker. They will be sent during Tick */
    void AppendAssignments(TArrayView<FPackageData*> Assignments, ECookDirectorThread TickThread);

    void TickFromSchedulerThread();

    /** How many package assignments is the remote CookWorker supposed to save but hasn't yet? */
    int32 NumAssignments() const;

    /** Called when the Server detects all packages are complete. Tell the CookWorker to flush messages and exit. */
    void SignalCookComplete(ECookDirectorThread TickThread);

    /** Kill the Client process (non-graceful termination), and close the connection resources. */
    void ShutdownRemoteProcess();

private:
    enum class EConnectStatus
    {
        Uninitialized,
        WaitForConnect,
        Connected,
        ShuttingDown,
        LostConnection,
    };
    enum class EWorkerDetachType
    {
        Dismissed,
        StillRunning,
        ForceTerminated,
        Crashed,
    };
    /** Send this into the given state. Update any state-dependent variables. */
    void SendToState(EConnectStatus TargetStatus);
    /** Close the connection and connection resources to the remote process. Does not kill the process. */
    void DetachFromRemoteProcess(EWorkerDetachType DetachType);

    /** Helper for PumpReceiveMessages: dispatch the messages received from the socket. */
    void HandleReceiveMessages();
    void HandleSendMessages();

    /** Add results from the client to the local CookOnTheFlyServer. */
    void RecordResults(FPackageResultsMessage& Message);
    void QueueDiscoveredPackage(FDiscoveredPackagesMessage& Message);
    void CookSideEffectPackage(FCookSideEffectMessage& Message);

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

    void ReassignPackages();

private:
    FCookDirector& Director;
    UCookOnTheFlyServer& COTFS;
    FWorkerId WorkerId = FWorkerId::Invalid();
    FProcHandle CookWorkerHandle;
    FString CrashDiagnosticsError;
    EConnectStatus ConnectStatus = EConnectStatus::Uninitialized;
    bool bNeedCrashDiagnostics   = false;
    uint32 CookWorkerProcessId   = 0;
    uint32 SyncFenceVersion = 0, SendSyncFenceVersion = 0;
    double ConnectStartTimeSeconds = 0.;
    mutable FCriticalSection CommunicationLock;

    FCriticalSection MessagesLock;
    TArray<ITargetPlatform*> OrderedSessionPlatforms;
    TArray<FPackageData*> PackagesToAssign;
    TSet<FPackageData*> AssignedPackages;

    FAssignPackagesMessage PendingAssignMessage;
    FAssignPackagesMessage PendingSendAssignMessage;
    TArray<FMarshalledMessage> PendingMessages;
};
}  // namespace Cook
}  // namespace UE