#include "seobeo/seobeo.h"
#include <time.h>

#define SEOBEO_WS_GUID "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

static void Seobeo_WebSocket_Parse_Url(const char *url, char **p_host, char **p_port, char **p_path, boolean *p_use_tls, Dowa_Arena *p_arena)
{
  if (!url)
    return;

  const char *start = url;
  *p_use_tls = FALSE;

  if (strncmp(url, "wss://", 6) == 0)
  {
    *p_use_tls = TRUE;
    start = url + 6;
  }
  else if (strncmp(url, "ws://", 5) == 0)
  {
    *p_use_tls = FALSE;
    start = url + 5;
  }

  const char *slash = strchr(start, '/');
  const char *colon = strchr(start, ':');

  if (colon && (!slash || colon < slash))
  {
    size_t host_len = colon - start;
    *p_host = Dowa_Arena_Allocate(p_arena, host_len + 1);
    memcpy(*p_host, start, host_len);
    (*p_host)[host_len] = '\0';

    const char *port_start = colon + 1;
    size_t port_len = slash ? (slash - port_start) : strlen(port_start);
    *p_port = Dowa_Arena_Allocate(p_arena, port_len + 1);
    memcpy(*p_port, port_start, port_len);
    (*p_port)[port_len] = '\0';
  }
  else
  {
    size_t host_len = slash ? (slash - start) : strlen(start);
    *p_host = Dowa_Arena_Allocate(p_arena, host_len + 1);
    memcpy(*p_host, start, host_len);
    (*p_host)[host_len] = '\0';

    *p_port = Dowa_Arena_Allocate(p_arena, 8);
    strcpy(*p_port, *p_use_tls ? "443" : "80");
  }

  if (slash)
  {
    size_t path_len = strlen(slash);
    *p_path = Dowa_Arena_Allocate(p_arena, path_len + 1);
    strcpy(*p_path, slash);
  }
  else
  {
    *p_path = Dowa_Arena_Allocate(p_arena, 2);
    strcpy(*p_path, "/");
  }
}

static void Seobeo_WebSocket_Generate_Key(char *key_out, size_t key_size)
{
  static const char base64_chars[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

  srand((unsigned int)time(NULL) ^ (unsigned int)getpid());

  uint8 random_bytes[16];
  for (int i = 0; i < 16; i++)
    random_bytes[i] = (uint8)(rand() % 256);

  int out_idx = 0;
  for (int i = 0; i < 16; i += 3)
  {
    uint32 triple = (random_bytes[i] << 16) |
                    (i + 1 < 16 ? random_bytes[i + 1] << 8 : 0) |
                    (i + 2 < 16 ? random_bytes[i + 2] : 0);

    key_out[out_idx++] = base64_chars[(triple >> 18) & 0x3F];
    key_out[out_idx++] = base64_chars[(triple >> 12) & 0x3F];
    key_out[out_idx++] = base64_chars[(triple >> 6) & 0x3F];
    key_out[out_idx++] = base64_chars[triple & 0x3F];
  }

  key_out[out_idx] = '\0';
}

Seobeo_WebSocket *Seobeo_WebSocket_Connect_With_Headers(const char *url, const char *headers)
{
  Seobeo_WebSocket *p_ws = malloc(sizeof(Seobeo_WebSocket));
  if (!p_ws)
    return NULL;

  memset(p_ws, 0, sizeof(Seobeo_WebSocket));

  p_ws->p_arena = Dowa_Arena_Create(1024 * 1024);
  if (!p_ws->p_arena)
  {
    free(p_ws);
    return NULL;
  }

  size_t url_len = strlen(url);
  p_ws->url = Dowa_Arena_Allocate(p_ws->p_arena, url_len + 1);
  strcpy(p_ws->url, url);

  Seobeo_WebSocket_Parse_Url(url, &p_ws->host, &p_ws->port, &p_ws->path, &p_ws->use_tls, p_ws->p_arena);

  p_ws->p_handle = Seobeo_Stream_Handle_Client_Create(p_ws->host, p_ws->port, p_ws->use_tls);
  if (!p_ws->p_handle || p_ws->p_handle->socket < 0)
  {
    Seobeo_Log(SEOBEO_ERROR, "Failed to create socket connection\n");
    if (p_ws->p_handle)
      Seobeo_Handle_Destroy(p_ws->p_handle);
    Dowa_Arena_Free(p_ws->p_arena);
    free(p_ws);
    return NULL;
  }

  char ws_key[32];
  Seobeo_WebSocket_Generate_Key(ws_key, sizeof(ws_key));

  // Build handshake with optional custom headers
  char handshake[4096];
  int handshake_len = snprintf(handshake, sizeof(handshake),
    "GET %s HTTP/1.1\r\n"
    "Host: %s\r\n"
    "Upgrade: websocket\r\n"
    "Connection: Upgrade\r\n"
    "Sec-WebSocket-Key: %s\r\n"
    "Sec-WebSocket-Version: 13\r\n",
    p_ws->path, p_ws->host, ws_key);

  if (headers && strlen(headers) > 0)
  {
    const char *line_start = headers;
    while (*line_start)
    {
      while (*line_start == ' ' || *line_start == '\t') line_start++;
      if (*line_start == '\0') break;

      const char *line_end = line_start;
      while (*line_end && *line_end != '\n') line_end++;

      size_t line_len = line_end - line_start;
      if (line_len > 0 && memchr(line_start, ':', line_len) != NULL)
      {
        while (line_len > 0 && (line_start[line_len-1] == '\r' ||
               line_start[line_len-1] == ' ' || line_start[line_len-1] == '\t'))
          line_len--;

        if (line_len > 0 && handshake_len + line_len + 4 < sizeof(handshake))
        {
          memcpy(handshake + handshake_len, line_start, line_len);
          handshake_len += line_len;
          handshake[handshake_len++] = '\r';
          handshake[handshake_len++] = '\n';
        }
      }

      line_start = (*line_end == '\n') ? line_end + 1 : line_end;
    }
  }

  handshake[handshake_len++] = '\r';
  handshake[handshake_len++] = '\n';
  handshake[handshake_len] = '\0';

  Seobeo_Handle_Queue(p_ws->p_handle, (uint8*)handshake, (uint32)handshake_len);
  if (Seobeo_Handle_Flush(p_ws->p_handle) < 0)
  {
    Seobeo_Log(SEOBEO_ERROR, "Failed to send WebSocket handshake\n");
    Seobeo_Handle_Destroy(p_ws->p_handle);
    Dowa_Arena_Free(p_ws->p_arena);
    free(p_ws);
    return NULL;
  }

  while (1)
  {
    int r = Seobeo_Handle_Read(p_ws->p_handle);
    if (r < 0)
    {
      Seobeo_Log(SEOBEO_ERROR, "Failed to read handshake response\n");
      Seobeo_Handle_Destroy(p_ws->p_handle);
      Dowa_Arena_Free(p_ws->p_arena);
      free(p_ws);
      return NULL;
    }

    if (p_ws->p_handle->read_buffer_len >= 4 &&
        strstr((char*)p_ws->p_handle->read_buffer, "\r\n\r\n") != NULL)
      break;

    if (r == 0)
      continue;
  }

  char *response = (char*)p_ws->p_handle->read_buffer;
  if (strstr(response, "HTTP/1.1 101") == NULL)
  {
    Seobeo_Log(SEOBEO_ERROR, "WebSocket handshake failed: %s\n", response);
    Seobeo_Handle_Destroy(p_ws->p_handle);
    Dowa_Arena_Free(p_ws->p_arena);
    free(p_ws);
    return NULL;
  }

  char *end_of_headers = strstr(response, "\r\n\r\n");
  if (end_of_headers)
  {
    size_t header_len = end_of_headers - response + 4;
    Seobeo_Handle_Consume(p_ws->p_handle, (uint32)header_len);
  }

  p_ws->state = SEOBEO_WS_STATE_OPEN;
  p_ws->fragment_capacity = 4096;
  p_ws->fragment_buffer = malloc(p_ws->fragment_capacity);

  Seobeo_Log(SEOBEO_INFO, "WebSocket connected to %s\n", url);

  return p_ws;
}

Seobeo_WebSocket *Seobeo_WebSocket_Connect(const char *url)
{
  return Seobeo_WebSocket_Connect_With_Headers(url, NULL);
}

static int32 Seobeo_WebSocket_Send_Frame(Seobeo_WebSocket *p_ws, Seobeo_WebSocket_Opcode opcode,
    const uint8 *payload, size_t payload_length, boolean fin)
{
  if (!p_ws || p_ws->state != SEOBEO_WS_STATE_OPEN)
    return -1;

  uint8 frame[14]; 
  size_t frame_len = 0;

  // Big endian
  frame[0] = (fin ? 0x80 : 0x00) | (opcode & 0x0F);
  frame_len++;

  uint8 mask_key[4];
  for (int i = 0; i < 4; i++)
    mask_key[i] = (uint8)(rand() % 256);

  // within 1 byte
  if (payload_length < 126)
  {
    frame[1] = 0x80 | (uint8)payload_length;
    frame_len++;
  }
  // from now frame 1 is thrown away just keeping it 1 
  // within 4 bytes
  else if (payload_length < MAX_INT_16)
  {
    frame[1] = 0x80 | 126; 
    frame[2] = (uint8)((payload_length >> 8) & 0xFF); 
    frame[3] = (uint8)(payload_length & 0xFF); 
    frame_len += 3;
  }
  // within 8 bytes
  else
  {
    frame[1] = 0x80 | 127;
    for (int i = 0; i < 8; i++)
      frame[2 + i] = (uint8)((payload_length >> (56 - i * 8)) & 0xFF);
    frame_len += 9;
  }

  memcpy(frame + frame_len, mask_key, 4);
  frame_len += 4;

  Seobeo_Handle_Queue(p_ws->p_handle, frame, (uint32)frame_len);

  if (payload_length > 0)
  {
    uint8 *masked_payload = malloc(payload_length);
    memcpy(masked_payload, payload, payload_length);
    Seobeo_WebSocket_Mask_Data(masked_payload, payload_length, mask_key);
    Seobeo_Handle_Queue(p_ws->p_handle, masked_payload, (uint32)payload_length);
    free(masked_payload);
  }

  return Seobeo_Handle_Flush(p_ws->p_handle);
}

static int32 Seobeo_WebSocket_Send_Fragmented(Seobeo_WebSocket *p_ws, Seobeo_WebSocket_Opcode opcode, 
    const uint8 *payload, size_t total_length)
{
  if (!payload || total_length == 0)
    return -1;

  if (total_length <= MAX_FRAGMENT_SIZE)
    return Seobeo_WebSocket_Send_Frame(p_ws, opcode, payload, total_length, TRUE);

  size_t sent = 0;
  int32 result;

  result = Seobeo_WebSocket_Send_Frame(p_ws, opcode, payload, MAX_FRAGMENT_SIZE, FALSE);
  if (result < 0)
    return result;

  sent += MAX_FRAGMENT_SIZE;

  while (sent + MAX_FRAGMENT_SIZE < total_length)
  {
    result = Seobeo_WebSocket_Send_Frame(p_ws, SEOBEO_WS_OPCODE_CONTINUATION, payload + sent, MAX_FRAGMENT_SIZE, FALSE);
    if (result < 0)
      return result;
    sent += MAX_FRAGMENT_SIZE;
  }

  size_t remaining = total_length - sent;
  return Seobeo_WebSocket_Send_Frame(p_ws, SEOBEO_WS_OPCODE_CONTINUATION, payload + sent, remaining, TRUE);
}

int32 Seobeo_WebSocket_Send_Text(Seobeo_WebSocket *p_ws, const char *text)
{
  if (!text)
    return -1;

  return Seobeo_WebSocket_Send_Fragmented(p_ws, SEOBEO_WS_OPCODE_TEXT, (const uint8*)text, strlen(text));
}

int32 Seobeo_WebSocket_Send_Binary(Seobeo_WebSocket *p_ws, const uint8 *data, size_t length)
{
  if (!data)
    return -1;

  return Seobeo_WebSocket_Send_Fragmented(p_ws, SEOBEO_WS_OPCODE_BINARY, data, length);
}

int32 Seobeo_WebSocket_Send_Ping(Seobeo_WebSocket *p_ws, const char *payload)
{
  size_t len = payload ? strlen(payload) : 0;
  return Seobeo_WebSocket_Send_Frame(p_ws, SEOBEO_WS_OPCODE_PING, (const uint8*)payload, len, TRUE);
}

int32 Seobeo_WebSocket_Send_Pong(Seobeo_WebSocket *p_ws, const char *payload)
{
  size_t len = payload ? strlen(payload) : 0;
  return Seobeo_WebSocket_Send_Frame(p_ws, SEOBEO_WS_OPCODE_PONG, (const uint8*)payload, len, TRUE);
}

Seobeo_WebSocket_Message *Seobeo_WebSocket_Receive(Seobeo_WebSocket *p_ws)
{
  if (!p_ws || p_ws->state == SEOBEO_WS_STATE_CLOSED)
    return NULL;

  int r = Seobeo_Handle_Read(p_ws->p_handle);
  if (r < 0)
  {
    Seobeo_Log(SEOBEO_ERROR, "WebSocket read error\n");
    p_ws->state = SEOBEO_WS_STATE_CLOSED;
    return NULL;
  }

  if (r == -2)
  {
    Seobeo_Log(SEOBEO_INFO, "WebSocket connection closed\n");
    p_ws->state = SEOBEO_WS_STATE_CLOSED;
    return NULL;
  }

  if (p_ws->p_handle->read_buffer_len < 2)
    return NULL;

  uint8 *buf = p_ws->p_handle->read_buffer;

  uint8 byte1 = buf[0];
  uint8 byte2 = buf[1];

  boolean fin = (byte1 & 0x80) != 0;
  Seobeo_WebSocket_Opcode opcode = (Seobeo_WebSocket_Opcode)(byte1 & 0x0F);
  boolean masked = (byte2 & 0x80) != 0;
  uint64 payload_len = byte2 & 0x7F;

  size_t header_len = 2;

  if (payload_len == 126)
  {
    if (p_ws->p_handle->read_buffer_len < 4)
      return NULL;
    payload_len = (buf[2] << 8) | buf[3];
    header_len += 2;
  }
  else if (payload_len == 127)
  {
    if (p_ws->p_handle->read_buffer_len < 10)
      return NULL;
    payload_len = 0;
    for (int i = 0; i < 8; i++)
      payload_len = (payload_len << 8) | buf[2 + i];
    header_len += 8;
  }

  uint8 mask_key[4] = {0};
  if (masked)
  {
    if (p_ws->p_handle->read_buffer_len < header_len + 4)
      return NULL;
    memcpy(mask_key, buf + header_len, 4);
    header_len += 4;
  }

  if (p_ws->p_handle->read_buffer_len < header_len + payload_len)
    return NULL;

  uint8 *payload = NULL;
  if (payload_len > 0)
  {
    payload = malloc(payload_len);
    memcpy(payload, buf + header_len, payload_len);

    if (masked)
      Seobeo_WebSocket_Mask_Data(payload, payload_len, mask_key);
  }

  Seobeo_Handle_Consume(p_ws->p_handle, (uint32)(header_len + payload_len));

  if (opcode == SEOBEO_WS_OPCODE_PING)
  {
    Seobeo_WebSocket_Send_Pong(p_ws, (char*)payload);
    if (payload)
      free(payload);
    return NULL;
  }

  if (opcode == SEOBEO_WS_OPCODE_PONG)
  {
    if (payload)
      free(payload);
    return NULL;
  }

  if (opcode == SEOBEO_WS_OPCODE_CLOSE)
  {
    uint16 close_code = 1000;
    if (payload_len >= 2)
      close_code = (payload[0] << 8) | payload[1];

    Seobeo_Log(SEOBEO_INFO, "WebSocket close received with code %d\n", close_code);
    p_ws->state = SEOBEO_WS_STATE_CLOSING;

    Seobeo_WebSocket_Send_Frame(p_ws, SEOBEO_WS_OPCODE_CLOSE, payload, payload_len, TRUE);
    p_ws->state = SEOBEO_WS_STATE_CLOSED;

    if (payload)
      free(payload);
    return NULL;
  }

  if (opcode == SEOBEO_WS_OPCODE_CONTINUATION)
  {
    if (p_ws->fragment_length + payload_len > p_ws->fragment_capacity)
    {
      p_ws->fragment_capacity = (p_ws->fragment_length + payload_len) * 2;
      p_ws->fragment_buffer = realloc(p_ws->fragment_buffer, p_ws->fragment_capacity);
    }

    if (payload_len > 0)
    {
      memcpy(p_ws->fragment_buffer + p_ws->fragment_length, payload, payload_len);
      p_ws->fragment_length += payload_len;
    }

    if (payload)
      free(payload);

    if (!fin)
      return NULL;

    Seobeo_WebSocket_Message *p_msg = malloc(sizeof(Seobeo_WebSocket_Message));
    p_msg->opcode = p_ws->fragment_opcode;
    p_msg->data = malloc(p_ws->fragment_length);
    memcpy(p_msg->data, p_ws->fragment_buffer, p_ws->fragment_length);
    p_msg->length = p_ws->fragment_length;
    p_msg->is_final = TRUE;

    p_ws->fragment_length = 0;

    return p_msg;
  }

  if (!fin)
  {
    p_ws->fragment_opcode = opcode;
    p_ws->fragment_length = 0;

    if (payload_len > 0)
    {
      if (payload_len > p_ws->fragment_capacity)
      {
        p_ws->fragment_capacity = payload_len * 2;
        p_ws->fragment_buffer = realloc(p_ws->fragment_buffer, p_ws->fragment_capacity);
      }

      memcpy(p_ws->fragment_buffer, payload, payload_len);
      p_ws->fragment_length = payload_len;
    }

    if (payload)
      free(payload);

    return NULL;
  }

  Seobeo_WebSocket_Message *p_msg = malloc(sizeof(Seobeo_WebSocket_Message));
  p_msg->opcode = opcode;
  p_msg->data = payload;
  p_msg->length = payload_len;
  p_msg->is_final = fin;

  return p_msg;
}

int32 Seobeo_WebSocket_Close(Seobeo_WebSocket *p_ws, uint16 code, const char *reason)
{
  if (!p_ws || p_ws->state == SEOBEO_WS_STATE_CLOSED)
    return -1;

  p_ws->state = SEOBEO_WS_STATE_CLOSING;

  size_t reason_len = reason ? strlen(reason) : 0;
  size_t payload_len = 2 + reason_len;
  uint8 *payload = malloc(payload_len);

  payload[0] = (uint8)((code >> 8) & 0xFF);
  payload[1] = (uint8)(code & 0xFF);

  if (reason_len > 0)
    memcpy(payload + 2, reason, reason_len);

  int32 result = Seobeo_WebSocket_Send_Frame(p_ws, SEOBEO_WS_OPCODE_CLOSE, payload, payload_len, TRUE);

  free(payload);

  p_ws->state = SEOBEO_WS_STATE_CLOSED;

  return result;
}

void Seobeo_WebSocket_Destroy(Seobeo_WebSocket *p_ws)
{
  if (!p_ws)
    return;

  if (p_ws->state == SEOBEO_WS_STATE_OPEN)
    Seobeo_WebSocket_Close(p_ws, 1000, "Normal closure");

  if (p_ws->p_handle)
    Seobeo_Handle_Destroy(p_ws->p_handle);

  if (p_ws->fragment_buffer)
    free(p_ws->fragment_buffer);

  if (p_ws->p_arena)
    Dowa_Arena_Free(p_ws->p_arena);

  free(p_ws);
}