#include "wj_server.h"
#include "wj_input.h"
#include "wj_lanip.h"
#include "wj_page.h"
#include "civetweb.h"

#include <atomic>
#include <cstring>
#include <cstdio>
#include <mutex>
#include <random>
#include <unordered_map>

#include "c_cvars.h"
#include "c_dispatch.h"
#include "printf.h"

CVAR(Int, wj_port, 5029, CVAR_ARCHIVE)

namespace
{
	mg_context *g_ctx = nullptr;
	std::string g_sessionToken;
	std::vector<std::string> g_lanAddresses;
	int g_listenPort = 0;
	std::atomic<int> g_nextConnToken{ 1 };

	// Per-WebSocket-connection state, owned by the connection via
	// mg_set_user_connection_data / mg_get_user_connection_data. `slot` is
	// -1 until the main thread admits this connection to a player slot
	// (WJ_OnSlotAssigned); only ever read/written through the atomic, since
	// the data handler (a civetweb thread) reads it while the main thread
	// may be writing it.
	struct ConnState
	{
		int connToken;
		std::atomic<int> slot{ -1 };
	};

	// connToken -> ConnState*, so WJ_OnSlotAssigned/WJ_OnSlotFreed (called
	// from the main thread via wj_ticcmd.cpp) can find a connection's state
	// without holding a raw mg_connection* that might already have been
	// torn down by civetweb on another thread. Entries are added in the
	// websocket ready handler and removed (then the ConnState is deleted)
	// in the close handler, both under g_connMutex.
	std::mutex g_connMutex;
	std::unordered_map<int, ConnState *> g_liveConns;

	std::string MakeSessionToken()
	{
		std::random_device rd;
		std::mt19937_64 gen(rd());
		std::uniform_int_distribution<uint64_t> dist;
		char buf[32];
		snprintf(buf, sizeof(buf), "%016llx", (unsigned long long)dist(gen));
		return buf;
	}

	// Pulls the value of `key` out of a "k=v&k2=v2" query string. Small and
	// case-sensitive on purpose -- this only ever needs to read our own `t=`
	// token that our own page generates.
	std::string QueryParam(const char *query, const char *key)
	{
		if (!query) return {};
		std::string q(query);
		std::string needle = std::string(key) + "=";
		size_t pos = 0;
		while (pos < q.size())
		{
			size_t amp = q.find('&', pos);
			if (amp == std::string::npos) amp = q.size();
			std::string kv = q.substr(pos, amp - pos);
			if (kv.compare(0, needle.size(), needle) == 0)
				return kv.substr(needle.size());
			pos = amp + 1;
		}
		return {};
	}

	int RequestHandler_JoinPage(mg_connection *conn, void *)
	{
		const mg_request_info *ri = mg_get_request_info(conn);
		std::string token = QueryParam(ri->query_string, "t");
		if (token != g_sessionToken)
		{
			mg_send_http_error(conn, 403, "%s", "invalid or missing join token");
			return 403;
		}

		const char *body = WJ_GetPageHTML();
		mg_send_http_ok(conn, "text/html; charset=utf-8", (long long)strlen(body));
		mg_write(conn, body, strlen(body));
		return 200;
	}

	int WS_Connect(const mg_connection *conn, void *)
	{
		const mg_request_info *ri = mg_get_request_info(conn);
		std::string token = QueryParam(ri->query_string, "t");
		if (token != g_sessionToken)
			return 1; // reject: close immediately
		return 0; // accept
	}

	void WS_Ready(mg_connection *conn, void *)
	{
		auto *cs = new ConnState();
		cs->connToken = g_nextConnToken.fetch_add(1);
		mg_set_user_connection_data(conn, cs);

		{
			std::lock_guard<std::mutex> lock(g_connMutex);
			g_liveConns[cs->connToken] = cs;
		}
		{
			std::lock_guard<std::mutex> lock(wj_joinQueue.mtx);
			wj_joinQueue.pending.push_back(WJPendingJoin{ cs->connToken, std::string() });
		}
	}

	int WS_Data(mg_connection *conn, int opcode, char *data, size_t len, void *)
	{
		if ((opcode & 0xf) != MG_WEBSOCKET_OPCODE_BINARY)
			return 1; // ignore text/ping/etc, keep connection open

		auto *cs = static_cast<ConnState *>(mg_get_user_connection_data(conn));
		if (!cs || len < sizeof(WJInputPacket))
			return 1;

		int slot = cs->slot.load();
		if (slot < 0 || slot >= MAXPLAYERS)
			return 1; // not admitted to a player slot yet

		WJInputPacket packet;
		memcpy(&packet, data, sizeof(packet));

		WJRemoteInput &ri = wj_remoteInput[slot];
		std::lock_guard<std::mutex> lock(ri.mtx);
		ri.connected = true;
		ri.packet = packet;
		return 1;
	}

	void WS_Close(const mg_connection *conn, void *)
	{
		auto *cs = static_cast<ConnState *>(mg_get_user_connection_data(conn));
		if (!cs) return;

		int slot = cs->slot.load();
		if (slot >= 0 && slot < MAXPLAYERS)
		{
			WJRemoteInput &ri = wj_remoteInput[slot];
			std::lock_guard<std::mutex> lock(ri.mtx);
			ri.connected = false;
		}

		{
			std::lock_guard<std::mutex> lock(g_connMutex);
			g_liveConns.erase(cs->connToken);
		}
		delete cs;
	}
}

bool WJ_OnSlotAssigned(int connToken, int slot)
{
	std::lock_guard<std::mutex> lock(g_connMutex);
	auto it = g_liveConns.find(connToken);
	if (it == g_liveConns.end())
		return false;
	it->second->slot.store(slot);
	return true;
}

bool WJ_StartServer(int port)
{
	if (g_ctx != nullptr)
		return true;

	if (port <= 0)
		port = *wj_port;

	g_sessionToken = MakeSessionToken();
	g_lanAddresses = WJ_GetLocalIPv4Addresses();

	char portStr[16];
	snprintf(portStr, sizeof(portStr), "%d", port);

	const char *options[] = {
		"listening_ports", portStr,
		"num_threads", "40",
		"enable_keep_alive", "yes",
		nullptr
	};

	mg_callbacks callbacks;
	memset(&callbacks, 0, sizeof(callbacks));

	g_ctx = mg_start(&callbacks, nullptr, options);
	if (!g_ctx)
	{
		Printf("webjoin: failed to start server on port %d\n", port);
		return false;
	}

	g_listenPort = port;
	mg_set_request_handler(g_ctx, "/", RequestHandler_JoinPage, nullptr);
	mg_set_websocket_handler(g_ctx, "/ws", WS_Connect, WS_Ready, WS_Data, WS_Close, nullptr);

	Printf("webjoin: listening on port %d (%d LAN address%s found)\n",
		port, (int)g_lanAddresses.size(), g_lanAddresses.size() == 1 ? "" : "es");
	return true;
}

void WJ_StopServer()
{
	if (!g_ctx) return;
	mg_stop(g_ctx);
	g_ctx = nullptr;

	// Any connections civetweb had are gone now; drop whatever slots were
	// still marked connected so stale input doesn't linger.
	for (int i = 0; i < MAXPLAYERS; i++)
	{
		WJRemoteInput &ri = wj_remoteInput[i];
		std::lock_guard<std::mutex> lock(ri.mtx);
		ri.connected = false;
		ri.packet = {};
	}

	std::lock_guard<std::mutex> lock(g_connMutex);
	g_liveConns.clear();
}

bool WJ_IsServerRunning()
{
	return g_ctx != nullptr;
}

int WJ_GetJoinURLCount()
{
	return (int)g_lanAddresses.size();
}

std::string WJ_GetJoinURL(int ipIndex)
{
	if (!g_ctx || ipIndex < 0 || ipIndex >= (int)g_lanAddresses.size())
		return {};

	char buf[128];
	snprintf(buf, sizeof(buf), "http://%s:%d/?t=%s",
		g_lanAddresses[ipIndex].c_str(), g_listenPort, g_sessionToken.c_str());
	return buf;
}

CCMD(webjoin_start)
{
	int port = argv.argc() > 1 ? atoi(argv[1]) : 0;
	if (WJ_StartServer(port))
	{
		int n = WJ_GetJoinURLCount();
		if (n == 0)
			Printf("webjoin: started, but no LAN IPv4 address was found to show a join URL for.\n");
		for (int i = 0; i < n; i++)
			Printf("webjoin: join URL: %s\n", WJ_GetJoinURL(i).c_str());
	}
}

CCMD(webjoin_stop)
{
	WJ_StopServer();
	Printf("webjoin: stopped\n");
}