#ifndef _DEFAULT_SOURCE
#define _DEFAULT_SOURCE
#endif

#include "cst_net.h"

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

#ifdef _WIN32
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#if !defined(_WIN32_WINNT) || _WIN32_WINNT < 0x0600
#undef _WIN32_WINNT
#define _WIN32_WINNT 0x0600 /* WSAPoll */
#endif
#include <winsock2.h>
#include <ws2tcpip.h>
#else
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <poll.h>
#include <signal.h>
#include <sys/socket.h>
#include <unistd.h>
#endif

#ifndef MSG_NOSIGNAL
#define MSG_NOSIGNAL 0
#endif

#define NET_SEND_TIMEOUT_MS 20000

#ifdef _WIN32
static void set_err_win(char *err, size_t err_sz, const char *what) {
  if (err && err_sz)
    snprintf(err, err_sz, "%s (winsock error %d)", what, WSAGetLastError());
}
#define SET_ERR(e, n, w) set_err_win(e, n, w)
#define NET_FD(s) ((SOCKET)(s))
static void net_close_raw(os_socket_t s) { closesocket((SOCKET)s); }
static int net_last_would_block(void) {
  int e = WSAGetLastError();
  return e == WSAEWOULDBLOCK || e == WSAEINTR;
}
static int net_last_conn_reset(void) {
  int e = WSAGetLastError();
  return e == WSAECONNRESET || e == WSAECONNABORTED || e == WSAENOTSOCK;
}
#else
static void set_err(char *err, size_t err_sz, const char *what) {
  if (err && err_sz)
    snprintf(err, err_sz, "%s", what);
}
#define SET_ERR(e, n, w) set_err(e, n, w)
#define NET_FD(s) ((int)(s))
static void net_close_raw(os_socket_t s) { close(s); }
static int net_last_would_block(void) { return errno == EINTR; }
static int net_last_conn_reset(void) {
  return errno == ECONNRESET || errno == EPIPE || errno == EBADF;
}
#endif

int os_net_init(void) {
#ifdef _WIN32
  WSADATA w;
  if (WSAStartup(MAKEWORD(2, 2), &w) != 0)
    return -1;
#else
  /* A hung-up client must not kill the process with SIGPIPE. */
  signal(SIGPIPE, SIG_IGN);
#endif
  return 0;
}

void os_net_cleanup(void) {
#ifdef _WIN32
  WSACleanup();
#endif
}

/* poll, not select: FD_SET past FD_SETSIZE writes out of bounds, and a
 * failed connect shows up as an error event, not as writable, on Windows. */
static int wait_io(os_socket_t s, int want_write, int timeout_ms) {
#ifdef _WIN32
  WSAPOLLFD p;
  p.fd = NET_FD(s);
  p.events = want_write ? POLLWRNORM : POLLRDNORM;
  p.revents = 0;
  int rc = WSAPoll(&p, 1, timeout_ms < 0 ? -1 : timeout_ms);
#else
  struct pollfd p;
  p.fd = NET_FD(s);
  p.events = want_write ? POLLOUT : POLLIN;
  p.revents = 0;
  int rc = poll(&p, 1, timeout_ms < 0 ? -1 : timeout_ms);
#endif
  if (rc > 0)
    return 1;
  if (rc == 0)
    return 0; /* timed out */
  return -1;
}

/* Host to bind: AF_INET unless the name is an IPv6 literal, so "localhost"
 * and "127.0.0.1" always mean the IPv4 loopback a browser will reach. */
static int resolve_bind(const char *host, int port, struct sockaddr_storage *sa,
                        socklen_t *sa_len) {
  memset(sa, 0, sizeof(*sa));
  int family = AF_INET;
  const char *h = host && host[0] ? host : "127.0.0.1";
  if (strchr(h, ':'))
    family = AF_INET6;
  struct addrinfo hints, *res = NULL;
  memset(&hints, 0, sizeof(hints));
  hints.ai_family = family;
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_flags = AI_PASSIVE;
  char portstr[16];
  snprintf(portstr, sizeof(portstr), "%d", port);
  if (getaddrinfo(h, portstr, &hints, &res) != 0 || !res)
    return -1;
  memcpy(sa, res->ai_addr, res->ai_addrlen);
  *sa_len = (socklen_t)res->ai_addrlen;
  freeaddrinfo(res);
  return 0;
}

os_socket_t os_net_listen(const char *host, int port, char *err, size_t err_sz) {
  struct sockaddr_storage sa;
  socklen_t sa_len;
  if (resolve_bind(host, port, &sa, &sa_len) != 0) {
    SET_ERR(err, err_sz, "cannot resolve the bind address");
    return OS_SOCKET_BAD;
  }
  os_socket_t s = (os_socket_t)socket(sa.ss_family, SOCK_STREAM, 0);
  if (s == OS_SOCKET_BAD) {
    SET_ERR(err, err_sz, "cannot create a socket");
    return OS_SOCKET_BAD;
  }
  int one = 1;
#ifdef _WIN32
  /* SO_REUSEADDR on Windows lets another process bind the same port and
   * take over our connections */
  setsockopt(NET_FD(s), SOL_SOCKET, SO_EXCLUSIVEADDRUSE, (const char *)&one,
             sizeof(one));
#else
  setsockopt(NET_FD(s), SOL_SOCKET, SO_REUSEADDR, (const char *)&one,
             sizeof(one));
#endif
#ifdef IPV6_V6ONLY
  if (sa.ss_family == AF_INET6)
    setsockopt(NET_FD(s), IPPROTO_IPV6, IPV6_V6ONLY, (const char *)&one,
               sizeof(one));
#endif
  if (bind(NET_FD(s), (const struct sockaddr *)&sa, sa_len) != 0) {
    SET_ERR(err, err_sz, "the address is already in use");
    net_close_raw(s);
    return OS_SOCKET_BAD;
  }
  if (listen(NET_FD(s), 32) != 0) {
    SET_ERR(err, err_sz, "listen failed");
    net_close_raw(s);
    return OS_SOCKET_BAD;
  }
  return s;
}

int os_net_port(os_socket_t s) {
  struct sockaddr_storage sa;
  socklen_t len = sizeof(sa);
  if (getsockname(NET_FD(s), (struct sockaddr *)&sa, &len) != 0)
    return 0;
  if (sa.ss_family == AF_INET)
    return ntohs(((struct sockaddr_in *)&sa)->sin_port);
  if (sa.ss_family == AF_INET6)
    return ntohs(((struct sockaddr_in6 *)&sa)->sin6_port);
  return 0;
}

void os_net_peer(os_socket_t s, char *out, size_t out_sz) {
  if (!out || out_sz == 0)
    return;
  out[0] = 0;
  struct sockaddr_storage sa;
  socklen_t len = sizeof(sa);
  if (getpeername(NET_FD(s), (struct sockaddr *)&sa, &len) != 0)
    return;
  char ip[64] = "";
  int port = 0;
  if (sa.ss_family == AF_INET) {
    if (!inet_ntop(AF_INET, &((struct sockaddr_in *)&sa)->sin_addr, ip,
                   sizeof(ip)))
      return;
    port = ntohs(((struct sockaddr_in *)&sa)->sin_port);
  } else if (sa.ss_family == AF_INET6) {
    if (!inet_ntop(AF_INET6, &((struct sockaddr_in6 *)&sa)->sin6_addr, ip,
                   sizeof(ip)))
      return;
    port = ntohs(((struct sockaddr_in6 *)&sa)->sin6_port);
  } else {
    return;
  }
  if (strchr(ip, ':'))
    snprintf(out, out_sz, "[%s]:%d", ip, port);
  else
    snprintf(out, out_sz, "%s:%d", ip, port);
}

os_socket_t os_net_connect(const char *host, int port, int timeout_ms) {
  const char *h = host && host[0] ? host : "127.0.0.1";
  struct addrinfo hints, *res = NULL;
  memset(&hints, 0, sizeof(hints));
  hints.ai_family = strchr(h, ':') ? AF_INET6 : AF_INET;
  hints.ai_socktype = SOCK_STREAM;
  char portstr[16];
  snprintf(portstr, sizeof(portstr), "%d", port);
  if (getaddrinfo(h, portstr, &hints, &res) != 0 || !res)
    return OS_SOCKET_BAD;
  os_socket_t s = (os_socket_t)socket(res->ai_family, SOCK_STREAM, 0);
  if (s == OS_SOCKET_BAD) {
    freeaddrinfo(res);
    return OS_SOCKET_BAD;
  }
  /* Bounded wait: connect() non-blocking, poll, then read SO_ERROR. A dead
   * server must not hang the caller. */
#ifdef _WIN32
  u_long nb = 1;
  ioctlsocket(NET_FD(s), FIONBIO, &nb);
#else
  int fl = fcntl(s, F_GETFL, 0);
  fcntl(s, F_SETFL, fl | O_NONBLOCK);
#endif
  int rc = connect(NET_FD(s), res->ai_addr, (int)res->ai_addrlen);
  freeaddrinfo(res);
  if (rc != 0) {
#ifdef _WIN32
    int in_progress = WSAGetLastError() == WSAEWOULDBLOCK;
#else
    int in_progress = errno == EINPROGRESS || errno == EINTR;
#endif
    if (!in_progress || wait_io(s, 1, timeout_ms) <= 0) {
      net_close_raw(s);
      return OS_SOCKET_BAD;
    }
    int soerr = 0;
    socklen_t elen = sizeof(soerr);
    if (getsockopt(NET_FD(s), SOL_SOCKET, SO_ERROR, (char *)&soerr, &elen) != 0 ||
        soerr != 0) {
      net_close_raw(s);
      return OS_SOCKET_BAD;
    }
  }
#ifdef _WIN32
  nb = 0;
  ioctlsocket(NET_FD(s), FIONBIO, &nb);
#else
  fcntl(s, F_SETFL, fl);
#endif
  return s;
}

os_socket_t os_net_accept(os_socket_t s, int timeout_ms, int *timed_out) {
  if (timed_out)
    *timed_out = 0;
  int w = wait_io(s, 0, timeout_ms);
  if (w == 0) {
    if (timed_out)
      *timed_out = 1;
    return OS_SOCKET_BAD;
  }
  if (w < 0)
    return OS_SOCKET_BAD;
  os_socket_t c = (os_socket_t)accept(NET_FD(s), NULL, NULL);
  if (c == OS_SOCKET_BAD)
    return OS_SOCKET_BAD;
  int one = 1;
  setsockopt(NET_FD(c), IPPROTO_TCP, TCP_NODELAY, (const char *)&one,
             sizeof(one));
  return c;
}

long os_net_recv(os_socket_t s, void *buf, size_t n, int timeout_ms) {
  int w = wait_io(s, 0, timeout_ms);
  if (w == 0)
    return -2;
  if (w < 0)
    return -1;
  size_t want = n > (size_t)(1 << 20) ? (size_t)(1 << 20) : n;
  int r = (int)recv(NET_FD(s), (char *)buf, (int)want, 0);
  if (r == 0)
    return 0;
  if (r < 0)
    return net_last_conn_reset() ? 0 : (net_last_would_block() ? -2 : -1);
  return r;
}

int os_net_send_all(os_socket_t s, const void *buf, size_t n, int timeout_ms) {
  const char *p = buf;
  size_t off = 0;
  while (off < n) {
    int w = wait_io(s, 1, timeout_ms ? timeout_ms : NET_SEND_TIMEOUT_MS);
    if (w <= 0)
      return -1;
    long r = (long)send(NET_FD(s), p + off, (int)(n - off), MSG_NOSIGNAL);
    if (r <= 0) {
      if (net_last_would_block())
        continue;
      return -1;
    }
    off += (size_t)r;
  }
  return 0;
}

void os_net_shutdown(os_socket_t s) {
#ifdef _WIN32
  shutdown((SOCKET)s, SD_BOTH);
#else
  shutdown(s, SHUT_RDWR);
#endif
}

void os_net_close(os_socket_t s) {
  if (s == OS_SOCKET_BAD)
    return;
  net_close_raw(s);
}