#include "seobeo/seobeo.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static const char *SEOBEO_SSE_HEADERS =
    "HTTP/1.1 200 OK\r\n"
    "Content-Type: text/event-stream\r\n"
    "Cache-Control: no-cache\r\n"
    "Connection: keep-alive\r\n"
    "X-Accel-Buffering: no\r\n"
    "\r\n";

static pthread_mutex_t g_sse_mutex = PTHREAD_MUTEX_INITIALIZER;
static Seobeo_SSE_Stream **g_sse_streams = NULL;

typedef struct {
  char *data;
  uint32 size;
} Seobeo_SSE_Pending_Frame;

#define SEOBEO_SSE_MAX_PENDING 32

static boolean sse_is_open_unlocked(const Seobeo_SSE_Stream *p_stream)
{
  return p_stream &&
      p_stream->started &&
      !p_stream->closed &&
      p_stream->p_handle &&
      p_stream->p_handle->socket >= 0 &&
      !atomic_load(&p_stream->p_handle->destroyed);
}

static void sse_clear_pending_unlocked(Seobeo_SSE_Stream *p_stream)
{
  Seobeo_SSE_Pending_Frame *p_frames = p_stream->p_pending_frames;
  size_t count = Dowa_Array_Length(p_frames);
  for (size_t i = p_stream->pending_offset; i < count; i++)
    Dowa_Free(p_frames[i].data);
  Dowa_Array_Free(p_frames);
  p_stream->p_pending_frames = NULL;
  p_stream->pending_offset = 0;
}

static void sse_release_unlocked(Seobeo_SSE_Stream *p_stream)
{
  if (!p_stream)
    return;
  if (atomic_fetch_sub(&p_stream->ref_count, 1) != 1)
    return;
  sse_clear_pending_unlocked(p_stream);
  Dowa_Free(p_stream->path);
  Dowa_Free(p_stream);
}

static int32 sse_drain_pending_unlocked(Seobeo_SSE_Stream *p_stream)
{
  if (!sse_is_open_unlocked(p_stream))
    return -1;
  int32 flush_result = Seobeo_Handle_Flush(p_stream->p_handle);
  if (flush_result != 0)
    return flush_result;

  Seobeo_SSE_Pending_Frame *p_frames = p_stream->p_pending_frames;
  size_t count = Dowa_Array_Length(p_frames);
  while (p_stream->pending_offset < count)
  {
    Seobeo_SSE_Pending_Frame *p_frame =
        &p_frames[p_stream->pending_offset];
    int32 queue_result = Seobeo_Handle_Queue(
        p_stream->p_handle,
        (const uint8 *)p_frame->data,
        p_frame->size);
    if (queue_result != 0)
      return queue_result;
    Dowa_Free(p_frame->data);
    p_stream->pending_offset++;
    flush_result = Seobeo_Handle_Flush(p_stream->p_handle);
    if (flush_result != 0)
      return flush_result;
  }

  Dowa_Array_Free(p_frames);
  p_stream->p_pending_frames = NULL;
  p_stream->pending_offset = 0;
  return 0;
}

static int32 sse_enqueue_frame_unlocked(
    Seobeo_SSE_Stream *p_stream,
    char *frame,
    uint32 frame_size)
{
  Seobeo_SSE_Pending_Frame *p_frames = p_stream->p_pending_frames;
  size_t pending = Dowa_Array_Length(p_frames) - p_stream->pending_offset;
  if (pending >= SEOBEO_SSE_MAX_PENDING)
    return SEOBEO_SSE_BACKPRESSURE;
  Seobeo_SSE_Pending_Frame pending_frame = {
    .data = frame,
    .size = frame_size,
  };
  Dowa_Array_Push(p_frames, pending_frame);
  p_stream->p_pending_frames = p_frames;
  return sse_drain_pending_unlocked(p_stream);
}

static boolean sse_field_is_safe(const char *value)
{
  return !value || (strchr(value, '\r') == NULL && strchr(value, '\n') == NULL);
}

static size_t sse_multiline_size(const char *prefix, const char *value)
{
  const char *line = value ? value : "";
  size_t prefix_length = strlen(prefix);
  size_t total = 0;

  while (TRUE)
  {
    const char *end = line;
    while (*end && *end != '\r' && *end != '\n')
      end++;
    size_t line_length = (size_t)(end - line);
    total += prefix_length + (line_length ? 1 : 0) + line_length + 1;
    if (*end == '\0')
      break;
    if (*end == '\r' && end[1] == '\n')
      end++;
    line = end + 1;
  }
  return total;
}

static char *sse_append_multiline(
    char *output,
    const char *prefix,
    const char *value)
{
  const char *line = value ? value : "";
  size_t prefix_length = strlen(prefix);

  while (TRUE)
  {
    const char *end = line;
    while (*end && *end != '\r' && *end != '\n')
      end++;
    size_t line_length = (size_t)(end - line);
    memcpy(output, prefix, prefix_length);
    output += prefix_length;
    if (line_length)
      *output++ = ' ';
    memcpy(output, line, line_length);
    output += line_length;
    *output++ = '\n';
    if (*end == '\0')
      break;
    if (*end == '\r' && end[1] == '\n')
      end++;
    line = end + 1;
  }
  return output;
}

static int32 sse_send_frame(
    Seobeo_SSE_Stream *p_stream,
    const char *frame,
    size_t frame_size)
{
  if (!sse_is_open_unlocked(p_stream) || !frame)
    return -1;
  if (frame_size > p_stream->p_handle->write_buffer_capacity)
    return SEOBEO_SSE_EVENT_TOO_LARGE;

  int32 flush_result = Seobeo_Handle_Flush(p_stream->p_handle);
  if (flush_result != 0)
    return flush_result;

  int32 queue_result = Seobeo_Handle_Queue(
      p_stream->p_handle,
      (const uint8 *)frame,
      (uint32)frame_size);
  if (queue_result != 0)
    return queue_result;
  return Seobeo_Handle_Flush(p_stream->p_handle);
}

int32 Seobeo_SSE_Start(
    Seobeo_SSE_Stream *p_stream,
    Seobeo_Handle *p_handle)
{
  if (!p_stream || !p_handle || atomic_load(&p_handle->destroyed))
    return -1;
  p_stream->p_handle = p_handle;
  p_stream->started = TRUE;
  p_stream->closed = FALSE;
  return sse_send_frame(
      p_stream,
      SEOBEO_SSE_HEADERS,
      strlen(SEOBEO_SSE_HEADERS));
}

static int32 sse_send_event_unlocked(
    Seobeo_SSE_Stream *p_stream,
    const Seobeo_SSE_Event *p_event)
{
  if (!sse_is_open_unlocked(p_stream))
    return -1;
  if (!p_event || !sse_field_is_safe(p_event->event) ||
      !sse_field_is_safe(p_event->id) ||
      p_event->retry_ms < SEOBEO_SSE_RETRY_NONE)
  {
    return SEOBEO_SSE_INVALID_EVENT;
  }

  size_t frame_size = 1 + sse_multiline_size("data:", p_event->data);
  if (p_event->id)
    frame_size += strlen("id: ") + strlen(p_event->id) + 1;
  if (p_event->event)
    frame_size += strlen("event: ") + strlen(p_event->event) + 1;

  char retry[32] = {0};
  if (p_event->retry_ms != SEOBEO_SSE_RETRY_NONE)
  {
    int written = snprintf(
        retry,
        sizeof(retry),
        "retry: %d\n",
        p_event->retry_ms);
    if (written < 0 || (size_t)written >= sizeof(retry))
      return SEOBEO_SSE_INVALID_EVENT;
    frame_size += (size_t)written;
  }

  if (frame_size > p_stream->p_handle->write_buffer_capacity)
    return SEOBEO_SSE_EVENT_TOO_LARGE;
  char *frame = malloc(frame_size);
  if (!frame)
    return -1;

  char *output = frame;
  if (p_event->id)
    output += sprintf(output, "id: %s\n", p_event->id);
  if (p_event->event)
    output += sprintf(output, "event: %s\n", p_event->event);
  if (retry[0])
  {
    size_t retry_length = strlen(retry);
    memcpy(output, retry, retry_length);
    output += retry_length;
  }
  output = sse_append_multiline(output, "data:", p_event->data);
  *output++ = '\n';

  int32 result = sse_enqueue_frame_unlocked(
      p_stream,
      frame,
      (uint32)(output - frame));
  if (result == SEOBEO_SSE_BACKPRESSURE)
    Dowa_Free(frame);
  return result;
}

int32 Seobeo_SSE_Send(
    Seobeo_SSE_Stream *p_stream,
    const Seobeo_SSE_Event *p_event)
{
  pthread_mutex_lock(&g_sse_mutex);
  int32 result = sse_send_event_unlocked(p_stream, p_event);
  pthread_mutex_unlock(&g_sse_mutex);
  return result;
}

int32 Seobeo_SSE_Send_Data(
    Seobeo_SSE_Stream *p_stream,
    const char *data)
{
  Seobeo_SSE_Event event = {
    .data = data,
    .retry_ms = SEOBEO_SSE_RETRY_NONE,
  };
  return Seobeo_SSE_Send(p_stream, &event);
}

int32 Seobeo_SSE_Send_Comment(
    Seobeo_SSE_Stream *p_stream,
    const char *comment)
{
  pthread_mutex_lock(&g_sse_mutex);
  if (!sse_is_open_unlocked(p_stream))
  {
    pthread_mutex_unlock(&g_sse_mutex);
    return -1;
  }
  size_t frame_size = sse_multiline_size(":", comment) + 1;
  if (frame_size > p_stream->p_handle->write_buffer_capacity)
  {
    pthread_mutex_unlock(&g_sse_mutex);
    return SEOBEO_SSE_EVENT_TOO_LARGE;
  }
  char *frame = malloc(frame_size);
  if (!frame)
  {
    pthread_mutex_unlock(&g_sse_mutex);
    return -1;
  }
  char *output = sse_append_multiline(frame, ":", comment);
  *output++ = '\n';
  int32 result = sse_enqueue_frame_unlocked(
      p_stream,
      frame,
      (uint32)(output - frame));
  if (result == SEOBEO_SSE_BACKPRESSURE)
    Dowa_Free(frame);
  pthread_mutex_unlock(&g_sse_mutex);
  return result;
}

int32 Seobeo_SSE_Flush(Seobeo_SSE_Stream *p_stream)
{
  pthread_mutex_lock(&g_sse_mutex);
  int32 result = sse_is_open_unlocked(p_stream)
      ? sse_drain_pending_unlocked(p_stream)
      : -1;
  pthread_mutex_unlock(&g_sse_mutex);
  return result;
}

boolean Seobeo_SSE_Is_Open(const Seobeo_SSE_Stream *p_stream)
{
  pthread_mutex_lock(&g_sse_mutex);
  boolean open = sse_is_open_unlocked(p_stream);
  pthread_mutex_unlock(&g_sse_mutex);
  return open;
}

boolean Seobeo_SSE_Retain(Seobeo_SSE_Stream *p_stream)
{
  if (!p_stream)
    return FALSE;
  pthread_mutex_lock(&g_sse_mutex);
  unsigned int references = atomic_load(&p_stream->ref_count);
  if (references == 0)
  {
    pthread_mutex_unlock(&g_sse_mutex);
    return FALSE;
  }
  atomic_fetch_add(&p_stream->ref_count, 1);
  pthread_mutex_unlock(&g_sse_mutex);
  return TRUE;
}

void Seobeo_SSE_Release(Seobeo_SSE_Stream *p_stream)
{
  if (!p_stream)
    return;
  pthread_mutex_lock(&g_sse_mutex);
  sse_release_unlocked(p_stream);
  pthread_mutex_unlock(&g_sse_mutex);
}

void Seobeo_SSE_Close(Seobeo_SSE_Stream *p_stream)
{
  if (!p_stream)
    return;
  pthread_mutex_lock(&g_sse_mutex);
  if (p_stream->started && !p_stream->closed && p_stream->p_handle)
  {
    sse_drain_pending_unlocked(p_stream);
    if (p_stream->managed)
      shutdown(p_stream->p_handle->socket, SHUT_RDWR);
  }
  sse_clear_pending_unlocked(p_stream);
  p_stream->closed = TRUE;
  pthread_mutex_unlock(&g_sse_mutex);
}

Seobeo_SSE_Stream *Seobeo_SSE_Server_Attach(
    Seobeo_Handle *p_handle,
    const char *path)
{
  if (!p_handle || !path)
    return NULL;
  Seobeo_SSE_Stream *p_stream = malloc(sizeof(*p_stream));
  if (!p_stream)
    return NULL;
  memset(p_stream, 0, sizeof(*p_stream));
  p_stream->path = strdup(path);
  if (!p_stream->path)
  {
    Dowa_Free(p_stream);
    return NULL;
  }
  atomic_init(&p_stream->ref_count, 1);
  if (Seobeo_SSE_Start(p_stream, p_handle) < 0)
  {
    Dowa_Free(p_stream->path);
    Dowa_Free(p_stream);
    return NULL;
  }
  p_stream->managed = TRUE;
  p_handle->is_sse = TRUE;

  pthread_mutex_lock(&g_sse_mutex);
  Dowa_Array_Push(g_sse_streams, p_stream);
  pthread_mutex_unlock(&g_sse_mutex);
  return p_stream;
}

void Seobeo_SSE_Server_Detach_Handle(Seobeo_Handle *p_handle)
{
  if (!p_handle)
    return;
  p_handle->is_sse = FALSE;
  pthread_mutex_lock(&g_sse_mutex);
  size_t count = Dowa_Array_Length(g_sse_streams);
  for (size_t i = 0; i < count; i++)
  {
    Seobeo_SSE_Stream *p_stream = g_sse_streams[i];
    if (p_stream->p_handle != p_handle)
      continue;
    p_stream->closed = TRUE;
    p_stream->p_handle = NULL;
    sse_clear_pending_unlocked(p_stream);
    g_sse_streams[i] = Dowa_Array_Pop(g_sse_streams);
    sse_release_unlocked(p_stream);
    break;
  }
  pthread_mutex_unlock(&g_sse_mutex);
}

void Seobeo_SSE_Server_Destroy(void)
{
  pthread_mutex_lock(&g_sse_mutex);
  size_t count = Dowa_Array_Length(g_sse_streams);
  for (size_t i = 0; i < count; i++)
  {
    Seobeo_SSE_Stream *p_stream = g_sse_streams[i];
    if (!p_stream)
      continue;
    p_stream->closed = TRUE;
    if (p_stream->p_handle)
      p_stream->p_handle->is_sse = FALSE;
    p_stream->p_handle = NULL;
    sse_clear_pending_unlocked(p_stream);
    sse_release_unlocked(p_stream);
  }
  Dowa_Array_Free(g_sse_streams);
  pthread_mutex_unlock(&g_sse_mutex);
}