/*
 * usage example:
 * ./simpleserve 8080
 * ./simpleserve 8080 /home/bean
 *
 * 
 * build: gcc -O2 -Wall -o simpleserve simpleserve.c
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <ctype.h>
#include <stdint.h>
#include <limits.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <dirent.h>

#define MAX_REQUEST 16384

static void write_all(int fd, const void *buf, size_t len) {
    const char *p = buf;
    while (len > 0) {
        ssize_t w = write(fd, p, len);
        if (w < 0) {
            if (errno == EINTR) continue;
            return;
        }
        p += w;
        len -= (size_t)w;
    }
}

static void buf_append(char **buf, size_t *len, size_t *cap, const char *s) {
    size_t sl = strlen(s);
    if (*len + sl + 1 > *cap) {
        size_t ncap = (*cap == 0) ? (sl + 1) * 2 : (*cap + sl + 1) * 2;
        char *nbuf = realloc(*buf, ncap);
        if (!nbuf) { perror("realloc"); exit(1); }
        *buf = nbuf;
        *cap = ncap;
    }
    memcpy(*buf + *len, s, sl);
    *len += sl;
    (*buf)[*len] = '\0';
}

static int url_decode(const char *src, char *dst, size_t dst_size) {
    size_t j = 0;
    for (size_t i = 0; src[i] && j + 1 < dst_size; i++) {
        if (src[i] == '%') {
            if (!isxdigit((unsigned char)src[i+1]) || !isxdigit((unsigned char)src[i+2]))
                return 0;
            char hex[3] = { src[i+1], src[i+2], '\0' };
            dst[j++] = (char)strtol(hex, NULL, 16);
            i += 2;
        } else if (src[i] == '+') {
            dst[j++] = ' ';
        } else {
            dst[j++] = src[i];
        }
    }
    dst[j] = '\0';
    return 1;
}

static void url_encode(const char *src, char *dst, size_t dst_size) {
    static const char hex[] = "0123456789ABCDEF";
    size_t j = 0;
    for (size_t i = 0; src[i] && j + 3 < dst_size; i++) {
        unsigned char c = (unsigned char)src[i];
        if (isalnum(c) || c == '-' || c == '_' || c == '.' || c == '~') {
            dst[j++] = (char)c;
        } else {
            dst[j++] = '%';
            dst[j++] = hex[c >> 4];
            dst[j++] = hex[c & 0x0F];
        }
    }
    dst[j] = '\0';
}

static void html_escape(const char *src, char *dst, size_t dst_size) {
    size_t j = 0;
    for (size_t i = 0; src[i] && j + 6 < dst_size; i++) {
        switch (src[i]) {
            case '&':  memcpy(dst + j, "&amp;",  5); j += 5; break;
            case '<':  memcpy(dst + j, "&lt;",   4); j += 4; break;
            case '>':  memcpy(dst + j, "&gt;",   4); j += 4; break;
            case '"':  memcpy(dst + j, "&quot;", 6); j += 6; break;
            case '\'': memcpy(dst + j, "&#39;",  5); j += 5; break;
            default:   dst[j++] = src[i];
        }
    }
    dst[j] = '\0';
}

static int has_dotdot(const char *path) {
    const char *p = path;
    while (*p) {
        while (*p == '/') p++;
        const char *start = p;
        while (*p && *p != '/') p++;
        if (p - start == 2 && start[0] == '.' && start[1] == '.')
            return 1;
    }
    return 0;
}

static int resolve_path(const char *root, const char *url_path,
                        char *out, size_t out_size) {
    char root_real[PATH_MAX];
    if (!realpath(root, root_real)) return 0;
    size_t rlen = strlen(root_real);
    if (rlen > 1 && root_real[rlen - 1] == '/') root_real[--rlen] = '\0';

    char candidate[PATH_MAX];
    snprintf(candidate, sizeof(candidate), "%s%s", root_real, url_path);

    char resolved[PATH_MAX];
    if (!realpath(candidate, resolved)) return 0;

    if (strncmp(resolved, root_real, rlen) != 0) return 0;
    if (resolved[rlen] != '\0' && resolved[rlen] != '/') return 0;

    snprintf(out, out_size, "%s", resolved);
    return 1;
}

static int open_path(const char *root, const char *url_path) {
    int fd = open(root, O_RDONLY | O_DIRECTORY | O_CLOEXEC);
    if (fd < 0) return -1;

    const char *p = url_path;
    while (*p == '/') p++;

    while (*p) {
        const char *start = p;
        while (*p && *p != '/') p++;
        size_t n = (size_t)(p - start);

        if (n == 0) {
            while (*p == '/') p++;
            continue;
        }

        if (n >= NAME_MAX) {
            close(fd);
            return -1;
        }

        char name[NAME_MAX + 1];
        memcpy(name, start, n);
        name[n] = '\0';

        while (*p == '/') p++;

        int flags = O_RDONLY | O_CLOEXEC | O_NOFOLLOW;
        if (*p) flags |= O_DIRECTORY;

        int next = openat(fd, name, flags);
        if (next < 0) {
            close(fd);
            return -1;
        }

        close(fd);
        fd = next;
    }

    return fd;
}

static const char *mime_type(const char *path) {
    const char *dot = strrchr(path, '.');
    if (!dot) return "application/octet-stream";
    if (!strcasecmp(dot, ".html") || !strcasecmp(dot, ".htm")) return "text/html; charset=utf-8";
    if (!strcasecmp(dot, ".css"))  return "text/css; charset=utf-8";
    if (!strcasecmp(dot, ".js") || !strcasecmp(dot, ".mjs")) return "application/javascript; charset=utf-8";
    if (!strcasecmp(dot, ".json")) return "application/json";
    if (!strcasecmp(dot, ".png"))  return "image/png";
    if (!strcasecmp(dot, ".jpg") || !strcasecmp(dot, ".jpeg")) return "image/jpeg";
    if (!strcasecmp(dot, ".gif"))  return "image/gif";
    if (!strcasecmp(dot, ".svg"))  return "image/svg+xml";
    if (!strcasecmp(dot, ".ico"))  return "image/x-icon";
    if (!strcasecmp(dot, ".webp")) return "image/webp";
    if (!strcasecmp(dot, ".pdf"))  return "application/pdf";
    if (!strcasecmp(dot, ".txt"))  return "text/plain; charset=utf-8";
    if (!strcasecmp(dot, ".xml"))  return "application/xml";
    if (!strcasecmp(dot, ".zip"))  return "application/zip";
    if (!strcasecmp(dot, ".gz"))   return "application/gzip";
    if (!strcasecmp(dot, ".tar"))  return "application/x-tar";
    if (!strcasecmp(dot, ".mp3"))  return "audio/mpeg";
    if (!strcasecmp(dot, ".mp4"))  return "video/mp4";
    if (!strcasecmp(dot, ".webm")) return "video/webm";
    if (!strcasecmp(dot, ".woff"))  return "font/woff";
    if (!strcasecmp(dot, ".woff2")) return "font/woff2";
    if (!strcasecmp(dot, ".ttf"))   return "font/ttf";
    if (!strcasecmp(dot, ".wasm"))  return "application/wasm";
    return "application/octet-stream";
}

static void send_response(int fd, int status, const char *reason,
                          const char *ctype, long length) {
    char hdr[512];
    int n = snprintf(hdr, sizeof(hdr),
        "HTTP/1.1 %d %s\r\n"
        "Server: simpleserve/1.0\r\n"
        "Content-Type: %s\r\n"
        "Content-Length: %ld\r\n"
        "Connection: close\r\n"
        "\r\n",
        status, reason, ctype ? ctype : "text/html; charset=utf-8", length);
    write_all(fd, hdr, (size_t)n);
}

static void send_error(int fd, int status, const char *reason, const char *msg) {
    char body[1024];
    int n = snprintf(body, sizeof(body),
        "<!DOCTYPE html><html><head><title>%d %s</title></head>"
        "<body><h1>%d %s</h1><p>%s</p></body></html>\n",
        status, reason, status, reason, msg);
    send_response(fd, status, reason, "text/html; charset=utf-8", n);
    write_all(fd, body, (size_t)n);
}

static void send_file(int fd, int f, const char *path, int head_only) {
    struct stat st;
    if (fstat(f, &st) != 0 || !S_ISREG(st.st_mode)) {
        send_error(fd, 404, "not found", "file not found. perhaps you typed the url wrong?");
        return;
    }
    send_response(fd, 200, "OK", mime_type(path), (long)st.st_size);
    if (head_only) return;

    char buf[65536];
    ssize_t r;
    while ((r = read(f, buf, sizeof(buf))) > 0)
        write_all(fd, buf, (size_t)r);
}

static char *make_listing(int dirfd, const char *urlpath, size_t *out_len) {
    int scanfd = dup(dirfd);
    if (scanfd < 0) return NULL;

    DIR *dir = fdopendir(scanfd);
    if (!dir) {
        close(scanfd);
        return NULL;
    }

    char *html = NULL;
    size_t len = 0, cap = 0;

    char title[512];
    snprintf(title, sizeof(title), "index of %s", urlpath);
    buf_append(&html, &len, &cap, "<!DOCTYPE HTML PUBLIC \"-//W3C//DTD HTML 3.2 Final//EN\">\n");
    buf_append(&html, &len, &cap, "<html>\n<head>\n<title>");
    buf_append(&html, &len, &cap, title);
    buf_append(&html, &len, &cap, "</title>\n</head>\n<body>\n<h1>");
    buf_append(&html, &len, &cap, title);
    buf_append(&html, &len, &cap, "</h1>\n<hr>\n<pre>");

    if (strcmp(urlpath, "/") != 0)
        buf_append(&html, &len, &cap, "<a href=\"../\">../</a>\n");

    char base[4096];
    snprintf(base, sizeof(base), "%s%s",
             urlpath, (urlpath[strlen(urlpath) - 1] == '/') ? "" : "/");

    struct dirent *de;
    while ((de = readdir(dir)) != NULL) {
        if (!strcmp(de->d_name, ".") || !strcmp(de->d_name, ".."))
            continue;

        struct stat es;
        int is_dir = (fstatat(dirfd, de->d_name, &es, AT_SYMLINK_NOFOLLOW) == 0 &&
                      S_ISDIR(es.st_mode));

        char enc[2048], esc[2048];
        url_encode(de->d_name, enc, sizeof(enc));
        html_escape(de->d_name, esc, sizeof(esc));

        char line[8192];
        snprintf(line, sizeof(line),
                 "<a href=\"%s%s%s\">%s%s</a>\n",
                 base, enc, is_dir ? "/" : "", esc, is_dir ? "/" : "");
        buf_append(&html, &len, &cap, line);
    }
    closedir(dir);

    buf_append(&html, &len, &cap, "</pre>\n<hr>\n</body>\n</html>\n");
    *out_len = len;
    return html;
}

static void handle_client(int fd, const char *root) {
    char buf[MAX_REQUEST];
    size_t got = 0;

    while (got < sizeof(buf) - 1) {
        ssize_t r = read(fd, buf + got, sizeof(buf) - 1 - got);
        if (r <= 0) break;
        got += (size_t)r;
        buf[got] = '\0';
        if (strstr(buf, "\r\n\r\n") || strstr(buf, "\n\n")) break;
    }
    if (got == 0) { close(fd); return; }
    buf[got] = '\0';

    char method[16], raw_path[4096], version[16];
    if (sscanf(buf, "%15s %4095s %15s", method, raw_path, version) != 3) {
        send_error(fd, 400, "bad request", "malformed request line.");
        close(fd);
        return;
    }

    int head_only = 0;
    if (!strcmp(method, "GET"))        head_only = 0;
    else if (!strcmp(method, "HEAD"))  head_only = 1;
    else {
        send_error(fd, 501, "not implemented", "only get and head (heh, head) are supported.");
        close(fd);
        return;
    }

    char *q = strchr(raw_path, '?');
    if (q) *q = '\0';
    q = strchr(raw_path, '#');
    if (q) *q = '\0';

    char decoded[4096];
    if (!url_decode(raw_path, decoded, sizeof(decoded))) {
        send_error(fd, 400, "bad request", "invalid percent-encoding.");
        close(fd);
        return;
    }
    if (has_dotdot(decoded) || strchr(decoded, '\\') || decoded[0] != '/') {
        send_error(fd, 403, "forbidden", "you are not allowed to view this content.");
        close(fd);
        return;
    }

    int full = open_path(root, decoded);
    if (full < 0) {
        send_error(fd, 404, "not found", "the requested resource does not exist. perhaps you typed the url wrong?");
        close(fd);
        return;
    }

    struct stat st;
    if (fstat(full, &st) == 0 && S_ISDIR(st.st_mode)) {
        int index = openat(full, "index.html", O_RDONLY | O_CLOEXEC | O_NOFOLLOW);
        if (index >= 0) {
            send_file(fd, index, "index.html", head_only);
            close(index);
        } else {
            size_t llen = 0;
            char *listing = make_listing(full, decoded, &llen);
            if (listing) {
                send_response(fd, 200, "OK", "text/html; charset=utf-8", (long)llen);
                if (!head_only) write_all(fd, listing, llen);
                free(listing);
            } else {
                send_error(fd, 500, "internal server error", "could not read directory. check permissions?");
            }
        }
    } else if (fstat(full, &st) == 0 && S_ISREG(st.st_mode)) {
        send_file(fd, full, decoded, head_only);
    } else {
        send_error(fd, 404, "not found", "the requested resource does not exist. perhaps you may have typed the url wrong?");
    }
    close(full);
    close(fd);
}

int main(int argc, char **argv) {
    if (argc < 2 || argc > 3) {
        fprintf(stderr, "usage: %s <port> [root-directory]\n", argv[0]);
        return 1;
    }

    char *end = NULL;
    long port = strtol(argv[1], &end, 10);
    if (!end || *end != '\0' || port < 1 || port > 65535) {
        fprintf(stderr, "error: invalid port '%s'\n", argv[1]);
        return 1;
    }

    const char *root = (argc == 3) ? argv[2] : ".";
    char root_abs[PATH_MAX];
    if (!realpath(root, root_abs)) {
        fprintf(stderr, "error: cannot resolve root '%s': %s\n", root, strerror(errno));
        return 1;
    }
    struct stat st;
    if (stat(root_abs, &st) != 0 || !S_ISDIR(st.st_mode)) {
        fprintf(stderr, "error: '%s' is not a directory\n", root);
        return 1;
    }

    signal(SIGPIPE, SIG_IGN);
    signal(SIGCHLD, SIG_IGN);

    int lsock = socket(AF_INET, SOCK_STREAM, 0);
    if (lsock < 0) { perror("socket"); return 1; }

    int one = 1;
    setsockopt(lsock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    addr.sin_port = htons((uint16_t)port);

    if (bind(lsock, (struct sockaddr *)&addr, sizeof(addr)) < 0) { perror("bind"); return 1; }
    if (listen(lsock, 128) < 0) { perror("listen"); return 1; }

    printf("simpleserve serving %s on http://0.0.0.0:%ld/\n", root_abs, port);
    fflush(stdout);

    for (;;) {
        int cfd = accept(lsock, NULL, NULL);
        if (cfd < 0) {
            if (errno == EINTR) continue;
            perror("accept");
            continue;
        }
        pid_t pid = fork();
        if (pid == 0) {
            close(lsock);
            handle_client(cfd, root_abs);
            _exit(0);
        }
        close(cfd);
    }
}