#include "Base.h"

static constexpr wchar_t kPipeName[] = L"\\\\.\\pipe\\HB_PIPE_DEMO";
static constexpr DWORD kHeartbeatIntervalMs = 250;
static constexpr DWORD kHeartbeatTimeoutMs = 2000;
static constexpr uint32_t kProtocolVersion = 4;
static constexpr size_t kPayloadBytes = 1024;

enum class MsgType : uint32_t
{
    Hello = 1,
    HelloAck = 2,
    Heartbeat = 3,
    Exit = 4,
    Goodbye = 5,
    CommandRequest = 6,
    CommandReply = 7,
};

enum class CommandId : uint32_t
{
    None = 0,
    Printf = 1,
    LoadDriver = 2,
    UnloadDriver = 3,
    QueryStatus = 4,
    BreakpointRemoved = 5,
    CheckVMware = 6,
    CheckVmx = 7,
};

#pragma pack(push, 8)
struct Msg
{
    uint32_t type;
    uint32_t client_id;
    uint32_t pid;
    uint32_t version;
    uint32_t command;
    uint32_t status;
    uint32_t input_size;
    uint32_t output_size;
    uint32_t reserved;
    uint64_t tick;
    uint64_t session_id;
    uint64_t request_id;
    uint8_t payload[kPayloadBytes];
};

struct PrintfRequest
{
    char text[512];
};

struct LoadDriverRequest
{
    wchar_t service_name[128];

    wchar_t driver_path[MAX_PATH];

    DWORD status;
};

struct UnloadDriverRequest
{
    wchar_t service_name[128];
};

struct BreakpointRemovedRequest
{
    uint32_t thread_id;
    uint32_t reserved;
    uint64_t dr0;
    uint64_t dr7;
};
struct QueryStatusReply
{
    DWORD server_pid;
    DWORD already_open;
    char text[128];
};
#pragma pack(pop)

static_assert(sizeof(PrintfRequest) <= kPayloadBytes, "PrintfRequest too large");
static_assert(sizeof(LoadDriverRequest) <= kPayloadBytes, "LoadDriverRequest too large");
static_assert(sizeof(UnloadDriverRequest) <= kPayloadBytes, "UnloadDriverRequest too large");
static_assert(sizeof(BreakpointRemovedRequest) <= kPayloadBytes, "BreakpointRemovedRequest too large");
static_assert(sizeof(QueryStatusReply) <= kPayloadBytes, "QueryStatusReply too large");

inline LONG64 NowMs()
{
    return static_cast<LONG64>(GetTickCount64());
}

struct PendingCommand
{
    uint32_t command = static_cast<uint32_t>(CommandId::None);
    DWORD status = ERROR_GEN_FAILURE;
    uint64_t request_id = 0;
    std::vector<uint8_t> input;
    std::vector<uint8_t> output;
    HANDLE done_event = nullptr;
};

BOOL QueryServerAlreadyOpen(char* out_text, size_t out_text_size);

BOOL DllClientStart();

BOOL HBClientCall(
    uint32_t command,
    const void* input,
    uint32_t input_size,
    void* output,
    uint32_t output_size,
    uint32_t* bytes_returned);

bool WideToUtf8(const std::wstring& input, char* output, size_t output_size);

BOOL DllClientPrintfW(const wchar_t* text);

BOOL HBClientLoadDriverW(PDWORD status);

BOOL QueryServerAlreadyOpen();