#include "seobeo/seobeo.h"

static pthread_once_t g_sigpipe_once = PTHREAD_ONCE_INIT;

static void Seobeo_Process_Ignore_Sigpipe(void)
{
  signal(SIGPIPE, SIG_IGN);
}

static void Seobeo_Socket_Disable_Sigpipe(int socket_fd)
{
#ifdef SO_NOSIGPIPE
  int enabled = 1;
  setsockopt(socket_fd, SOL_SOCKET, SO_NOSIGPIPE, &enabled, sizeof(enabled));
#else
  (void)socket_fd;
#endif
}

static void Seobeo_Socket_Set_Close_On_Exec(int socket_fd)
{
  int flags = fcntl(socket_fd, F_GETFD, 0);
  if (flags >= 0)
    fcntl(socket_fd, F_SETFD, flags | FD_CLOEXEC);
}

static ssize_t Seobeo_Socket_Write(
    int socket_fd,
    const void *buffer,
    size_t length)
{
#ifdef MSG_NOSIGNAL
  return send(socket_fd, buffer, length, MSG_NOSIGNAL);
#else
  return send(socket_fd, buffer, length, 0);
#endif
}

Seobeo_Handle *Seobeo_Stream_Handle_Server_Create(const char *host,  const char* port)
{
  pthread_once(&g_sigpipe_once, Seobeo_Process_Ignore_Sigpipe);
  Seobeo_Handle *p_handle;
  struct addrinfo hints, *server_infos, *free_server_info;
  int32 socket_fd = -1;
  int32 yes = 1;

  memset(&hints, 0, sizeof hints);
  hints.ai_family   = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_protocol = IPPROTO_TCP;
  hints.ai_flags    = AI_PASSIVE;

  if (getaddrinfo(host, port, &hints, &server_infos) != 0)
  { perror("getaddrinfo"); return NULL; }

  for
  (
    free_server_info = server_infos;
    free_server_info != NULL;
    free_server_info = free_server_info->ai_next
  )
  {
    int socket_type = free_server_info->ai_socktype;
#ifdef SOCK_CLOEXEC
    socket_type |= SOCK_CLOEXEC;
#endif
    if((socket_fd = socket(free_server_info->ai_family,
                        socket_type, free_server_info->ai_protocol)) == -1)
    { perror("socket"); continue; }
    Seobeo_Socket_Set_Close_On_Exec(socket_fd);
    Seobeo_Socket_Disable_Sigpipe(socket_fd);

     if (setsockopt(socket_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes)) == -1)
     {
       perror("setsockopt SO_REUSEADDR");
       close(socket_fd);
       socket_fd = -1;
       continue;
     }

#ifdef SO_REUSEPORT
     // SO_REUSEPORT allows multiple threads/processes to bind to the same port
     // The kernel will distribute incoming connections among them
     if (setsockopt(socket_fd, SOL_SOCKET, SO_REUSEPORT, &yes, sizeof(yes)) == -1)
     {
       perror("setsockopt SO_REUSEPORT");
       close(socket_fd);
       socket_fd = -1;
       continue;
     }
#endif

     if (bind(socket_fd, free_server_info->ai_addr, free_server_info->ai_addrlen) == -1)
     {
       perror("bind");
       close(socket_fd);
       socket_fd = -1;
       continue;
     }

     break;
  }
  freeaddrinfo(server_infos);

  if (socket_fd < 0)
    return NULL;

  if (listen(socket_fd, 16) != 0)
  { 
    Seobeo_Log(SEOBEO_DEBUG, "Closing socket: %d\n", socket_fd);
    perror("listen"); close(socket_fd); return NULL; 
  }

  int flags = fcntl(socket_fd, F_GETFL, 0);
  if (flags < 0 ||
      fcntl(socket_fd, F_SETFL, flags | O_NONBLOCK) != 0)
  {
    perror("fcntl");
    close(socket_fd);
    return NULL;
  }

  p_handle = calloc(1, sizeof(*p_handle));
  if (!p_handle)
  {
    close(socket_fd);
    return NULL;
  }
  p_handle->socket = socket_fd;
  p_handle->type = SEOBEO_STREAM_TYPE_SERVER;
  p_handle->connected = FALSE;

  p_handle->host = strdup(host != NULL ? host : "localhost");
  p_handle->port = strdup(port);

  p_handle->ssl_ctx              = NULL;
  p_handle->ssl                  = NULL;


  p_handle->read_buffer = malloc(sizeof(*p_handle->read_buffer) * INITIAL_BUFFER_CAPACITY);
  p_handle->read_buffer_capacity = INITIAL_BUFFER_CAPACITY;
  p_handle->read_buffer_len = 0;

  p_handle->write_buffer = malloc(sizeof(*p_handle->read_buffer) * INITIAL_BUFFER_CAPACITY);
  p_handle->write_buffer_capacity = INITIAL_BUFFER_CAPACITY;
  p_handle->write_buffer_len = 0;

  p_handle->destroyed = FALSE;
  p_handle->is_sse = FALSE;

  if (!p_handle->host ||
      !p_handle->port ||
      !p_handle->read_buffer ||
      !p_handle->write_buffer)
  {
    Seobeo_Handle_Destroy(p_handle);
    return NULL;
  }

  return p_handle;
}


Seobeo_Handle *Seobeo_Stream_Handle_Client_Create(const char *host,  const char* port, boolean use_tls)
{
  pthread_once(&g_sigpipe_once, Seobeo_Process_Ignore_Sigpipe);
  Seobeo_Handle *p_handle;
  p_handle = malloc(sizeof(*p_handle));

  struct addrinfo hints, *server_infos;
  int32 socket_fd;  // Need this for setsockopt 

  memset(&hints, 0, sizeof hints);
  hints.ai_family   = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;

  if (getaddrinfo(host, port, &hints, &server_infos) != 0)
  { perror("getaddrinfo"); return NULL; }


  int socket_type = server_infos->ai_socktype;
#ifdef SOCK_CLOEXEC
  socket_type |= SOCK_CLOEXEC;
#endif
  if((socket_fd = socket(server_infos->ai_family,
                      socket_type, server_infos->ai_protocol)) == -1)
  { perror("socket"); return NULL; }
  Seobeo_Socket_Set_Close_On_Exec(socket_fd);
  Seobeo_Socket_Disable_Sigpipe(socket_fd);

  if (connect(socket_fd, server_infos->ai_addr, server_infos->ai_addrlen) != 0)
  { perror("connect"); return NULL; }
  freeaddrinfo(server_infos);

  // Set non-blocking
  int flags = fcntl(socket_fd, F_GETFL, 0);
  if (flags == -1 || fcntl(socket_fd, F_SETFL, flags | O_NONBLOCK) != 0)
  { perror("fcntl"); close(socket_fd); free(p_handle); return NULL; }

  p_handle->socket = socket_fd;
  p_handle->type = SEOBEO_STREAM_TYPE_CLIENT;

  p_handle->ssl_ctx = NULL;
  p_handle->ssl = NULL;

  if (use_tls)
  {
    if (Seobeo_SSL_Setup_Client(p_handle, host, socket_fd) != 0)
    {
      free(p_handle);
      return NULL;
    }
  }
  p_handle->connected = TRUE;

  p_handle->host = host != NULL ? strdup(host) : "localhost";
  p_handle->port = strdup(port);

  p_handle->read_buffer = malloc(sizeof(*p_handle->read_buffer) * INITIAL_BUFFER_CAPACITY);
  p_handle->read_buffer_capacity = INITIAL_BUFFER_CAPACITY;
  p_handle->read_buffer_len = 0;

  p_handle->write_buffer = malloc(sizeof(*p_handle->read_buffer) * INITIAL_BUFFER_CAPACITY);
  p_handle->write_buffer_capacity = INITIAL_BUFFER_CAPACITY;
  p_handle->write_buffer_len = 0;

  p_handle->destroyed = FALSE;
  p_handle->is_sse = FALSE;

  return p_handle;
}

Seobeo_Handle *Seobeo_Stream_Handle_Server_Accept(Seobeo_Handle *p_server_handle)
{
  struct sockaddr_storage addr;
  socklen_t addrlen = sizeof addr;
  char client_inet_addr[INET6_ADDRSTRLEN];

  int client_fd = accept(
      p_server_handle->socket,
      (struct sockaddr*)&addr,
      &addrlen);
  inet_ntop(
      addr.ss_family,
      Seobeo_Get_IP4_Or_IP6((struct sockaddr *)&addr),
      client_inet_addr, sizeof client_inet_addr);
  if (client_fd == -1) return NULL;
  Seobeo_Socket_Set_Close_On_Exec(client_fd);
  Seobeo_Socket_Disable_Sigpipe(client_fd);

  // Set non blocking...
  int flags = fcntl(client_fd, F_GETFL, 0);
  if (flags == -1) return NULL;
  fcntl(client_fd, F_SETFL, flags | O_NONBLOCK);

  Seobeo_Handle *p_client_handle        = malloc(sizeof *p_client_handle);

  p_client_handle->socket               = client_fd;
  p_client_handle->type                 = SEOBEO_STREAM_TYPE_CLIENT;
  p_client_handle->connected            = TRUE;

  // TODO: support SSL in the future.
  p_client_handle->ssl_ctx              = NULL;
  p_client_handle->ssl                  = NULL;

  p_client_handle->host                 = strdup(client_inet_addr);
  p_client_handle->port                 = NULL;

  p_client_handle->read_buffer_capacity = p_server_handle->read_buffer_capacity;
  p_client_handle->read_buffer_len      = 0;
  p_client_handle->read_buffer          = malloc(p_client_handle->read_buffer_capacity);

  p_client_handle->write_buffer_capacity = p_server_handle->write_buffer_capacity;
  p_client_handle->write_buffer_len      = 0;
  p_client_handle->write_buffer          = malloc(p_client_handle->write_buffer_capacity);

  p_client_handle->read_buffer_used      = 0;
  p_client_handle->file                  = NULL;
  p_client_handle->text_copy             = NULL;
  p_client_handle->file_name             = NULL;
  p_client_handle->destroyed             = FALSE;
  p_client_handle->is_sse                = FALSE;

  return p_client_handle;
}

void Seobeo_Handle_Destroy(Seobeo_Handle *p_handle)
{
  if (!p_handle) return;

  boolean expected = FALSE;
  // Need to check
  if (!atomic_compare_exchange_strong(&p_handle->destroyed, &expected, TRUE))
  {
    return;
  }

  Seobeo_SSE_Server_Detach_Handle(p_handle);

  if (p_handle->host) Dowa_Free(p_handle->host);
  if (p_handle->port) Dowa_Free(p_handle->port);

  Seobeo_SSL_Cleanup(p_handle);

  if (p_handle->socket >= 0) {
    Seobeo_Log(SEOBEO_DEBUG, "Closing handle socket: %d\n", p_handle->socket);
    close(p_handle->socket);
  }

  if (p_handle->read_buffer) Dowa_Free(p_handle->read_buffer);
  if (p_handle->write_buffer) Dowa_Free(p_handle->write_buffer);

  Dowa_Free(p_handle);
}


int32 Seobeo_Handle_Flush(Seobeo_Handle *p_handle)
{
  if (!p_handle || !p_handle->write_buffer)
    return -1;
  uint32 total = p_handle->write_buffer_len;
  uint32 sent  = 0;

  Seobeo_Log(SEOBEO_DEBUG, "Write buffer total: %d\n", p_handle->write_buffer_len);

  while (sent < total)
  {
    if (p_handle->ssl)
    {
      int n = Seobeo_SSL_Write(p_handle, p_handle->write_buffer + sent, total - sent);
      if (n < 0)
      {
        if (sent)
        {
          memmove(
              p_handle->write_buffer,
              p_handle->write_buffer + sent,
              total - sent);
          p_handle->write_buffer_len = total - sent;
        }
        return -1;
      }
      if (n == 0)
      {
        if (sent)
        {
          memmove(
              p_handle->write_buffer,
              p_handle->write_buffer + sent,
              total - sent);
          p_handle->write_buffer_len = total - sent;
        }
        return 1;
      }
      sent += (uint32)n;
    }else
    {
      Seobeo_Log(SEOBEO_DEBUG, "Flushing socket: %d\n", p_handle->socket);
      ssize_t n = Seobeo_Socket_Write(
        p_handle->socket,
        p_handle->write_buffer + sent,
        total - sent
      );
      if (n < 0) {
        if (errno == EINTR)  continue;
        if (errno == EAGAIN || errno == EWOULDBLOCK)
        {
          if (sent)
          {
            memmove(
                p_handle->write_buffer,
                p_handle->write_buffer + sent,
                total - sent);
            p_handle->write_buffer_len = total - sent;
          }
          return 1;
        }
        if (sent)
        {
          memmove(
              p_handle->write_buffer,
              p_handle->write_buffer + sent,
              total - sent);
          p_handle->write_buffer_len = total - sent;
        }
        return -1;
      }
      sent += (uint32)n;
    }
  }

  p_handle->write_buffer_len = 0;
  return 0;
}

int32 Seobeo_Handle_Queue(Seobeo_Handle *p_handle, const uint8 *data, uint32 data_size)
{
  if (p_handle->write_buffer_len + data_size > p_handle->write_buffer_capacity)
  {
    int32 rc = Seobeo_Handle_Flush(p_handle);
    if (rc < 0) return -1;
    if (rc > 0) return 1;
  }

  if (data_size > p_handle->write_buffer_capacity)
  {
    uint32 offset = 0;
    while (offset < data_size)
    {
      ssize_t n;
      if (p_handle->ssl)
      {
        n = Seobeo_SSL_Write(p_handle, data + offset, data_size - offset);
      }
      else
      {
        n = Seobeo_Socket_Write(
            p_handle->socket,
            data + offset,
            data_size - offset);
      }

      if (n==0)
      {
        // DEBUG
        Seobeo_Log(SEOBEO_DEBUG, "Write offset: %d\n", offset);
        break;
      }
      if (n < 0)
      {
        if (errno == EINTR || errno == EAGAIN)
        {
          // DEBUG
          // printf("Partial write, returning early (offset=%d)\n", offset);
          continue;
        }
        if (errno == EAGAIN) return 1;
        return -1;
      }
      offset += (uint32)n;
      // DEBUG
      Seobeo_Log(SEOBEO_DEBUG, "Write completed - offset: %d, data_size: %d\n", offset, data_size);
    }
    // DEBUG
    Seobeo_Log(SEOBEO_DEBUG, "Total bytes written: %d\n", offset);
    return 0;
  }

  if (!p_handle)
  {
    Seobeo_Log(SEOBEO_ERROR, "p_handle is NULL before memcpy\n");
    return -1;
  }
  
  if (!p_handle->write_buffer)
  {
    Seobeo_Log(SEOBEO_ERROR, "p_handle->write_buffer is NULL (len=%u, size=%u)\n",
            p_handle->write_buffer_len, data_size);
    return -1;
  }
  
  Seobeo_Log(SEOBEO_DEBUG, "memcpy -> dest=%p (write_buffer=%p + offset=%u), src=%p, size=%u\n",
          p_handle->write_buffer + p_handle->write_buffer_len,
          p_handle->write_buffer,
          p_handle->write_buffer_len,
          data,
          data_size); 
  memcpy(p_handle->write_buffer + p_handle->write_buffer_len,
         data,
         data_size);
  p_handle->write_buffer_len += data_size;
  return 0;
}

int32 Seobeo_Handle_Read(Seobeo_Handle *p_handle)
{
  int32 read_size;
  if (!p_handle) return -1;

  // How many bytes we can still read into the buffer
  uint32 free_space = p_handle->read_buffer_capacity - p_handle->read_buffer_len;
  if (free_space == 0)
    return -1;

  if (p_handle->ssl)
  {
    read_size = Seobeo_SSL_Read(p_handle, p_handle->read_buffer + p_handle->read_buffer_len, free_space);
    if (read_size < 0) return read_size;  // -1 for error, -2 for closed
    if (read_size == 0) return 0;  // would block
  }
  else
  {
    read_size = (int32)read(p_handle->socket,
                            p_handle->read_buffer + p_handle->read_buffer_len,
                            free_space);
    if (read_size == 0) return -2; 
    if (read_size < 0)
    {
      if (errno == EAGAIN || errno == EWOULDBLOCK) return 0;
      return -1;
    }
  }

  p_handle->read_buffer_len += (uint32)read_size;
  return read_size;
}

void Seobeo_Handle_Consume(Seobeo_Handle *p_handle, uint32 consumed)
{
  if (consumed >= p_handle->read_buffer_len)
  {
    p_handle->read_buffer_len = 0;
    return;
  }

  // Slide remaining bytes to the front
  memmove(
    p_handle->read_buffer,
    p_handle->read_buffer + consumed,
    p_handle->read_buffer_len - consumed
  );
  p_handle->read_buffer_len -= consumed;
}

void *Seobeo_Get_IP4_Or_IP6(struct sockaddr *sa)
{
  if (sa->sa_family == AF_INET) 
  {
    return &(((struct sockaddr_in*)sa)->sin_addr);
  }

  return &(((struct sockaddr_in6*)sa)->sin6_addr);
}