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

#ifndef SEOBEO_NO_SSL
#include <openssl/sha.h>
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/buffer.h>
#endif

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

static Seobeo_WebSocket_Server_Route *g_ws_routes = NULL;
static Seobeo_WebSocket_Server_Connection *g_ws_connections = NULL;

void Seobeo_WebSocket_Server_Init()
{
  Dowa_Array_Reserve(g_ws_routes, 10);
}

void Seobeo_WebSocket_Server_Register(const char *path, Seobeo_WebSocket_Server_Handler handler, void *p_user_data)
{
  Seobeo_WebSocket_Server_Route route = {0};
  route.path = strdup(path);
  route.handler = handler;
  route.p_user_data = p_user_data;

  Dowa_Array_Push(g_ws_routes, route);
}

static void Seobeo_WebSocket_Server_Compute_Accept_Key(const char *client_key, char *accept_key, size_t accept_key_size)
{
#ifdef SEOBEO_NO_SSL
  snprintf(accept_key, accept_key_size, "dGhlIHNhbXBsZSBub25jZQ==");
  (void)client_key;
#else
  char concatenated[256];
  snprintf(concatenated, sizeof(concatenated), "%s%s", client_key, SEOBEO_WS_GUID);

  unsigned char hash[SHA_DIGEST_LENGTH];
  SHA1((unsigned char*)concatenated, strlen(concatenated), hash);

  BIO *b64 = BIO_new(BIO_f_base64());
  BIO *bio = BIO_new(BIO_s_mem());
  bio = BIO_push(b64, bio);
  BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
  BIO_write(bio, hash, SHA_DIGEST_LENGTH);
  BIO_flush(bio);

  BUF_MEM *buffer_ptr;
  BIO_get_mem_ptr(bio, &buffer_ptr);

  size_t copy_len = buffer_ptr->length < accept_key_size - 1 ? buffer_ptr->length : accept_key_size - 1;
  memcpy(accept_key, buffer_ptr->data, copy_len);
  accept_key[copy_len] = '\0';

  BIO_free_all(bio);
#endif
}

boolean Seobeo_WebSocket_Server_Handle_Upgrade(Seobeo_Handle *p_handle, Seobeo_Request_Entry *p_req_map, const char *path)
{
  void *p_upgrade_kv = Dowa_HashMap_Get_Ptr(p_req_map, "Upgrade");
  if (!p_upgrade_kv)
    return FALSE;

  const char *upgrade_value = ((Seobeo_Request_Entry*)p_upgrade_kv)->value;
  if (strcasecmp(upgrade_value, "websocket") != 0)
    return FALSE;

  void *p_connection_kv = Dowa_HashMap_Get_Ptr(p_req_map, "Connection");
  if (!p_connection_kv)
    return FALSE;

  const char *connection_value = ((Seobeo_Request_Entry*)p_connection_kv)->value;
  if (!strcasestr(connection_value, "upgrade"))
    return FALSE;

  void *p_key_kv = Dowa_HashMap_Get_Ptr(p_req_map, "Sec-WebSocket-Key");
  if (!p_key_kv)
    return FALSE;

  const char *client_key = ((Seobeo_Request_Entry*)p_key_kv)->value;

  char accept_key[64];
  Seobeo_WebSocket_Server_Compute_Accept_Key(client_key, accept_key, sizeof(accept_key));

  char response[512];
  int response_len = snprintf(response, sizeof(response),
    "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_key);

  Seobeo_Handle_Queue(p_handle, (uint8*)response, (uint32)response_len);
  if (Seobeo_Handle_Flush(p_handle) < 0)
    return FALSE;

  Seobeo_WebSocket_Server_Connection *p_conn = malloc(sizeof(Seobeo_WebSocket_Server_Connection));
  memset(p_conn, 0, sizeof(Seobeo_WebSocket_Server_Connection));

  p_conn->p_handle = p_handle;
  p_conn->is_active = TRUE;
  p_conn->fragment_capacity = 4096;
  p_conn->fragment_buffer = malloc(p_conn->fragment_capacity);

  char client_id[64];
  snprintf(client_id, sizeof(client_id), "client_%p", (void*)p_handle);
  p_conn->client_id = strdup(client_id);

  p_conn->next = g_ws_connections;
  g_ws_connections = p_conn;

  Seobeo_Log(SEOBEO_INFO, "WebSocket upgraded on path: %s\n", path);

  Seobeo_WebSocket_Server_Handle_Connection(p_conn);

  return TRUE;
}

// Seobeo_WebSocket_Unmask_Data removed - using Seobeo_WebSocket_Mask_Data from s_websocket_common.c (XOR is symmetric)

static int32 Seobeo_WebSocket_Server_Send_Frame(Seobeo_WebSocket_Server_Connection *p_conn, Seobeo_WebSocket_Opcode opcode, const uint8 *payload, size_t payload_length, boolean fin)
{
  if (!p_conn || !p_conn->is_active)
    return -1;

  uint8 frame[14];
  size_t frame_len = 0;

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

  // within 1 byte, so max chunk would be 125 bytes 
  if (payload_length < 126)
  {
    frame[1] = (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] = 126;
    frame[2] = (uint8)((payload_length >> 8) & 0xFF);
    frame[3] = (uint8)(payload_length & 0xFF);
    frame_len += 3;
  }
  else
  {
    frame[1] = 127;
    for (int i = 0; i < 8; i++)
      frame[2 + i] = (uint8)((payload_length >> (56 - i * 8)) & 0xFF);
    frame_len += 9;
  }

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

  if (payload_length > 0)
    Seobeo_Handle_Queue(p_conn->p_handle, payload, (uint32)payload_length);

  return Seobeo_Handle_Flush(p_conn->p_handle);
}

static int32 Seobeo_WebSocket_Server_Send_Fragmented(Seobeo_WebSocket_Server_Connection *p_conn, 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_Server_Send_Frame(p_conn, opcode, payload, total_length, TRUE);

  size_t sent = 0;
  int32 result;

  result = Seobeo_WebSocket_Server_Send_Frame(p_conn, 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_Server_Send_Frame(p_conn, 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_Server_Send_Frame(p_conn, SEOBEO_WS_OPCODE_CONTINUATION, payload + sent, remaining, TRUE);
}

int32 Seobeo_WebSocket_Server_Send_Text(Seobeo_WebSocket_Server_Connection *p_conn, const char *text)
{
  if (!text)
    return -1;

  return Seobeo_WebSocket_Server_Send_Fragmented(p_conn, SEOBEO_WS_OPCODE_TEXT, (const uint8*)text, strlen(text));
}

int32 Seobeo_WebSocket_Server_Send_Binary(Seobeo_WebSocket_Server_Connection *p_conn, const uint8 *data, size_t length)
{
  if (!data)
    return -1;

  return Seobeo_WebSocket_Server_Send_Fragmented(p_conn, SEOBEO_WS_OPCODE_BINARY, data, length);
}

static Seobeo_WebSocket_Message *Seobeo_WebSocket_Server_Receive_Frame(Seobeo_WebSocket_Server_Connection *p_conn)
{
  if (!p_conn || !p_conn->is_active)
    return NULL;

  int r = Seobeo_Handle_Read(p_conn->p_handle);
  if (r < 0)
  {
    Seobeo_Log(SEOBEO_ERROR, "WebSocket server read error\n");
    p_conn->is_active = FALSE;
    return NULL;
  }

  if (r == -2)
  {
    Seobeo_Log(SEOBEO_INFO, "WebSocket client disconnected\n");
    p_conn->is_active = FALSE;
    return NULL;
  }

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

  uint8 *buf = p_conn->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_conn->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_conn->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_conn->p_handle->read_buffer_len < header_len + 4)
      return NULL;
    memcpy(mask_key, buf + header_len, 4);
    header_len += 4;
  }

  if (p_conn->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_conn->p_handle, (uint32)(header_len + payload_len));

  if (opcode == SEOBEO_WS_OPCODE_PING)
  {
    Seobeo_WebSocket_Server_Send_Frame(p_conn, SEOBEO_WS_OPCODE_PONG, payload, payload_len, TRUE);
    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 from client with code %d\n", close_code);

    Seobeo_WebSocket_Server_Send_Frame(p_conn, SEOBEO_WS_OPCODE_CLOSE, payload, payload_len, TRUE);
    p_conn->is_active = FALSE;

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

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

    if (payload_len > 0)
    {
      memcpy(p_conn->fragment_buffer + p_conn->fragment_length, payload, payload_len);
      p_conn->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_conn->fragment_opcode;
    p_msg->data = malloc(p_conn->fragment_length);
    memcpy(p_msg->data, p_conn->fragment_buffer, p_conn->fragment_length);
    p_msg->length = p_conn->fragment_length;
    p_msg->is_final = TRUE;

    p_conn->fragment_length = 0;

    return p_msg;
  }

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

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

      memcpy(p_conn->fragment_buffer, payload, payload_len);
      p_conn->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;
}

void Seobeo_WebSocket_Server_Handle_Connection(Seobeo_WebSocket_Server_Connection *p_conn)
{
  if (!g_ws_routes)
    return;

  size_t route_count = Dowa_Array_Length(g_ws_routes);
  if (route_count == 0)
    return;

  Seobeo_WebSocket_Server_Handler handler = g_ws_routes[0].handler;
  void *p_user_data = g_ws_routes[0].p_user_data;

  while (p_conn->is_active)
  {
    Seobeo_WebSocket_Message *p_msg = Seobeo_WebSocket_Server_Receive_Frame(p_conn);
    if (p_msg)
    {
      if (handler)
        handler(p_conn, p_msg, p_user_data);

      Seobeo_WebSocket_Message_Destroy(p_msg);
    }

    usleep(1000);
  }

  Seobeo_WebSocket_Server_Connection_Destroy(p_conn);
}

void Seobeo_WebSocket_Server_Broadcast_Text(const char *text, Seobeo_WebSocket_Server_Connection *origin_p_conn)
{
  if (!text)
    return;

  Seobeo_WebSocket_Server_Connection *p_conn = g_ws_connections;
  while (p_conn)
  {
    if (p_conn->is_active && p_conn != origin_p_conn)
      Seobeo_WebSocket_Server_Send_Text(p_conn, text);

    p_conn = p_conn->next;
  }
}

void Seobeo_WebSocket_Server_Broadcast_Binary(const uint8 *data, size_t length)
{
  if (!data)
    return;

  Seobeo_WebSocket_Server_Connection *p_conn = g_ws_connections;
  while (p_conn)
  {
    if (p_conn->is_active)
      Seobeo_WebSocket_Server_Send_Binary(p_conn, data, length);

    p_conn = p_conn->next;
  }
}

void Seobeo_WebSocket_Server_Connection_Close(Seobeo_WebSocket_Server_Connection *p_conn, uint16 code, const char *reason)
{
  if (!p_conn || !p_conn->is_active)
    return;

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

  Seobeo_WebSocket_Server_Send_Frame(p_conn, SEOBEO_WS_OPCODE_CLOSE, payload, payload_len, TRUE);

  free(payload);

  p_conn->is_active = FALSE;
}

void Seobeo_WebSocket_Server_Connection_Destroy(Seobeo_WebSocket_Server_Connection *p_conn)
{
  if (!p_conn)
    return;

  if (p_conn->is_active)
    Seobeo_WebSocket_Server_Connection_Close(p_conn, 1000, "Server closing connection");

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

  if (p_conn->client_id)
    free(p_conn->client_id);

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

  Seobeo_WebSocket_Server_Connection **pp = &g_ws_connections;
  while (*pp)
  {
    if (*pp == p_conn)
    {
      *pp = p_conn->next;
      break;
    }
    pp = &(*pp)->next;
  }

  free(p_conn);
}