// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#if WITH_EDITOR
#include "Async/Future.h"
#include "Containers/Array.h"
#include "Containers/ArrayView.h"
#include "Containers/UnrealString.h"
#include "HAL/Platform.h"
#include "Misc/Guid.h"
#include "Templates/Function.h"
#include "Templates/RefCounting.h"
#include "Cooker/CookTypes.h"
#include "Cooker/MPCookSideEffect.h"
#endif

#if WITH_EDITOR

class FSocket;
class UCookOnTheFlyServer;
namespace UE
{
namespace Cook
{
/**
 *  Identifier for a CookWorker process launched from a Director process, or for the local process.
 *  A director can have multiple CookWorkers.
 */
struct FWorkerId
{
public:
    FWorkerId() { Id = InvalidId; }
    constexpr static FWorkerId Invalid() { return FWorkerId(InvalidId); }
    constexpr static FWorkerId Local() { return FWorkerId(LocalId); }
    static FWorkerId FromRemoteIndex(uint8 Index)
    {
        check(Index < InvalidId - 1U);
        return FWorkerId(Index + 1U);
    }
    static FWorkerId FromLocalOrRemoteIndex(uint8 Index)
    {
        check(Index < InvalidId);
        return FWorkerId(Index);
    }
    static int32 GetMaxCookWorkerCount() { return static_cast<int32>(InvalidId - 1); }

    bool IsValid() const { return Id != InvalidId; }
    bool IsInvalid() const { return Id == InvalidId; }
    bool IsLocal() const { return Id == LocalId; }
    bool IsRemote() const { return Id != InvalidId && Id != LocalId; }
    uint8 GetRemoteIndex() const
    {
        check(IsRemote());
        return Id - 1U;
    }
    uint8 GetMultiprocessId() const
    {
        check(IsValid());
        return Id;
    }
    bool operator==(const FWorkerId& Other) const { return Id == Other.Id; }
    bool operator!=(const FWorkerId& Other) const { return Id != Other.Id; }
    bool operator<(const FWorkerId& Other) const { return Id < Other.Id; }
    inline friend int32 GetTypeHash(const FWorkerId& WorkerId) { return WorkerId.Id; }

private:
    constexpr explicit FWorkerId(uint8 InId): Id(InId) {}

private:
    uint8 Id;

    constexpr static uint8 InvalidId = 255;
    constexpr static uint8 LocalId   = 0;
};

// CookUniqueBuffer.h
namespace Sockets
{
constexpr int32 COOKDIRECTOR_DEFAULT_REQUEST_CONNECTION_PORT = 41897;
}

class FCookUniqueBuffer
{
public:
    FCookUniqueBuffer()                                    = default;
    FCookUniqueBuffer(const FCookUniqueBuffer&)            = delete;
    FCookUniqueBuffer& operator=(const FCookUniqueBuffer&) = delete;

    FCookUniqueBuffer(FCookUniqueBuffer&&)                 = default;
    FCookUniqueBuffer& operator=(FCookUniqueBuffer&&)      = default;

    /** Make an uninitialized owned buffer of the specified size. */
    static FCookUniqueBuffer Alloc(uint32 Size)
    {
        FCookUniqueBuffer NewBuffer;
        NewBuffer.Data.SetNumUninitialized(Size);
        return NewBuffer;
    }

    /** Returns a pointer to the start of the buffer. */
    void* GetData()
    {
        return Data.GetData();
    }
    const void* GetData() const
    {
        return Data.GetData();
    }

    /** Returns the size of the buffer in bytes. */
    uint32 GetSize() const
    {
        return Data.Num();
    }

    /** Returns true if this buffer is null (empty). */
    bool IsNull() const
    {
        return Data.Num() == 0;
    }

    /**
     * Convert this to a shared buffer (TSharedPtr<TArray<uint8>>), leaving this null.
     * This is the simulated Zero-Copy transfer point.
     * * @return TSharedPtr<TArray<uint8>> An immutable-like shared reference.
     */
    bool MoveToShared(TArray<uint8>& SharedPtr)
    {
        if (IsNull())
        {
            return false;
        }
        SharedPtr = MoveTemp(Data);
        return true;
    }

private:
    TArray<uint8> Data;
};

struct FPacketHeader
{
    uint32 Magic;
    uint32 Size;

    FPacketHeader(uint32 InSize);

    FPacketHeader(): Magic(0), Size(0) {};
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex = 0);
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex = 0);
    bool Check(uint32 MaxPacketSize);
};

constexpr uint32 MessageHeaderExpectedMagic = 0x854EBA92;
/**
 * Holds the state of the CompactBinaryTCP communication read from a Socket. Expects the socket to be a
 * series of packets of messages and keeps track of the next packet when a packet is as yet only partially received.
 */
struct FReceiveBuffer
{
private:
    FCookUniqueBuffer Payload;
    uint32 BytesRead   = 0;
    bool bParsedHeader = false;
    friend class FCookSocket;
};

struct ICookMessage
{
public:
    virtual ~ICookMessage() {};
    virtual FGuid GetMessageType() const                                            = 0;
    virtual void Serialize(TArray<uint8>& SerializedData, int32 StartIndex)         = 0;
    virtual bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) = 0;
    void Serialize(TArray<uint8>& SerializedData)
    {
        Serialize(SerializedData, 0);
    }
    bool Unserialize(const TArray<uint8>& SerializedData)
    {
        return Unserialize(SerializedData, 0);
    }
};

struct FMarshalledMessage;
struct FShutdownRequestMessage;
struct FShutdownAcknowledgeMessage;
class FCookSocket
{
public:
    virtual ~FCookSocket();
    FSocket* DetachSocket();
    void InitSocket(FSocket* InSocket);
    void CloseSocket();
    void BlockSocket(bool IsBlocking);

    virtual bool WriteBuffer(TArray<uint8>& Buffer, int32 StartIndex = 0);
    virtual bool WritePacket(ICookMessage& Message);
    virtual bool ReadPacket(ICookMessage& Message);
    virtual EConnectionStatus TryReadPacket(TArray<FMarshalledMessage>& Messages, uint32 MaxPacketSize = 0);
    virtual bool WriteMarshalledPacket(ICookMessage& Message);
    virtual bool ReadMarshalledPacket(FMarshalledMessage& Messages, ICookMessage& Message);

    bool WriteMarshalledPacket(FShutdownAcknowledgeMessage&& Message);
    bool WriteMarshalledPacket(FShutdownRequestMessage&& Message);
    static int32 MaxSendPackageNum;

private:
    int32 HeaderSize = 0;
    FSocket* Socket  = nullptr;
    TArray<uint8> SendBuffer;
    TArray<uint8> PendingBuffer;
    FReceiveBuffer ReceiveBuffer;
};

/** The atom of data that can be sent through CompactBinaryTCP: guid identifier and compact binary payload. */
struct FMarshalledMessage: public ICookMessage
{
    FGuid MessageType;
    TArray<uint8> MessageBuffer;
    int32 BufferIndex = 0;

    static int32 SerializeMessageType(TArray<uint8>& SerializedData, FGuid MessageType, int32 StartIndex);
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;

    int32 GetBufferIndex()
    {
        return BufferIndex;
    }
    virtual FGuid GetMessageType() const override;
    void Clear();
    bool MakeFrom(FCookUniqueBuffer& Payload);
};

/** Message sent from a CookWorker to the Director to report that it is ready for setup messages and cooking. */
struct FWorkerConnectMessage: public ICookMessage
{
public:
    virtual FGuid GetMessageType() const override { return MessageType; };
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;

public:
    int32 RemoteIndex = 0;
    static FGuid MessageType;
};

struct FInitialConfigMessage: public ICookMessage
{
public:
    virtual FGuid GetMessageType() const override { return MessageType; };
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;

    void ReadFromLocal(const UCookOnTheFlyServer& COTFS, const TConstArrayView<const ITargetPlatform*>& InOrderedSessionPlatforms);

    FCookByTheBookOptions&& ConsumeCookByTheBookOptions() { return MoveTemp(CookByTheBookOptions); }
    FCookByTheBookStartupOptions&& ConsumeCookByTheBookStartupOptions() { return MoveTemp(CookByTheBookStartupOptions); }
    void WriteToTargetPlatform(TArray<ITargetPlatform*>& InOrderedSessionPlatforms);

    ECookMode::Type GetDirectorCookMode() const { return DirectorCookMode; }

public:
    FString OutputDirectoryOverride;
    int32 MaxConcurrentShaderJobs = 0;
    bool bCookIncremental         = true;

    FCookByTheBookOptions CookByTheBookOptions;
    FCookByTheBookStartupOptions CookByTheBookStartupOptions;
    TArray<FString> OrderedSessionPlatforms;
    ECookMode::Type DirectorCookMode = ECookMode::CookByTheBook;
    static FGuid MessageType;
};

/** Information about a PackageData assigned to a CookWorker from the director. */
struct FAssignPackageData
{
    FName PackageName;
    FName NormalizedFileName;
};

/** Message from Server to Client to cook the given packages. */
struct FAssignPackagesMessage: public ICookMessage
{
public:
    FAssignPackagesMessage() = default;
    FAssignPackagesMessage(TArray<FAssignPackageData>&& InPackageDatas);
    virtual FGuid GetMessageType() const override { return MessageType; };
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;

    void MessageSliceTo(int32 BatchCount, FAssignPackagesMessage& PendingMessage);

public:
    TArray<FAssignPackageData> Packages;
    uint32 FenceVersion = 0;
    static FGuid MessageType;
};

/** Information about a discovered package sent from a CookWorker to the Director. */
struct FDiscoveredPackageData
{
    FName PackageName;
    FName NormalizedFileName;
    EPackageState TargetState;
};

/**
 * Message from CookWorker to Director that reports dependency packages discovered during load/save of
 * a package that were not found in the earlier traversal of the packages dependencies.
 */
struct FDiscoveredPackagesMessage: public ICookMessage
{
public:
    virtual FGuid GetMessageType() const override { return MessageType; }
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;
    void MessageSliceTo(int32 BatchCount, FDiscoveredPackagesMessage& PendingMessage);

public:
    TArray<FDiscoveredPackageData> Packages;
    static FGuid MessageType;
};

enum class EMPPackageResult : uint8
{
    Success = 0,
    Failed  = 2,
};
struct FPackageCookResultData
{
    FName PackageName;
    EMPPackageResult PackageResult;
    TArray<FAssetPackageData> AssetPackages;
    TArray<ECookResult> CookResult;
};

/** Message from Client to Server giving the results for saved or refused-to-cook packages. */
struct FPackageResultsMessage: public ICookMessage
{
public:
    virtual FGuid GetMessageType() const override { return MessageType; }
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;
    void MessageSliceTo(int32 BatchCount, FPackageResultsMessage& PendingMessage);

public:
    TArray<FPackageCookResultData> Results;
    static FGuid MessageType;
};

struct FCookSideEffectMessage: public ICookMessage
{
public:
    virtual FGuid GetMessageType() const override { return MessageType; }
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;

    TMap<FName, TArray<FCookAssetData>> AssetRegistryDataMap;
    TMap<FName, TArray<FAssetPackageData>> AssetPackageDataMap;
    FCookSideEffects CookSideEffects;
    static FGuid MessageType;
};

struct FSyncFenceMessage: public ICookMessage
{
public:
    virtual FGuid GetMessageType() const override { return MessageType; }
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;

    uint32 FenceVersion = 0;
    static FGuid MessageType;
};

/** Message from Server to Client to request a graceful shutdown. */
struct FShutdownRequestMessage: public ICookMessage
{
public:
    virtual FGuid GetMessageType() const override { return MessageType; }
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;

public:
    static FGuid MessageType;
};

/** Message from Client to Server to acknowledge the shutdown request. */
struct FShutdownAcknowledgeMessage: public ICookMessage
{
public:
    virtual FGuid GetMessageType() const override { return MessageType; }
    void Serialize(TArray<uint8>& SerializedData, int32 StartIndex) override;
    bool Unserialize(const TArray<uint8>& SerializedData, int32 StartIndex) override;

public:
    static FGuid MessageType;
};
}  // namespace Cook
}  // namespace UE

#endif  // WITH_EDITOR