#include "HttpControlServer.h"

#include "../json/JsonWriter.h"
#include "../logging/Logger.h"

#include <ws2tcpip.h>

#include <algorithm>
#include <array>
#include <cctype>
#include <cstdint>
#include <fstream>
#include <sstream>
#include <vector>

namespace RenderCadenceCompositor
{
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;
}

bool IsKnownPostEndpoint(const std::string& path)
{
	return path == "/api/layers/add"
		|| path == "/api/layers/remove"
		|| path == "/api/layers/move"
		|| path == "/api/layers/reorder"
		|| path == "/api/layers/set-bypass"
		|| path == "/api/layers/set-shader"
		|| path == "/api/layers/update-parameter"
		|| path == "/api/layers/reset-parameters"
		|| path == "/api/stack-presets/save"
		|| path == "/api/stack-presets/load"
		|| path == "/api/reload"
		|| path == "/api/screenshot";
}

std::array<uint8_t, 20> Sha1(const std::string& input)
{
	auto leftRotate = [](uint32_t value, uint32_t bits) {
		return (value << bits) | (value >> (32U - bits));
	};

	std::vector<uint8_t> data(input.begin(), input.end());
	const uint64_t bitLength = static_cast<uint64_t>(data.size()) * 8ULL;
	data.push_back(0x80);
	while ((data.size() % 64) != 56)
		data.push_back(0);
	for (int shift = 56; shift >= 0; shift -= 8)
		data.push_back(static_cast<uint8_t>((bitLength >> shift) & 0xff));

	uint32_t h0 = 0x67452301;
	uint32_t h1 = 0xefcdab89;
	uint32_t h2 = 0x98badcfe;
	uint32_t h3 = 0x10325476;
	uint32_t h4 = 0xc3d2e1f0;

	for (std::size_t offset = 0; offset < data.size(); offset += 64)
	{
		uint32_t words[80] = {};
		for (std::size_t i = 0; i < 16; ++i)
		{
			const std::size_t index = offset + i * 4;
			words[i] = (static_cast<uint32_t>(data[index]) << 24)
				| (static_cast<uint32_t>(data[index + 1]) << 16)
				| (static_cast<uint32_t>(data[index + 2]) << 8)
				| static_cast<uint32_t>(data[index + 3]);
		}
		for (std::size_t i = 16; i < 80; ++i)
			words[i] = leftRotate(words[i - 3] ^ words[i - 8] ^ words[i - 14] ^ words[i - 16], 1);

		uint32_t a = h0;
		uint32_t b = h1;
		uint32_t c = h2;
		uint32_t d = h3;
		uint32_t e = h4;

		for (std::size_t i = 0; i < 80; ++i)
		{
			uint32_t f = 0;
			uint32_t k = 0;
			if (i < 20)
			{
				f = (b & c) | ((~b) & d);
				k = 0x5a827999;
			}
			else if (i < 40)
			{
				f = b ^ c ^ d;
				k = 0x6ed9eba1;
			}
			else if (i < 60)
			{
				f = (b & c) | (b & d) | (c & d);
				k = 0x8f1bbcdc;
			}
			else
			{
				f = b ^ c ^ d;
				k = 0xca62c1d6;
			}

			const uint32_t temp = leftRotate(a, 5) + f + e + k + words[i];
			e = d;
			d = c;
			c = leftRotate(b, 30);
			b = a;
			a = temp;
		}

		h0 += a;
		h1 += b;
		h2 += c;
		h3 += d;
		h4 += e;
	}

	std::array<uint8_t, 20> digest = {};
	const uint32_t parts[] = { h0, h1, h2, h3, h4 };
	for (std::size_t i = 0; i < 5; ++i)
	{
		digest[i * 4] = static_cast<uint8_t>((parts[i] >> 24) & 0xff);
		digest[i * 4 + 1] = static_cast<uint8_t>((parts[i] >> 16) & 0xff);
		digest[i * 4 + 2] = static_cast<uint8_t>((parts[i] >> 8) & 0xff);
		digest[i * 4 + 3] = static_cast<uint8_t>(parts[i] & 0xff);
	}
	return digest;
}

std::string Base64Encode(const uint8_t* data, std::size_t size)
{
	static constexpr char kAlphabet[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
	std::string output;
	output.reserve(((size + 2) / 3) * 4);
	for (std::size_t i = 0; i < size; i += 3)
	{
		const uint32_t a = data[i];
		const uint32_t b = i + 1 < size ? data[i + 1] : 0;
		const uint32_t c = i + 2 < size ? data[i + 2] : 0;
		const uint32_t triple = (a << 16) | (b << 8) | c;
		output.push_back(kAlphabet[(triple >> 18) & 0x3f]);
		output.push_back(kAlphabet[(triple >> 12) & 0x3f]);
		output.push_back(i + 1 < size ? kAlphabet[(triple >> 6) & 0x3f] : '=');
		output.push_back(i + 2 < size ? kAlphabet[triple & 0x3f] : '=');
	}
	return output;
}
}

UniqueSocket::UniqueSocket(SOCKET socket) :
	mSocket(socket)
{
}

UniqueSocket::~UniqueSocket()
{
	reset();
}

UniqueSocket::UniqueSocket(UniqueSocket&& other) noexcept :
	mSocket(other.release())
{
}

UniqueSocket& UniqueSocket::operator=(UniqueSocket&& other) noexcept
{
	if (this != &other)
		reset(other.release());
	return *this;
}

SOCKET UniqueSocket::release()
{
	const SOCKET socket = mSocket;
	mSocket = INVALID_SOCKET;
	return socket;
}

void UniqueSocket::reset(SOCKET socket)
{
	if (valid())
		closesocket(mSocket);
	mSocket = socket;
}

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

bool HttpControlServer::Start(
	const std::filesystem::path& uiRoot,
	const std::filesystem::path& docsRoot,
	HttpControlServerConfig config,
	HttpControlServerCallbacks callbacks,
	std::string& error)
{
	Stop();

	if (!InitializeWinsock(error))
		return false;
	mWinsockStarted = true;

	mUiRoot = uiRoot;
	mDocsRoot = docsRoot;
	mConfig = config;
	mCallbacks = std::move(callbacks);

	mListenSocket.reset(socket(AF_INET, SOCK_STREAM, IPPROTO_TCP));
	if (!mListenSocket.valid())
	{
		error = "Could not create HTTP control server socket.";
		Stop();
		return false;
	}

	u_long nonBlocking = 1;
	ioctlsocket(mListenSocket.get(), FIONBIO, &nonBlocking);

	sockaddr_in address = {};
	address.sin_family = AF_INET;
	address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);

	bool bound = false;
	for (unsigned short offset = 0; offset < mConfig.portSearchCount; ++offset)
	{
		address.sin_port = htons(static_cast<u_short>(mConfig.preferredPort + offset));
		if (bind(mListenSocket.get(), reinterpret_cast<sockaddr*>(&address), sizeof(address)) == 0)
		{
			mPort = static_cast<unsigned short>(mConfig.preferredPort + offset);
			bound = true;
			break;
		}
	}

	if (!bound)
	{
		error = "Could not bind HTTP control server to loopback.";
		Stop();
		return false;
	}

	if (listen(mListenSocket.get(), SOMAXCONN) != 0)
	{
		error = "Could not listen on HTTP control server socket.";
		Stop();
		return false;
	}

	mRunning.store(true, std::memory_order_release);
	mThread = std::thread([this]() { ThreadMain(); });
	Log("http", "HTTP control server listening on http://127.0.0.1:" + std::to_string(mPort));
	return true;
}

void HttpControlServer::Stop()
{
	mRunning.store(false, std::memory_order_release);
	mListenSocket.reset();

	if (mThread.joinable())
		mThread.join();

	std::vector<std::thread> clientThreads;
	{
		std::lock_guard<std::mutex> lock(mClientThreadsMutex);
		clientThreads.swap(mClientThreads);
		for (std::thread& thread : mFinishedClientThreads)
			clientThreads.push_back(std::move(thread));
		mFinishedClientThreads.clear();
	}
	for (std::thread& thread : clientThreads)
	{
		if (thread.joinable())
			thread.join();
	}

	if (mWinsockStarted)
	{
		WSACleanup();
		mWinsockStarted = false;
	}
	mPort = 0;
}

HttpControlServer::HttpResponse HttpControlServer::RouteRequestForTest(const HttpRequest& request) const
{
	return RouteRequest(request);
}

void HttpControlServer::SetCallbacksForTest(HttpControlServerCallbacks callbacks)
{
	mCallbacks = std::move(callbacks);
}

void HttpControlServer::SetRootsForTest(const std::filesystem::path& uiRoot, const std::filesystem::path& docsRoot)
{
	mUiRoot = uiRoot;
	mDocsRoot = docsRoot;
}

void HttpControlServer::ThreadMain()
{
	while (mRunning.load(std::memory_order_acquire))
	{
		JoinFinishedClientThreads();
		TryAcceptClient();
		std::this_thread::sleep_for(mConfig.idleSleep);
	}
}

bool HttpControlServer::TryAcceptClient()
{
	sockaddr_in clientAddress = {};
	int addressSize = sizeof(clientAddress);
	UniqueSocket clientSocket(accept(mListenSocket.get(), reinterpret_cast<sockaddr*>(&clientAddress), &addressSize));
	if (!clientSocket.valid())
		return false;

	return HandleClient(std::move(clientSocket));
}

bool HttpControlServer::HandleClient(UniqueSocket clientSocket)
{
	char buffer[16384];
	const int received = recv(clientSocket.get(), buffer, sizeof(buffer), 0);
	if (received <= 0)
		return false;

	HttpRequest request;
	if (!ParseHttpRequest(std::string(buffer, buffer + received), request))
		return SendResponse(clientSocket.get(), TextResponse("400 Bad Request", "Bad Request"));

	if (request.path == "/ws")
		return HandleWebSocketClient(std::move(clientSocket), request);

	return SendResponse(clientSocket.get(), RouteRequest(request));
}

bool HttpControlServer::HandleWebSocketClient(UniqueSocket clientSocket, const HttpRequest& request)
{
	const auto keyIt = request.headers.find("sec-websocket-key");
	if (keyIt == request.headers.end() || keyIt->second.empty())
		return SendResponse(clientSocket.get(), TextResponse("400 Bad Request", "Missing WebSocket key"));

	std::ostringstream stream;
	stream << "HTTP/1.1 101 Switching Protocols\r\n"
		<< "Upgrade: websocket\r\n"
		<< "Connection: Upgrade\r\n"
		<< "Sec-WebSocket-Accept: " << WebSocketAcceptKey(keyIt->second) << "\r\n\r\n";
	const std::string response = stream.str();
	if (send(clientSocket.get(), response.c_str(), static_cast<int>(response.size()), 0) != static_cast<int>(response.size()))
		return false;

	u_long nonBlocking = 1;
	ioctlsocket(clientSocket.get(), FIONBIO, &nonBlocking);

	std::thread thread([this, socket = std::move(clientSocket)]() mutable {
		WebSocketClientMain(std::move(socket));
	});
	{
		std::lock_guard<std::mutex> lock(mClientThreadsMutex);
		mClientThreads.push_back(std::move(thread));
	}
	return true;
}

void HttpControlServer::WebSocketClientMain(UniqueSocket clientSocket)
{
	std::string previousState;
	while (mRunning.load(std::memory_order_acquire))
	{
		const std::string state = mCallbacks.getStateJson ? mCallbacks.getStateJson() : "{}";
		if (state != previousState)
		{
			if (!SendWebSocketText(clientSocket.get(), state))
				break;
			previousState = state;
		}
		std::this_thread::sleep_for(std::chrono::milliseconds(250));
	}

	std::lock_guard<std::mutex> lock(mClientThreadsMutex);
	const std::thread::id currentId = std::this_thread::get_id();
	for (auto it = mClientThreads.begin(); it != mClientThreads.end(); ++it)
	{
		if (it->get_id() != currentId)
			continue;
		mFinishedClientThreads.push_back(std::move(*it));
		mClientThreads.erase(it);
		break;
	}
}

void HttpControlServer::JoinFinishedClientThreads()
{
	std::vector<std::thread> finished;
	{
		std::lock_guard<std::mutex> lock(mClientThreadsMutex);
		finished.swap(mFinishedClientThreads);
	}
	for (std::thread& thread : finished)
	{
		if (thread.joinable())
			thread.join();
	}
}

bool HttpControlServer::SendResponse(SOCKET clientSocket, const HttpResponse& response) const
{
	std::ostringstream stream;
	stream << "HTTP/1.1 " << response.status << "\r\n"
		<< "Content-Type: " << response.contentType << "\r\n"
		<< "Content-Length: " << response.body.size() << "\r\n"
		<< "Access-Control-Allow-Origin: *\r\n"
		<< "Connection: close\r\n\r\n"
		<< response.body;

	const std::string payload = stream.str();
	return send(clientSocket, payload.c_str(), static_cast<int>(payload.size()), 0) == static_cast<int>(payload.size());
}

HttpControlServer::HttpResponse HttpControlServer::RouteRequest(const HttpRequest& request) const
{
	if (request.method == "GET")
		return ServeGet(request);
	if (request.method == "POST")
		return ServePost(request);
	if (request.method == "OPTIONS")
		return TextResponse("204 No Content", std::string());
	return TextResponse("404 Not Found", "Not Found");
}

HttpControlServer::HttpResponse HttpControlServer::ServeGet(const HttpRequest& request) const
{
	if (request.path == "/api/state")
		return JsonResponse("200 OK", mCallbacks.getStateJson ? mCallbacks.getStateJson() : "{}");
	if (request.path == "/openapi.yaml" || request.path == "/docs/openapi.yaml")
		return ServeOpenApiSpec();
	if (request.path == "/docs" || request.path == "/docs/")
		return ServeSwaggerDocs();
	if (request.path == "/" || request.path == "/index.html")
		return ServeUiAsset("index.html");
	if (request.path.rfind("/assets/", 0) == 0)
		return ServeUiAsset(request.path.substr(1));
	if (request.path.size() > 1)
	{
		const HttpResponse asset = ServeUiAsset(request.path.substr(1));
		if (asset.status != "404 Not Found")
			return asset;
	}
	return ServeUiAsset("index.html");
}

HttpControlServer::HttpResponse HttpControlServer::ServePost(const HttpRequest& request) const
{
	if (!IsKnownPostEndpoint(request.path))
		return TextResponse("404 Not Found", "Not Found");

	if (request.path == "/api/layers/add" && mCallbacks.addLayer)
	{
		const ControlActionResult result = mCallbacks.addLayer(request.body);
		return JsonResponse(result.ok ? "200 OK" : "400 Bad Request", ActionResponse(result.ok, result.error));
	}

	if (request.path == "/api/layers/remove" && mCallbacks.removeLayer)
	{
		const ControlActionResult result = mCallbacks.removeLayer(request.body);
		return JsonResponse(result.ok ? "200 OK" : "400 Bad Request", ActionResponse(result.ok, result.error));
	}

	return {
		"400 Bad Request",
		"application/json",
		ActionResponse(false, "Endpoint is not implemented in RenderCadenceCompositor yet.")
	};
}

HttpControlServer::HttpResponse HttpControlServer::ServeOpenApiSpec() const
{
	const std::filesystem::path path = mDocsRoot / "openapi.yaml";
	const std::string body = LoadTextFile(path);
	return body.empty()
		? TextResponse("404 Not Found", "OpenAPI spec not found")
		: HttpResponse{ "200 OK", GuessContentType(path), body };
}

HttpControlServer::HttpResponse HttpControlServer::ServeSwaggerDocs() const
{
	std::ostringstream html;
	html << "<!doctype html>\n"
		<< "<html><head><meta charset=\"utf-8\"><title>Video Shader Toys API Docs</title>\n"
		<< "<link rel=\"stylesheet\" href=\"https://unpkg.com/swagger-ui-dist@5/swagger-ui.css\"></head>\n"
		<< "<body><div id=\"swagger-ui\"></div>\n"
		<< "<script src=\"https://unpkg.com/swagger-ui-dist@5/swagger-ui-bundle.js\"></script>\n"
		<< "<script>SwaggerUIBundle({url:'/docs/openapi.yaml',dom_id:'#swagger-ui'});</script>\n"
		<< "</body></html>\n";
	return { "200 OK", "text/html", html.str() };
}

HttpControlServer::HttpResponse HttpControlServer::ServeUiAsset(const std::string& relativePath) const
{
	if (mUiRoot.empty())
		return TextResponse("404 Not Found", "UI root is not configured");

	const std::filesystem::path sanitizedPath = std::filesystem::path(relativePath).lexically_normal();
	if (!IsSafeRelativePath(sanitizedPath))
		return TextResponse("404 Not Found", "Not Found");

	const std::filesystem::path path = mUiRoot / sanitizedPath;
	const std::string body = LoadTextFile(path);
	if (body.empty())
		return TextResponse("404 Not Found", "Not Found");
	return { "200 OK", GuessContentType(path), body };
}

std::string HttpControlServer::LoadTextFile(const std::filesystem::path& path) const
{
	std::ifstream input(path, std::ios::binary);
	if (!input)
		return std::string();

	std::ostringstream buffer;
	buffer << input.rdbuf();
	return buffer.str();
}

HttpControlServer::HttpResponse HttpControlServer::JsonResponse(const std::string& status, const std::string& body)
{
	return { status, "application/json", body };
}

HttpControlServer::HttpResponse HttpControlServer::TextResponse(const std::string& status, const std::string& body)
{
	return { status, "text/plain", body };
}

HttpControlServer::HttpResponse HttpControlServer::HtmlResponse(const std::string& status, const std::string& body)
{
	return { status, "text/html", body };
}

std::string HttpControlServer::ActionResponse(bool ok, const std::string& error)
{
	JsonWriter writer;
	writer.BeginObject();
	writer.KeyBool("ok", ok);
	if (!error.empty())
		writer.KeyString("error", error);
	writer.EndObject();
	return writer.StringValue();
}

bool HttpControlServer::SendWebSocketText(SOCKET clientSocket, const std::string& text)
{
	if (clientSocket == INVALID_SOCKET)
		return false;

	std::vector<unsigned char> frame;
	frame.reserve(text.size() + 16);
	frame.push_back(0x81);
	if (text.size() <= 125)
	{
		frame.push_back(static_cast<unsigned char>(text.size()));
	}
	else if (text.size() <= 0xffff)
	{
		frame.push_back(126);
		frame.push_back(static_cast<unsigned char>((text.size() >> 8) & 0xff));
		frame.push_back(static_cast<unsigned char>(text.size() & 0xff));
	}
	else
	{
		frame.push_back(127);
		const uint64_t length = static_cast<uint64_t>(text.size());
		for (int shift = 56; shift >= 0; shift -= 8)
			frame.push_back(static_cast<unsigned char>((length >> shift) & 0xff));
	}
	frame.insert(frame.end(), text.begin(), text.end());

	const char* data = reinterpret_cast<const char*>(frame.data());
	int remaining = static_cast<int>(frame.size());
	while (remaining > 0)
	{
		const int sent = send(clientSocket, data, remaining, 0);
		if (sent <= 0)
		{
			const int error = WSAGetLastError();
			if (error == WSAEWOULDBLOCK)
			{
				std::this_thread::sleep_for(std::chrono::milliseconds(2));
				continue;
			}
			return false;
		}
		data += sent;
		remaining -= sent;
	}
	return true;
}

std::string HttpControlServer::WebSocketAcceptKey(const std::string& clientKey)
{
	static constexpr const char* kWebSocketGuid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
	const std::array<uint8_t, 20> digest = Sha1(clientKey + kWebSocketGuid);
	return Base64Encode(digest.data(), digest.size());
}

std::string HttpControlServer::GuessContentType(const std::filesystem::path& path)
{
	const std::string extension = ToLower(path.extension().string());
	if (extension == ".yaml" || extension == ".yml")
		return "application/yaml";
	if (extension == ".json")
		return "application/json";
	if (extension == ".js" || extension == ".mjs")
		return "text/javascript";
	if (extension == ".css")
		return "text/css";
	if (extension == ".html" || extension == ".htm")
		return "text/html";
	if (extension == ".svg")
		return "image/svg+xml";
	if (extension == ".png")
		return "image/png";
	if (extension == ".jpg" || extension == ".jpeg")
		return "image/jpeg";
	if (extension == ".ico")
		return "image/x-icon";
	if (extension == ".map")
		return "application/json";
	return "text/plain";
}

bool HttpControlServer::IsSafeRelativePath(const std::filesystem::path& path)
{
	if (path.empty() || path.is_absolute())
		return false;

	for (const std::filesystem::path& part : path)
	{
		if (part == "..")
			return false;
	}
	return true;
}

std::string HttpControlServer::ToLower(std::string text)
{
	std::transform(text.begin(), text.end(), text.begin(), [](unsigned char character) {
		return static_cast<char>(std::tolower(character));
	});
	return text;
}

bool HttpControlServer::ParseHttpRequest(const std::string& rawRequest, HttpRequest& request)
{
	const std::size_t requestLineEnd = rawRequest.find("\r\n");
	if (requestLineEnd == std::string::npos)
		return false;

	const std::string requestLine = rawRequest.substr(0, requestLineEnd);
	const std::size_t methodEnd = requestLine.find(' ');
	if (methodEnd == std::string::npos)
		return false;

	const std::size_t pathEnd = requestLine.find(' ', methodEnd + 1);
	if (pathEnd == std::string::npos)
		return false;

	request.method = requestLine.substr(0, methodEnd);
	request.path = requestLine.substr(methodEnd + 1, pathEnd - methodEnd - 1);
	request.headers.clear();

	const std::size_t queryStart = request.path.find('?');
	if (queryStart != std::string::npos)
		request.path = request.path.substr(0, queryStart);

	const std::size_t headersStart = requestLineEnd + 2;
	const std::size_t bodySeparator = rawRequest.find("\r\n\r\n", headersStart);
	const std::size_t headersEnd = bodySeparator == std::string::npos ? rawRequest.size() : bodySeparator;

	for (std::size_t lineStart = headersStart; lineStart < headersEnd;)
	{
		const std::size_t lineEnd = rawRequest.find("\r\n", lineStart);
		const std::size_t currentLineEnd = lineEnd == std::string::npos ? headersEnd : (std::min)(lineEnd, headersEnd);
		const std::string line = rawRequest.substr(lineStart, currentLineEnd - lineStart);
		const std::size_t separator = line.find(':');
		if (separator != std::string::npos)
		{
			const std::string key = ToLower(line.substr(0, separator));
			std::string value = line.substr(separator + 1);
			const std::size_t first = value.find_first_not_of(" \t");
			const std::size_t last = value.find_last_not_of(" \t");
			request.headers[key] = first == std::string::npos ? std::string() : value.substr(first, last - first + 1);
		}

		if (lineEnd == std::string::npos || lineEnd >= headersEnd)
			break;
		lineStart = lineEnd + 2;
	}

	request.body = bodySeparator == std::string::npos ? std::string() : rawRequest.substr(bodySeparator + 4);
	return !request.method.empty() && !request.path.empty();
}
}