/*
 * ws.c - WebSocket protocol implementation.
 *
 * Handles the opening handshake, frame parsing/building, control frames
 * (ping/pong/close), and data frame echoing.
 */

#include "ws.h"
#include "base64.h"
#include "log.h"
#include "sha1.h"

#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>

// Magic GUID concatenated with the client key (RFC 6455)
#define WS_MAGIC_GUID "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

struct ws_client *ws_client_create(int fd)
{
	struct ws_client *c;

	c = malloc(sizeof(*c));
	if (!c)
		return NULL;

	c->fd = fd;
	c->session = NULL;
	c->buf = malloc(WSD_READ_BUF_SIZE);
	if (!c->buf) {
		free(c);
		return NULL;
	}

	return c;
}

void ws_client_destroy(struct ws_client *client)
{
	if (!client)
		return;

	if (client->fd >= 0)
		close(client->fd);

	free(client->buf);
	free(client);
}

struct ws_session *ws_session_create(struct ws_client *client)
{
	struct ws_session *s;

	s = malloc(sizeof(*s));
	if (!s)
		return NULL;

	s->state  = WS_STATE_OPEN;

	return s;
}

void ws_session_destroy(struct ws_session *session)
{
	if (!session)
		return;

	free(session->frag_buf);
	free(session);
}

void ws_frame_free(struct ws_frame *frame)
{
	if (!frame)
		return;

	free(frame->payload);
	frame->payload = NULL;
	frame->payload_len = 0;
}

static int ws_send_all(int fd, const void *data, size_t len)
{
	const uint8_t *p = data;
	ssize_t n;

	while (len > 0) {
		n = write(fd, p, len);
		if (n < 0) {
			if (errno == EINTR)
				continue;
			return -1;
		}
		p   += n;
		len -= (size_t)n;
	}

	return 0;
}

int ws_handshake(struct ws_client *client,
		 const struct http_request *req)
{
	const char *key;
	const char *upgrade;
	const char *version;
	char *concat = NULL;
	char response[512] = {0};
	uint8_t digest[SHA1_DIGEST_SIZE];
	char accept_b64[64];
	size_t concat_len;
	int resp_len;

	// Validate required headers
	upgrade = http_get_header(req, "Upgrade");
	if (!upgrade || strcasecmp(upgrade, "websocket") != 0) {
		LOG_ERR("Missing or invalid Upgrade header");
		return -1;
	}

	key = http_get_header(req, "Sec-WebSocket-Key");
	if (!key) {
		LOG_ERR("Missing Sec-WebSocket-Key header");
		return -1;
	}

	version = http_get_header(req, "Sec-WebSocket-Version");
	if (!version || strcmp(version, "13") != 0) {
		LOG_ERR("Unsupported WebSocket version: %s",
			version ? version : "(none)");
		return -1;
	}

	// Compute Sec-WebSocket-Accept = Base64(SHA-1(key + GUID))
	concat_len = strlen(key) + strlen(WS_MAGIC_GUID) + 1;
	concat = malloc(concat_len);
	if (!concat)
		return -1;

	snprintf(concat, concat_len, "%s%s", key, WS_MAGIC_GUID);
	sha1(concat, strlen(concat), digest);
	base64_encode(digest, SHA1_DIGEST_SIZE, accept_b64, sizeof(accept_b64));
	free(concat);

	// Send back the HTTP 101 Switching Protocols response
	resp_len = snprintf(response, 512,
			    "HTTP/1.1 101 Switching Protocols\r\n"
			    "Upgrade: websocket\r\n"
			    "Connection: Upgrade\r\n"
			    "Sec-WebSocket-Accept: %s\r\n"
			    "\r\n",
			    accept_b64);

	if (resp_len < 0 || resp_len >= 512)
		return -1;

	if (ws_send_all(client->fd, response, (size_t)resp_len) < 0) {
		LOG_ERR("Failed to send handshake response: %s",
			strerror(errno));
		return -1;
	}

	// Create the WebSocket session
	client->session = ws_session_create(client);
	if (!client->session)
		return -1;

	LOG_INFO("WebSocket connection established (fd=%d)", client->fd);

	return 0;
}

int ws_send_frame(struct ws_client *client, uint8_t opcode,
		  const void *payload, size_t len)
{
	uint8_t *frame_buf;
	size_t hdr_len = 2;
	size_t total;
	int fd = client->fd;

	// Compute header size
	if (len <= 125) {
		hdr_len = 2;
	} else if (len <= 0xFFFF) {
		hdr_len = 4;
	} else {
		hdr_len = 10;
	}

	total = hdr_len + len;
	frame_buf = malloc(total);
	if (!frame_buf)
		return -1;

	frame_buf[0] = 0x80 | (opcode & 0x0F); // FIN + opcode

	if (len <= 125) {
		frame_buf[1] = (uint8_t)len;
	} else if (len <= 0xFFFF) {
		frame_buf[1] = 126;
		frame_buf[2] = (uint8_t)(len >> 8);
		frame_buf[3] = (uint8_t)(len);
	} else {
		frame_buf[1] = 127;
		frame_buf[2] = (uint8_t)((uint64_t)len >> 56);
		frame_buf[3] = (uint8_t)((uint64_t)len >> 48);
		frame_buf[4] = (uint8_t)((uint64_t)len >> 40);
		frame_buf[5] = (uint8_t)((uint64_t)len >> 32);
		frame_buf[6] = (uint8_t)((uint64_t)len >> 24);
		frame_buf[7] = (uint8_t)((uint64_t)len >> 16);
		frame_buf[8] = (uint8_t)((uint64_t)len >> 8);
		frame_buf[9] = (uint8_t)((uint64_t)len);
	}

	if (len > 0)
		memcpy(frame_buf + hdr_len, payload, len);

	int ret = ws_send_all(fd, frame_buf, total);

	free(frame_buf);
	return ret;
}

int ws_send_close(struct ws_client *client, uint16_t code)
{
	uint8_t payload[2];

	payload[0] = (uint8_t)(code >> 8);
	payload[1] = (uint8_t)(code);

	return ws_send_frame(client, WS_OP_CLOSE, payload, sizeof(payload));
}

// Parses a frame from the buffer. Returns bytes consumed, 0 if incomplete, or -1 on error.
static ssize_t ws_parse_frame(const uint8_t *buf, size_t avail,
			      struct ws_frame *frame)
{
	size_t offset = 2; // minimum header size
	uint64_t payload_len;
	size_t i;

	if (avail < 2)
		return 0;

	memset(frame, 0, sizeof(*frame));

	frame->fin    = (buf[0] & 0x80) != 0;
	frame->opcode = buf[0] & 0x0F;
	frame->masked = (buf[1] & 0x80) != 0;

	payload_len = buf[1] & 0x7F;

	if (payload_len == 126) {
		if (avail < 4)
			return 0;
		payload_len = ((uint64_t)buf[2] << 8) | buf[3];
		offset = 4;
	} else if (payload_len == 127) {
		if (avail < 10)
			return 0;
		payload_len = ((uint64_t)buf[2] << 56) |
			      ((uint64_t)buf[3] << 48) |
			      ((uint64_t)buf[4] << 40) |
			      ((uint64_t)buf[5] << 32) |
			      ((uint64_t)buf[6] << 24) |
			      ((uint64_t)buf[7] << 16) |
			      ((uint64_t)buf[8] << 8)  |
			      ((uint64_t)buf[9]);
		offset = 10;
	}

	if (payload_len > WSD_MAX_FRAME_SIZE) {
		LOG_ERR("Frame payload too large: %llu bytes",
			(unsigned long long)payload_len);
		return -1;
	}

	// Client-to-server frames MUST be masked
	if (frame->masked) {
		if (avail < offset + 4)
			return 0;
		memcpy(frame->mask_key, buf + offset, 4);
		offset += 4;
	}

	// Wait until the full payload is buffered
	if (avail < offset + payload_len)
		return 0;

	frame->payload_len = payload_len;

	// Allocate and copy the payload so the client buffer can be reused
	if (payload_len > 0) {
		frame->payload = malloc((size_t)payload_len);
		if (!frame->payload)
			return -1;
		memcpy(frame->payload, buf + offset, (size_t)payload_len);

		// Unmask the payload
		if (frame->masked) {
			for (i = 0; i < payload_len; i++)
				frame->payload[i] ^= frame->mask_key[i % 4];
		}
	}

	return (ssize_t)(offset + payload_len);
}

static int ws_handle_frame(struct ws_client *client, struct ws_frame *frame)
{
	struct ws_session *session = client->session;
	int fd = client->fd;

	// Control frames (>= 0x8) can be interleaved and are NOT fragmented
	if (frame->opcode >= 0x8) {
		switch (frame->opcode) {
		case WS_OP_PING:
			LOG_DBG("PING received (fd=%d)", fd);
			return ws_send_frame(client, WS_OP_PONG,
					     frame->payload, frame->payload_len);
		case WS_OP_PONG:
			LOG_DBG("PONG received (fd=%d)", fd);
			return 0;
		case WS_OP_CLOSE:
			LOG_INFO("Close frame received (fd=%d)", fd);
			if (session->state == WS_STATE_OPEN) {
				uint16_t code = WS_CLOSE_NORMAL;
				if (frame->payload_len >= 2) {
					code = ((uint16_t)frame->payload[0] << 8) |
					       frame->payload[1];
				}
				ws_send_close(client, code);
				session->state = WS_STATE_CLOSED;
			} else {
				session->state = WS_STATE_CLOSED;
			}
			return -1;
		default:
			LOG_ERR("Unknown control opcode 0x%x (fd=%d)",
				frame->opcode, fd);
			ws_send_close(client, WS_CLOSE_PROTOCOL_ERR);
			session->state = WS_STATE_CLOSED;
			return -1;
		}
	}

	// Data frames (0x0, 0x1, 0x2)
	if (frame->opcode == WS_OP_TEXT || frame->opcode == WS_OP_BINARY) {
		if (session->frag_buf) {
			LOG_ERR("Received new data frame while fragmented message is pending");
			ws_send_close(client, WS_CLOSE_PROTOCOL_ERR);
			return -1;
		}

		if (frame->fin) {
			// Unfragmented message, echo directly
			LOG_DBG("Echo %s frame (%llu bytes, fd=%d)",
				frame->opcode == WS_OP_TEXT ? "text" : "binary",
				(unsigned long long)frame->payload_len, fd);
			return ws_send_frame(client, frame->opcode,
					     frame->payload, frame->payload_len);
		} else {
			// Start of a fragmented message
			session->frag_len = frame->payload_len;
			session->frag_opcode = frame->opcode;
			if (session->frag_len > 0) {
				session->frag_buf = malloc(session->frag_len);
				if (!session->frag_buf)
					return -1;
				memcpy(session->frag_buf, frame->payload, session->frag_len);
			} else {
				session->frag_buf = NULL;
			}
			return 0;
		}
	} else if (frame->opcode == WS_OP_CONTINUATION) {
		// Append payload to the fragmentation buffer
		if (frame->payload_len > 0) {
			size_t new_len = session->frag_len + frame->payload_len;
			if (new_len > WSD_MAX_FRAME_SIZE) {
				LOG_ERR("Fragmented message exceeds max size");
				ws_send_close(client, WS_CLOSE_TOO_BIG);
				return -1;
			}

			uint8_t *tmp = realloc(session->frag_buf, new_len);
			if (!tmp) {
				LOG_ERR("Failed to reallocate fragmentation buffer");
				ws_send_close(client, WS_CLOSE_ABNORMAL);
				return -1;
			}

			session->frag_buf = tmp;
			memcpy(session->frag_buf + session->frag_len,
			       frame->payload, frame->payload_len);
			session->frag_len = new_len;
		}

		if (frame->fin) {
			// End of fragmented message, echo back the fully assembled buffer
			// TODO: Implement actually useful logic here
			LOG_DBG("Echo fragmented frame (%zu bytes assembled, fd=%d)",
				session->frag_len, fd);

			int ret = ws_send_frame(client, session->frag_opcode,
						session->frag_buf, session->frag_len);

			// Clean up fragmentation state
			free(session->frag_buf);
			session->frag_buf = NULL;
			session->frag_len = 0;

			return ret;
		}

		return 0;
	}

	LOG_ERR("Unknown opcode 0x%x (fd=%d)", frame->opcode, fd);
	ws_send_close(client, WS_CLOSE_PROTOCOL_ERR);
	return -1;
}

// Data ingestion
int ws_on_data(struct ws_client *client, size_t len)
{
	struct ws_frame frame;
	ssize_t consumed;
	int ret;

	memset(&frame, 0, sizeof(frame));
	consumed = ws_parse_frame(client->buf, len, &frame);

	if (consumed < 0) {
		ws_frame_free(&frame);
		return -1; // protocol error
	}
	if (consumed == 0)
		return 0; // incomplete frame

	ret = ws_handle_frame(client, &frame);
	ws_frame_free(&frame);

	return ret;
}
