#include "stdafx.h"
#include "OscServer.h"

#include <ws2tcpip.h>

#include <array>
#include <cstring>
#include <iomanip>
#include <sstream>
#include <vector>

#pragma comment(lib, "Ws2_32.lib")

namespace
{
bool InitializeWinsock(std::string& error)
{
	WSADATA wsaData = {};
	const int result = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if (result != 0)
	{
		error = "WSAStartup failed.";
		return false;
	}
	return true;
}

std::vector<std::string> SplitAddress(const std::string& address)
{
	std::vector<std::string> parts;
	std::size_t start = !address.empty() && address[0] == '/' ? 1 : 0;

	while (start <= address.size())
	{
		const std::size_t slash = address.find('/', start);
		const std::size_t end = slash == std::string::npos ? address.size() : slash;
		if (end > start)
			parts.push_back(address.substr(start, end - start));
		if (slash == std::string::npos)
			break;
		start = slash + 1;
	}
	return parts;
}

}

OscServer::OscServer()
	: mPort(0), mRunning(false)
{
}

OscServer::~OscServer()
{
	Stop();
}

bool OscServer::Start(unsigned short port, const Callbacks& callbacks, std::string& error)
{
	if (port == 0)
		return true;

	mCallbacks = callbacks;
	mPort = port;

	if (!InitializeWinsock(error))
		return false;

	mSocket.reset(socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP));
	if (!mSocket.valid())
	{
		error = "Could not create OSC UDP socket.";
		return false;
	}

	DWORD timeoutMilliseconds = 100;
	setsockopt(mSocket.get(), SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<const char*>(&timeoutMilliseconds), sizeof(timeoutMilliseconds));

	sockaddr_in address = {};
	address.sin_family = AF_INET;
	address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	address.sin_port = htons(static_cast<u_short>(port));
	if (bind(mSocket.get(), reinterpret_cast<sockaddr*>(&address), sizeof(address)) != 0)
	{
		error = "Could not bind OSC listener to UDP port " + std::to_string(port) + ".";
		mSocket.reset();
		return false;
	}

	mRunning = true;
	mThread = std::thread(&OscServer::ServerLoop, this);
	return true;
}

void OscServer::Stop()
{
	mRunning = false;
	mSocket.reset();
	if (mThread.joinable())
		mThread.join();
}

void OscServer::ServerLoop()
{
	std::array<char, 4096> buffer = {};
	while (mRunning)
	{
		sockaddr_in sender = {};
		int senderLength = sizeof(sender);
		const int byteCount = recvfrom(mSocket.get(), buffer.data(), static_cast<int>(buffer.size()), 0,
			reinterpret_cast<sockaddr*>(&sender), &senderLength);
		if (byteCount <= 0)
			continue;

		OscMessage message;
		std::string error;
		if (DecodeMessage(buffer.data(), byteCount, message, error))
		{
			if (!DispatchMessage(message, error) && !error.empty())
				OutputDebugStringA(("OSC dispatch failed: " + error + "\n").c_str());
		}
		else if (!error.empty())
		{
			OutputDebugStringA(("OSC decode failed: " + error + "\n").c_str());
		}
	}
}

bool OscServer::DecodeMessage(const char* data, int byteCount, OscMessage& message, std::string& error) const
{
	int offset = 0;
	if (!ReadPaddedString(data, byteCount, offset, message.address) || message.address.empty() || message.address[0] != '/')
	{
		error = "Invalid OSC address.";
		return false;
	}

	std::string typeTags;
	if (!ReadPaddedString(data, byteCount, offset, typeTags) || typeTags.empty() || typeTags[0] != ',')
	{
		error = "Invalid OSC type tag string.";
		return false;
	}

	if (typeTags.size() < 2)
	{
		error = "OSC message has no parameter value.";
		return false;
	}

	std::vector<std::string> values;
	for (std::size_t index = 1; index < typeTags.size(); ++index)
	{
		std::string valueJson;
		if (!DecodeArgument(data, byteCount, offset, typeTags[index], valueJson))
		{
			error = "Unsupported or malformed OSC value type.";
			return false;
		}
		values.push_back(valueJson);
	}

	if (values.size() == 1)
	{
		message.valueJson = values.front();
		return true;
	}

	std::ostringstream arrayJson;
	arrayJson << "[";
	for (std::size_t index = 0; index < values.size(); ++index)
	{
		if (index > 0)
			arrayJson << ",";
		arrayJson << values[index];
	}
	arrayJson << "]";
	message.valueJson = arrayJson.str();
	return true;
}

bool OscServer::DispatchMessage(const OscMessage& message, std::string& error) const
{
	const std::vector<std::string> parts = SplitAddress(message.address);
	if (parts.size() != 3 || parts[0] != "VideoShaderToys")
	{
		error = "Unsupported OSC address: " + message.address;
		return false;
	}

	return mCallbacks.updateParameter &&
		mCallbacks.updateParameter(parts[1], parts[2], message.valueJson, error);
}

bool OscServer::DecodeArgument(const char* data, int byteCount, int& offset, char valueType, std::string& valueJson)
{
	if (valueType == 'f')
	{
		double value = 0.0;
		if (!ReadFloat32(data, byteCount, offset, value))
			return false;
		std::ostringstream stream;
		stream << std::setprecision(9) << value;
		valueJson = stream.str();
		return true;
	}

	if (valueType == 'd')
	{
		double value = 0.0;
		if (!ReadFloat64(data, byteCount, offset, value))
			return false;
		std::ostringstream stream;
		stream << std::setprecision(17) << value;
		valueJson = stream.str();
		return true;
	}

	if (valueType == 'i')
	{
		int value = 0;
		if (!ReadInt32(data, byteCount, offset, value))
			return false;
		valueJson = std::to_string(value);
		return true;
	}

	if (valueType == 's')
	{
		std::string value;
		if (!ReadPaddedString(data, byteCount, offset, value))
			return false;
		valueJson = BuildJsonString(value);
		return true;
	}

	if (valueType == 'T' || valueType == 'F')
	{
		valueJson = valueType == 'T' ? "true" : "false";
		return true;
	}

	return false;
}

bool OscServer::ReadPaddedString(const char* data, int byteCount, int& offset, std::string& value)
{
	if (offset < 0 || offset >= byteCount)
		return false;

	const int start = offset;
	while (offset < byteCount && data[offset] != '\0')
		++offset;
	if (offset >= byteCount)
		return false;

	value.assign(data + start, data + offset);
	++offset;
	while (offset % 4 != 0)
		++offset;
	return offset <= byteCount;
}

bool OscServer::ReadInt32(const char* data, int byteCount, int& offset, int& value)
{
	if (offset + 4 > byteCount)
		return false;
	const unsigned char* bytes = reinterpret_cast<const unsigned char*>(data + offset);
	value = static_cast<int>((bytes[0] << 24) | (bytes[1] << 16) | (bytes[2] << 8) | bytes[3]);
	offset += 4;
	return true;
}

bool OscServer::ReadFloat32(const char* data, int byteCount, int& offset, double& value)
{
	int bits = 0;
	if (!ReadInt32(data, byteCount, offset, bits))
		return false;

	float floatValue = 0.0f;
	const unsigned int unsignedBits = static_cast<unsigned int>(bits);
	std::memcpy(&floatValue, &unsignedBits, sizeof(floatValue));
	value = static_cast<double>(floatValue);
	return true;
}

bool OscServer::ReadFloat64(const char* data, int byteCount, int& offset, double& value)
{
	if (offset + 8 > byteCount)
		return false;

	const unsigned char* bytes = reinterpret_cast<const unsigned char*>(data + offset);
	uint64_t bits = 0;
	for (int index = 0; index < 8; ++index)
		bits = (bits << 8) | static_cast<uint64_t>(bytes[index]);

	std::memcpy(&value, &bits, sizeof(value));
	offset += 8;
	return true;
}

std::string OscServer::BuildJsonString(const std::string& value)
{
	std::ostringstream stream;
	stream << '"';
	for (char ch : value)
	{
		if (ch == '"' || ch == '\\')
			stream << '\\';
		stream << ch;
	}
	stream << '"';
	return stream.str();
}