// Discord Rich Presence - plugin made by Mathew Simon.

#include "DiscordIpcConnection.h"

#include "DiscordRichPresenceModule.h"
#include "DiscordRichPresenceTypes.h"

#include "Dom/JsonObject.h"
#include "Dom/JsonValue.h"
#include "Policies/CondensedJsonPrintPolicy.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonWriter.h"

#if PLATFORM_WINDOWS
#include "Windows/AllowWindowsPlatformTypes.h"
#include <Windows.h>
#include "Windows/HideWindowsPlatformTypes.h"
#endif

namespace DiscordIpc
{
	enum EOpCode : int32
	{
		Handshake = 0,
		Frame = 1,
		Close = 2,
		Ping = 3,
		Pong = 4,
	};

	/** Discord rejette les chaines de moins de 2 caracteres. */
	static bool IsUsableText(const FString& Text)
	{
		return Text.Len() >= 2;
	}
}

FDiscordIpcConnection::~FDiscordIpcConnection()
{
	Disconnect();
}

bool FDiscordIpcConnection::Connect(const FString& InClientId)
{
#if PLATFORM_WINDOWS
	Disconnect();

	if (InClientId.IsEmpty())
	{
		return false;
	}

	ClientId = InClientId;

	// Discord numerote ses pipes : la 1ere instance prend 0, la suivante 1, etc.
	for (int32 PipeIndex = 0; PipeIndex < 10; ++PipeIndex)
	{
		const FString PipeName = FString::Printf(TEXT("\\\\.\\pipe\\discord-ipc-%d"), PipeIndex);

		HANDLE Handle = ::CreateFileW(
			*PipeName,
			GENERIC_READ | GENERIC_WRITE,
			0,
			nullptr,
			OPEN_EXISTING,
			0,
			nullptr);

		if (Handle == INVALID_HANDLE_VALUE)
		{
			continue;
		}

		PipeHandle = Handle;
		bReady = false;

		const FString Handshake = FString::Printf(TEXT("{\"v\":1,\"client_id\":\"%s\"}"), *ClientId);
		if (!WriteFrame(DiscordIpc::Handshake, Handshake))
		{
			// Pipe occupe par un autre process : on tente le suivant.
			Disconnect();
			continue;
		}

		UE_LOG(LogDiscordRichPresence, Log, TEXT("Connecte a Discord sur %s"), *PipeName);
		return true;
	}

	return false;
#else
	return false;
#endif
}

void FDiscordIpcConnection::Disconnect()
{
#if PLATFORM_WINDOWS
	if (PipeHandle)
	{
		::CloseHandle(static_cast<HANDLE>(PipeHandle));
		PipeHandle = nullptr;
	}
#endif
	bReady = false;
}

bool FDiscordIpcConnection::WriteFrame(int32 OpCode, const FString& Payload)
{
#if PLATFORM_WINDOWS
	if (!PipeHandle)
	{
		return false;
	}

	FTCHARToUTF8 Utf8Payload(*Payload);
	const int32 PayloadLength = Utf8Payload.Length();

	TArray<uint8> Frame;
	Frame.Reserve(8 + PayloadLength);
	Frame.Append(reinterpret_cast<const uint8*>(&OpCode), 4);
	Frame.Append(reinterpret_cast<const uint8*>(&PayloadLength), 4);
	Frame.Append(reinterpret_cast<const uint8*>(Utf8Payload.Get()), PayloadLength);

	DWORD BytesWritten = 0;
	const BOOL bOk = ::WriteFile(
		static_cast<HANDLE>(PipeHandle),
		Frame.GetData(),
		static_cast<DWORD>(Frame.Num()),
		&BytesWritten,
		nullptr);

	if (!bOk || BytesWritten != static_cast<DWORD>(Frame.Num()))
	{
		UE_LOG(LogDiscordRichPresence, Warning, TEXT("Ecriture vers Discord echouee, deconnexion."));
		Disconnect();
		return false;
	}

	return true;
#else
	return false;
#endif
}

bool FDiscordIpcConnection::ReadExactly(void* Buffer, int32 NumBytes)
{
#if PLATFORM_WINDOWS
	uint8* Cursor = static_cast<uint8*>(Buffer);
	int32 Remaining = NumBytes;

	while (Remaining > 0)
	{
		DWORD BytesRead = 0;
		const BOOL bOk = ::ReadFile(
			static_cast<HANDLE>(PipeHandle),
			Cursor,
			static_cast<DWORD>(Remaining),
			&BytesRead,
			nullptr);

		if (!bOk || BytesRead == 0)
		{
			return false;
		}

		Cursor += BytesRead;
		Remaining -= static_cast<int32>(BytesRead);
	}

	return true;
#else
	return false;
#endif
}

bool FDiscordIpcConnection::Poll()
{
#if PLATFORM_WINDOWS
	if (!PipeHandle)
	{
		return false;
	}

	// On ne lit que ce qui est deja arrive : ReadFile sur ce pipe est bloquant.
	for (;;)
	{
		DWORD BytesAvailable = 0;
		if (!::PeekNamedPipe(static_cast<HANDLE>(PipeHandle), nullptr, 0, nullptr, &BytesAvailable, nullptr))
		{
			UE_LOG(LogDiscordRichPresence, Log, TEXT("Discord a ferme la connexion."));
			Disconnect();
			return false;
		}

		if (BytesAvailable < 8)
		{
			return true;
		}

		uint8 Header[8];
		if (!ReadExactly(Header, 8))
		{
			Disconnect();
			return false;
		}

		int32 OpCode = 0;
		int32 Length = 0;
		FMemory::Memcpy(&OpCode, Header, 4);
		FMemory::Memcpy(&Length, Header + 4, 4);

		if (Length < 0 || Length > 64 * 1024)
		{
			UE_LOG(LogDiscordRichPresence, Warning, TEXT("Frame Discord invalide (taille %d)."), Length);
			Disconnect();
			return false;
		}

		FString Payload;
		if (Length > 0)
		{
			TArray<uint8> PayloadBytes;
			PayloadBytes.SetNumUninitialized(Length + 1);
			if (!ReadExactly(PayloadBytes.GetData(), Length))
			{
				Disconnect();
				return false;
			}
			PayloadBytes[Length] = 0;
			Payload = UTF8_TO_TCHAR(reinterpret_cast<const ANSICHAR*>(PayloadBytes.GetData()));
		}

		HandleIncomingFrame(OpCode, Payload);

		if (!PipeHandle)
		{
			return false;
		}
	}
#else
	return false;
#endif
}

void FDiscordIpcConnection::HandleIncomingFrame(int32 OpCode, const FString& Payload)
{
	switch (OpCode)
	{
	case DiscordIpc::Ping:
		WriteFrame(DiscordIpc::Pong, Payload);
		return;

	case DiscordIpc::Close:
		UE_LOG(LogDiscordRichPresence, Log, TEXT("Discord a demande la fermeture : %s"), *Payload);
		Disconnect();
		return;

	case DiscordIpc::Frame:
		break;

	default:
		return;
	}

	TSharedPtr<FJsonObject> Root;
	const TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(Payload);
	if (!FJsonSerializer::Deserialize(Reader, Root) || !Root.IsValid())
	{
		return;
	}

	FString Event;
	if (!Root->TryGetStringField(TEXT("evt"), Event))
	{
		return;
	}

	if (Event == TEXT("READY"))
	{
		bReady = true;
		UE_LOG(LogDiscordRichPresence, Log, TEXT("Rich Presence prete."));
	}
	else if (Event == TEXT("ERROR"))
	{
		UE_LOG(LogDiscordRichPresence, Warning, TEXT("Discord a renvoye une erreur : %s"), *Payload);
	}
}

FString FDiscordIpcConnection::BuildActivityPayload(const FDiscordRichPresenceActivity* Activity, int64 StartUnixTime)
{
	const TSharedRef<FJsonObject> Args = MakeShared<FJsonObject>();
	Args->SetNumberField(TEXT("pid"), static_cast<double>(FPlatformProcess::GetCurrentProcessId()));

	if (Activity == nullptr)
	{
		// activity: null => Discord retire le jeu du statut.
		Args->SetField(TEXT("activity"), MakeShared<FJsonValueNull>());
	}
	else
	{
		const TSharedRef<FJsonObject> ActivityObject = MakeShared<FJsonObject>();

		if (DiscordIpc::IsUsableText(Activity->Details))
		{
			ActivityObject->SetStringField(TEXT("details"), Activity->Details);
		}
		if (DiscordIpc::IsUsableText(Activity->State))
		{
			ActivityObject->SetStringField(TEXT("state"), Activity->State);
		}

		if (Activity->bShowElapsedTime && StartUnixTime > 0)
		{
			const TSharedRef<FJsonObject> Timestamps = MakeShared<FJsonObject>();
			Timestamps->SetNumberField(TEXT("start"), static_cast<double>(StartUnixTime));
			ActivityObject->SetObjectField(TEXT("timestamps"), Timestamps);
		}

		const bool bHasAssets = !Activity->LargeImageKey.IsEmpty() || !Activity->SmallImageKey.IsEmpty();
		if (bHasAssets)
		{
			const TSharedRef<FJsonObject> Assets = MakeShared<FJsonObject>();
			if (!Activity->LargeImageKey.IsEmpty())
			{
				Assets->SetStringField(TEXT("large_image"), Activity->LargeImageKey);
				if (DiscordIpc::IsUsableText(Activity->LargeImageText))
				{
					Assets->SetStringField(TEXT("large_text"), Activity->LargeImageText);
				}
			}
			if (!Activity->SmallImageKey.IsEmpty())
			{
				Assets->SetStringField(TEXT("small_image"), Activity->SmallImageKey);
				if (DiscordIpc::IsUsableText(Activity->SmallImageText))
				{
					Assets->SetStringField(TEXT("small_text"), Activity->SmallImageText);
				}
			}
			ActivityObject->SetObjectField(TEXT("assets"), Assets);
		}

		if (Activity->PartySize > 0)
		{
			TArray<TSharedPtr<FJsonValue>> Sizes;
			Sizes.Add(MakeShared<FJsonValueNumber>(Activity->PartySize));
			Sizes.Add(MakeShared<FJsonValueNumber>(FMath::Max(Activity->PartyMax, Activity->PartySize)));

			const TSharedRef<FJsonObject> Party = MakeShared<FJsonObject>();
			Party->SetArrayField(TEXT("size"), Sizes);
			ActivityObject->SetObjectField(TEXT("party"), Party);
		}

		Args->SetObjectField(TEXT("activity"), ActivityObject);
	}

	const TSharedRef<FJsonObject> Root = MakeShared<FJsonObject>();
	Root->SetStringField(TEXT("cmd"), TEXT("SET_ACTIVITY"));
	Root->SetStringField(TEXT("nonce"), FString::Printf(TEXT("%d"), ++NonceCounter));
	Root->SetObjectField(TEXT("args"), Args);

	FString Output;
	const TSharedRef<TJsonWriter<TCHAR, TCondensedJsonPrintPolicy<TCHAR>>> Writer =
		TJsonWriterFactory<TCHAR, TCondensedJsonPrintPolicy<TCHAR>>::Create(&Output);
	FJsonSerializer::Serialize(Root, Writer);

	return Output;
}

bool FDiscordIpcConnection::SendActivity(const FDiscordRichPresenceActivity& Activity, int64 StartUnixTime)
{
	if (!IsConnected() || !bReady)
	{
		return false;
	}

	return WriteFrame(DiscordIpc::Frame, BuildActivityPayload(&Activity, StartUnixTime));
}

bool FDiscordIpcConnection::ClearActivity()
{
	if (!IsConnected() || !bReady)
	{
		return false;
	}

	return WriteFrame(DiscordIpc::Frame, BuildActivityPayload(nullptr, 0));
}
