#include "HttpControlServer.h"

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

#include <ws2tcpip.h>

#include <algorithm>
#include <cctype>
#include <fstream>
#include <sstream>

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";
}
}

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& docsRoot,
	HttpControlServerConfig config,
	HttpControlServerCallbacks callbacks,
	std::string& error)
{
	Stop();

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

	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();

	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::ThreadMain()
{
	while (mRunning.load(std::memory_order_acquire))
	{
		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"));

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

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 TextResponse("200 OK", "RenderCadenceCompositor control server");
	return TextResponse("404 Not Found", "Not Found");
}

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

	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() };
}

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 };
}

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();
}

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";
	return "text/plain";
}

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();
}
}