#include "../Shared/TriggerProtocol.h"
#include "../Shared/DvcServerIo.h"

#include <windows.h>
#include <wtsapi32.h>

#include <cstring>
#include <cwchar>
#include <vector>

namespace
{
    constexpr DWORD kAckTimeoutMs = 5000;

    enum class ExitCode : int
    {
        Success = 0,
        ChannelOpenFailed = 10,
        WriteFailed = 11,
        PartialWrite = 12,
        AckReadFailed = 13,
        StartRejected = 14,
        InvalidAck = 15
    };

    void DebugError(const wchar_t* step, DWORD error)
    {
        wchar_t buffer[256]{};
        swprintf_s(
            buffer,
            (sizeof(buffer) / sizeof(buffer[0])),
            L"PlandentRdpClientTrigger: %s fehlgeschlagen, Win32=%lu\n",
            step,
            error);
        OutputDebugStringW(buffer);
    }
}

int APIENTRY wWinMain(HINSTANCE, HINSTANCE, PWSTR, int)
{
    // WTS_CURRENT_SESSION stellt sicher, dass ausschliesslich der Client der
    // RDP-Sitzung angesprochen wird, in der diese EXE ausgefuehrt wird.
    HANDLE channel = WTSVirtualChannelOpenEx(
        WTS_CURRENT_SESSION,
        const_cast<LPSTR>(hook::protocol::kChannelName),
        WTS_CHANNEL_OPTION_DYNAMIC | WTS_CHANNEL_OPTION_DYNAMIC_PRI_MED);

    if (channel == nullptr)
    {
        DebugError(L"WTSVirtualChannelOpenEx", GetLastError());
        return static_cast<int>(ExitCode::ChannelOpenFailed);
    }

    ULONG bytesWritten = 0;
    const BOOL written = WTSVirtualChannelWrite(
        channel,
        const_cast<PCHAR>(hook::protocol::kStartMessage),
        static_cast<ULONG>(hook::protocol::kStartMessageLength),
        &bytesWritten);

    if (!written)
    {
        const DWORD error = GetLastError();
        WTSVirtualChannelClose(channel);
        DebugError(L"WTSVirtualChannelWrite", error);
        return static_cast<int>(ExitCode::WriteFailed);
    }

    if (bytesWritten != static_cast<ULONG>(hook::protocol::kStartMessageLength))
    {
        WTSVirtualChannelClose(channel);
        OutputDebugStringW(L"PlandentRdpClientTrigger: Trigger-Nachricht wurde nur teilweise geschrieben.\n");
        return static_cast<int>(ExitCode::PartialWrite);
    }

    // Wichtig: den DVC nicht direkt nach WTSVirtualChannelWrite schliessen.
    // Der Client quittiert START erst nach Verarbeitung. Dadurch bleibt der
    // Kanal lange genug bestehen, bis OnDataReceived die Nachricht erhalten hat.
    std::vector<char> response;
    DWORD readError = ERROR_SUCCESS;
    const bool readOk = hook::dvc::ReadMessage(channel, kAckTimeoutMs, response, readError);
    WTSVirtualChannelClose(channel);

    if (!readOk)
    {
        DebugError(L"START-Quittung lesen", readError);
        return static_cast<int>(ExitCode::AckReadFailed);
    }

    if (response.size() == hook::protocol::kStartOkResponseLength &&
        std::memcmp(
            response.data(),
            hook::protocol::kStartOkResponse,
            hook::protocol::kStartOkResponseLength) == 0)
    {
        return static_cast<int>(ExitCode::Success);
    }

    if (response.size() == hook::protocol::kStartFailedResponseLength &&
        std::memcmp(
            response.data(),
            hook::protocol::kStartFailedResponse,
            hook::protocol::kStartFailedResponseLength) == 0)
    {
        OutputDebugStringW(L"PlandentRdpClientTrigger: Client hat den Programmstart abgelehnt/als fehlgeschlagen gemeldet.\n");
        return static_cast<int>(ExitCode::StartRejected);
    }

    OutputDebugStringW(L"PlandentRdpClientTrigger: Ungueltige START-Quittung empfangen.\n");
    return static_cast<int>(ExitCode::InvalidAck);
}