#include "../beaker.h"
#include "beaker_globals.h"
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <poll.h>
#include <pthread.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <time.h>
#include <unistd.h>

#define MAX_PENDING_CONNECTIONS 128

static volatile sig_atomic_t g_shutdown_requested = 0;

static void signal_handler(int sig) {
  (void)sig;
  g_shutdown_requested = 1;
}

typedef struct {
  _Atomic(size_t) sequence;
  int socket;
} WorkSlot;

typedef struct {
  _Atomic(size_t) head;
  _Atomic(size_t) tail;
  _Atomic(int) shutdown;
  WorkSlot slots[MAX_PENDING_CONNECTIONS];
  pthread_mutex_t mutex;
  pthread_cond_t cond;
} WorkQueue;

static WorkQueue g_work_queue;

static void work_queue_init(WorkQueue *queue) {
  atomic_store(&queue->head, 0);
  atomic_store(&queue->tail, 0);
  atomic_store(&queue->shutdown, 0);
  for (int i = 0; i < MAX_PENDING_CONNECTIONS; i++) {
    atomic_store(&queue->slots[i].sequence, (size_t)i);
  }
  pthread_mutex_init(&queue->mutex, NULL);
  pthread_cond_init(&queue->cond, NULL);
}

static void work_queue_destroy(WorkQueue *queue) {
  pthread_mutex_destroy(&queue->mutex);
  pthread_cond_destroy(&queue->cond);
}

static int work_queue_push(WorkQueue *queue, int client_socket) {
  size_t tail = atomic_load(&queue->tail);

  for (;;) {
    WorkSlot *slot = &queue->slots[tail % MAX_PENDING_CONNECTIONS];
    size_t seq = atomic_load(&slot->sequence);
    intptr_t diff = (intptr_t)seq - (intptr_t)tail;

    if (diff == 0) {
      if (atomic_compare_exchange_weak(&queue->tail, &tail, tail + 1)) {
        slot->socket = client_socket;
        atomic_store(&slot->sequence, tail + 1);
        pthread_cond_signal(&queue->cond);
        return 0;
      }
    } else if (diff < 0) {
      return -1;
    } else {
      tail = atomic_load(&queue->tail);
    }
  }
}

static int work_queue_pop(WorkQueue *queue) {
  size_t head = atomic_load(&queue->head);

  for (;;) {
    if (atomic_load(&queue->shutdown))
      return -1;

    WorkSlot *slot = &queue->slots[head % MAX_PENDING_CONNECTIONS];
    size_t seq = atomic_load(&slot->sequence);
    intptr_t diff = (intptr_t)seq - (intptr_t)(head + 1);

    if (diff == 0) {
      if (atomic_compare_exchange_weak(&queue->head, &head, head + 1)) {
        int fd = slot->socket;
        atomic_store(&slot->sequence, head + MAX_PENDING_CONNECTIONS);
        return fd;
      }
    } else if (diff < 0) {
      pthread_mutex_lock(&queue->mutex);
      if (!atomic_load(&queue->shutdown) && atomic_load(&queue->head) == head) {
        pthread_cond_wait(&queue->cond, &queue->mutex);
      }
      pthread_mutex_unlock(&queue->mutex);
      head = atomic_load(&queue->head);
    } else {
      head = atomic_load(&queue->head);
    }
  }
}

static int get_optimal_thread_count(void) {
  long cores = sysconf(_SC_NPROCESSORS_ONLN);
  if (cores < 1)
    cores = 1;
  return (int)(cores * 2);
}

void handle_client_connection(int new_socket);

static void *worker_thread(void *arg) {
  (void)arg;
  while (1) {
    int client_socket = work_queue_pop(&g_work_queue);
    if (client_socket < 0) {
      break;
    }
    handle_client_connection(client_socket);
  }
  return NULL;
}

static int initialize_server_socket(const char *ip, int port,
                                    int *server_fd_out,
                                    struct sockaddr_in *address_out) {

  if ((*server_fd_out = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
    beaker_log_errno_format(
        "ERROR", "initialize_server_socket: Failed to create socket.\n");
    return -1;
  }

  int opt = 1;

  if (setsockopt(*server_fd_out, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt))) {
    beaker_log_errno("setsockopt SO_REUSEADDR failed");
  }

// Needed for FreeBSD support. On macOS this allows multiple processes to
// bind the same TCP port, which is surprising for a single-instance server.
#if defined(__FreeBSD__) && defined(SO_REUSEPORT)
  if (setsockopt(*server_fd_out, SOL_SOCKET, SO_REUSEPORT, &opt, sizeof(opt))) {
    beaker_log_errno_format(
        "ERROR", "initialize_server_socket: Failed to set SO_REUSEPORT.\n");
    close(*server_fd_out);
    return -1;
  }
#endif

  address_out->sin_family = AF_INET;
  address_out->sin_addr.s_addr = inet_addr(ip);
  address_out->sin_port = htons(port);

  if (bind(*server_fd_out, (struct sockaddr *)address_out,
           sizeof(*address_out)) < 0) {
    beaker_log_errno_format(
        "ERROR", "initialize_server_socket: Failed to bind socket to %s:%d.\n",
        ip, port);
    close(*server_fd_out);
    return -1;
  }

  if (listen(*server_fd_out, 10) < 0) {
    beaker_log_errno_format(
        "ERROR", "initialize_server_socket: Failed to listen on socket.\n");
    close(*server_fd_out);
    return -1;
  }

  int flags = fcntl(*server_fd_out, F_GETFL, 0);
  if (flags < 0 || fcntl(*server_fd_out, F_SETFL, flags | O_NONBLOCK) < 0) {
    beaker_log_errno("fcntl O_NONBLOCK failed");
    close(*server_fd_out);
    return -1;
  }

  beaker_log("INFO", "listening on %s:%d", ip, port);
  return 0;
}

static int set_socket_blocking(int fd) {
  int flags = fcntl(fd, F_GETFL, 0);
  if (flags < 0)
    return -1;

  if ((flags & O_NONBLOCK) == 0)
    return 0;

  return fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
}

static void finish_client_connection(int socket, char *requested_path,
                                     const char *method, const char *path,
                                     const struct timespec *started_at) {
  struct timespec finished_at;
  clock_gettime(CLOCK_MONOTONIC, &finished_at);
  double duration_ms = (finished_at.tv_sec - started_at->tv_sec) * 1000.0 +
                       (finished_at.tv_nsec - started_at->tv_nsec) / 1000000.0;
  beaker_log_request(current_request_info.remote_addr, method, path,
                     current_response_status, current_response_size,
                     duration_ms);
  free(requested_path);
  close(socket);
  current_client_socket = -1;
}

void handle_client_connection(int new_socket) {
  current_client_socket = new_socket;
  char buffer[BUFFER_SIZE] = {0};
  char method[16] = "-";
  char log_path[MAX_PATH_LEN] = "-";
  char *requested_path = NULL;
  struct timespec started_at;

  clock_gettime(CLOCK_MONOTONIC, &started_at);
  current_response_status = 0;
  current_response_size = 0;
  memset(&current_request_info, 0, sizeof(RequestInfo));

  struct sockaddr_in client_addr;
  socklen_t client_len = sizeof(client_addr);
  if (getpeername(new_socket, (struct sockaddr *)&client_addr, &client_len) !=
          0 ||
      inet_ntop(AF_INET, &client_addr.sin_addr,
                current_request_info.remote_addr,
                sizeof(current_request_info.remote_addr)) == NULL) {
    strcpy(current_request_info.remote_addr, "-");
  }

  ssize_t bytes_read = read(new_socket, buffer, BUFFER_SIZE - 1);
  if (bytes_read < 0) {
    beaker_log_errno_format(
        "ERROR",
        "handle_client_connection: Failed to read from client socket.\n");
    finish_client_connection(new_socket, requested_path, method, log_path,
                             &started_at);
    return;
  }
  buffer[bytes_read] = '\0';

  strncpy(current_request_buffer, buffer, BUFFER_SIZE - 1);
  current_request_buffer[BUFFER_SIZE - 1] = '\0';

  char request_line[MAX_PATH_LEN + 64];
  char *first_line_end = strstr(buffer, "\r\n");

  if (first_line_end == NULL) {
    beaker_log(
        "ERROR",
        "handle_client_connection: Invalid HTTP request: No CRLF found.\n");
    send_status("400 Bad Request");
    finish_client_connection(new_socket, requested_path, method, log_path,
                             &started_at);
    return;
  }
  size_t request_line_len = first_line_end - buffer;
  if (request_line_len >= sizeof(request_line)) {
    beaker_log("ERROR", "handle_client_connection: Request line too long.\n");
    send_status("400 Bad Request");
    finish_client_connection(new_socket, requested_path, method, log_path,
                             &started_at);
    return;
  }
  strncpy(request_line, buffer, request_line_len);
  request_line[request_line_len] = '\0';
  sscanf(request_line, "%15s %255s", method, log_path);

  UrlParams request_params;
  requested_path = parse_request_url(request_line, &request_params);

  if (requested_path == NULL) {
    beaker_log("ERROR", "handle_client_connection: Could not parse request "
                        "path. Sending 400 Bad Request.\n");
    send_status("400 Bad Request");
    finish_client_connection(new_socket, requested_path, method, log_path,
                             &started_at);
    return;
  }
  strncpy(log_path, requested_path, sizeof(log_path) - 1);
  log_path[sizeof(log_path) - 1] = '\0';

  bool handled = false;

  if (strncmp(requested_path, "/static/", strlen("/static/")) == 0) {

    if (serve_static_file(requested_path + strlen("/static/"))) {
      handled = true;
    }
  }

  if (!handled) {
    int best_match_handler_index = -1;
    size_t best_match_len = 0;

    for (int i = 0; i < handler_count; i++) {
      size_t handler_path_len = strlen(handlers[i].path);

      if (strncmp(requested_path, handlers[i].path, handler_path_len) == 0) {

        if (handler_path_len == strlen(requested_path) ||
            requested_path[handler_path_len] == '/') {

          if (handler_path_len > best_match_len) {
            best_match_len = handler_path_len;
            best_match_handler_index = i;
          }
        }
      }
    }

    if (best_match_handler_index != -1) {
      handlers[best_match_handler_index].handler(&request_params);
      handled = true;
    }
  }

  if (!handled) {
    const char *not_found_html = "<h1>404 Not Found</h1><p>The requested URL "
                                 "was not located on this server.</p>";
    char not_found_response[BUFFER_SIZE];
    snprintf(not_found_response, sizeof(not_found_response),
             "HTTP/1.1 404 Not Found\r\n"
             "Content-Type: text/html; charset=UTF-8\r\n"
             "Content-Length: %zu\r\n"
             "Connection: close\r\n"
             "\r\n%s",
             strlen(not_found_html), not_found_html);
    send(new_socket, not_found_response, strlen(not_found_response), 0);
    current_response_status = 404;
    current_response_size = strlen(not_found_html);
  }

  finish_client_connection(new_socket, requested_path, method, log_path,
                           &started_at);
}

int beaker_run(const char *ip, int port) {
  beaker_run_with_threads(ip, port, 0);
  return 0;
}

void beaker_run_with_threads(const char *ip, int port, int num_workers) {
  int server_fd;
  struct sockaddr_in address;
  int addrlen = sizeof(address);

  g_shutdown_requested = 0;

  struct sigaction sa;
  sa.sa_handler = signal_handler;
  sigemptyset(&sa.sa_mask);
  sa.sa_flags = 0;
  sigaction(SIGINT, &sa, NULL);
  sigaction(SIGTERM, &sa, NULL);

  if (num_workers <= 0) {
    num_workers = get_optimal_thread_count();
  }

  if (initialize_server_socket(ip, port, &server_fd, &address) != 0) {
    return;
  }

  work_queue_init(&g_work_queue);

  pthread_t threads[num_workers];
  for (int i = 0; i < num_workers; i++) {
    pthread_create(&threads[i], NULL, worker_thread, NULL);
  }

  beaker_log("DEBUG", "started %d worker threads", num_workers);

  struct pollfd pfd = {.fd = server_fd, .events = POLLIN};

  while (!g_shutdown_requested) {
    int ret = poll(&pfd, 1, 1000);
    if (ret < 0) {
      if (errno == EINTR)
        continue;
      beaker_log_errno("poll failed");
      break;
    }
    if (ret == 0)
      continue;

    int new_socket;
    while ((new_socket = accept(server_fd, (struct sockaddr *)&address,
                                (socklen_t *)&addrlen)) >= 0) {
      if (set_socket_blocking(new_socket) < 0) {
        beaker_log_errno("fcntl clear O_NONBLOCK failed");
        close(new_socket);
        continue;
      }
      if (work_queue_push(&g_work_queue, new_socket) < 0) {
        beaker_log("WARN", "work queue full; rejecting connection");
        const char *busy_response =
            "HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\n\r\n";
        send(new_socket, busy_response, strlen(busy_response), 0);
        close(new_socket);
      }
    }
  }

  beaker_log("INFO", "shutting down");

  atomic_store(&g_work_queue.shutdown, 1);
  pthread_cond_broadcast(&g_work_queue.cond);

  for (int i = 0; i < num_workers; i++) {
    pthread_join(threads[i], NULL);
  }

  work_queue_destroy(&g_work_queue);
  close(server_fd);
}

const char *beaker_get_remote_addr(void) {
  return current_request_info.remote_addr;
}

static __thread char g_header_value[MAX_VALUE_LEN];

const char *beaker_get_header(const char *name) {
  if (name == NULL)
    return "";

  size_t name_len = strlen(name);
  if (name_len == 0)
    return "";

  if (strstr(name, "\r\n") != NULL)
    return "";

  char *buffer = current_request_buffer;
  char *search_end = strstr(buffer, "\r\n\r\n");
  if (search_end == NULL)
    search_end = buffer + strlen(buffer);
  while (buffer < search_end) {
    if (strncasecmp(buffer, name, name_len) == 0 && buffer[name_len] == ':') {
      char *value_start = buffer + name_len + 1;
      while (*value_start == ' ')
        value_start++;

      char *value_end = strstr(value_start, "\r\n");
      if (value_end == NULL)
        value_end = search_end;
      size_t value_len = value_end - value_start;
      if (value_len > sizeof(g_header_value)) {
        value_len = sizeof(g_header_value) - 1;
      }
      strncpy(g_header_value, value_start, value_len);
      g_header_value[value_len] = '\0';
      return g_header_value;
    }
    buffer++;
  }

  return "";
}

void beaker_set_request_buffer(const char *buffer) {
  if (buffer == NULL)
    return;
  strncpy(current_request_buffer, buffer, BUFFER_SIZE - 1);
  current_request_buffer[BUFFER_SIZE - 1] = '\0';
}
