#pragma once
#include "CoreTypes.h"
#include "Cooker/MPCookMessage.h"

namespace UE
{
namespace Cook
{
/** Parameters parsed from commandline for how a CookWorker connects to the CooKDirector. */
struct FDirectorConnectionInfo
{
    bool TryParseCommandLine();

    FString HostURI;
    int32 RemoteIndex  = 0;
    int32 DirectorPort = 0;
};
class FCookWorkerServer;
/**
 * The categories of thread that can pump the communication with CookWorkers. It can be pumped either from
 * the cooker's scheduler thread (aka Unreal's game thread or main thread) or from a worker thread.
 */
enum class ECookDirectorThread : uint8
{
    SchedulerThread,
    CommunicateThread,
    Invalid,
};

struct FWorkerAssignment
{
    FWorkerId Id;
    int32 CurrentLoad = 0;
    TSet<FName> PackagesInProgress;
};

class FCookDirector
{
    friend class FCookWorkerServer;

public:
    FCookDirector(UCookOnTheFlyServer& InCOTFS, int32 CookProcessCount, bool bInCookProcessCountSetByCommandLine);
    ~FCookDirector();

    void InitializeWorkers();

    bool IsMultiprocessAvailable() const;

    /** Start the communication thread if not already started. */
    void LaunchCommunicationThread();
    /** Signal the communication thread to stop, wait for it to finish, and deallocate it. */
    void StopCommunicationThread();
    /** Entry point for the communication thread. */
    uint32 RunCommunicationThread();

    /**
     * Execute a single frame of communication with CookWorkers: send/receive to all CookWorkers,
     * including connecting, ongoing communication, and shutting down.
     */
    void TickCommunication(ECookDirectorThread TickThread);

    /** Tick helper: tick any workers that have not yet finished initialization. */
    void TickWorkerConnects(ECookDirectorThread TickThread);

    void StartCook();

    /** Initialization helper: create the listen socket. */
    bool TryCreateWorkerConnectSocket();

    /** Enum specifying how CookWorker log output should be shown. */
    enum EShowWorker
    {
        CombinedLogs,
        SeparateLogs,
        SeparateWindows,  // Implies SeparateLogs as well
    };
    EShowWorker GetShowWorkerOption() const { return ShowWorkerOption; }
    /** Data used by a CookWorkerServer to launch the remote process. */
    struct FLaunchInfo
    {
        EShowWorker ShowWorkerOption;
        FString CommandletExecutable;
        FString WorkerCommandLine;
    };
    FLaunchInfo GetLaunchInfo(FWorkerId WorkerId);
    FString GetWorkerCommandLine(FWorkerId WorkerId);
    /** The LogPath a worker process writes to. */
    FString GetWorkerLogFileName(FWorkerId WorkerId);
    FInitialConfigMessage* GetInitialConfigMessage();
    FWorkerId GetLocalWorkerId();
    bool AssignRequestTo(FPackageData* PackageData, FWorkerId WorkerId);
    /** Periodic tick function. Sends/Receives messages to CookWorkers. */
    void TickFromSchedulerThread();
    bool AssignRequests(TArrayView<FPackageData*> Batch, FWorkerId WorkerId);
    bool AssignRequests(TArrayView<FPackageData*> Batch, TArray<FWorkerId>& OutAssignments);
    int32 GetWorkerCount()
    {
        return RequestedCookWorkerCount + 1;
    }
    /**
     * Called when the COTFS Server has detected all packages are complete. Tells the CookWorkers to flush messages
     * and exit.
     */
    void PumpCookComplete(bool& bOutCompleted);

protected:
    /** Struct that implements the FRunnable interface and forwards it to to named functions on this FCookDirector. */
    struct FRunnableShunt: public FRunnable
    {
        FRunnableShunt(FCookDirector& InDirector): Director(InDirector) {}
        virtual uint32 Run() override;
        virtual void Stop() override;
        FCookDirector& Director;
    };
    // Synchronization primitives that can be used from any thread
    mutable FCriticalSection CommunicationLock;

    FEventRef ShutdownEvent{EEventMode::ManualReset};
    bool bWorkersActive      = false;
    bool bWorkersInitialized = false;
    bool bCookCompleteSent   = false;
    // Data only accessible from the SchedulerThread
    FRunnableShunt RunnableShunt;
    FRunnableThread* CommunicationThread = nullptr;
    FString WorkerConnectAuthority;
    int32 RequestedCookWorkerCount = 0;
    int32 WorkerConnectPort        = 0;
    FWorkerId LocalWorkerId;
    TArray<FWorkerAssignment> WorkerAssignments;
    EShowWorker ShowWorkerOption = EShowWorker::CombinedLogs;
    FString CommandletExecutablePath;

    UCookOnTheFlyServer& COTFS;
    // Data only accessible from the CommunicationThread (or if the CommunicationThread is inactive)
    FSocket* WorkerConnectSocket = nullptr;
    // Data shared between SchedulerThread and CommunicationThread that can only be accessed inside CommunicationLock
    TArray<TRefCountPtr<FCookWorkerServer>> RemoteWorkers;
    TArray<TRefCountPtr<FCookWorkerServer>> ShuttingDownWorkers;
    TUniquePtr<FInitialConfigMessage> InitialConfigMessage;
};
}  // namespace Cook
}  // namespace UE