#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <signal.h>
#include <sys/sendfile.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <unistd.h>

#include <algorithm>
#include <cctype>
#include <chrono>
#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <thread>

static std::string document_root = "src/frontend";

static const long DEFAULT_DOWNLOAD_MB = 30;
// Pretend-file ceiling: every served byte is synthesized from the repeating
// binary via sendfile, so advertising a huge Content-Length costs nothing.
static const long MAX_DOWNLOAD_MB = 1048576;

struct Mime {
    const char* file_extension;
    const char* content_type;
};

static const Mime MIME[] = {
    {".html", "text/html; charset=utf-8"},
    {".css", "text/css; charset=utf-8"},
    {".js", "text/javascript; charset=utf-8"},
    {".svg", "image/svg+xml"},
    {".json", "application/json"},
    {".bin", "application/octet-stream"},
    {".dat", "application/octet-stream"},
    {".woff2", "font/woff2"},
};

static const char* mime_for(const std::string& path) {
    for (const Mime& mime_entry : MIME)
        if (path.size() >= strlen(mime_entry.file_extension) &&
            path.compare(path.size() - strlen(mime_entry.file_extension), strlen(mime_entry.file_extension),
                         mime_entry.file_extension) == 0)
            return mime_entry.content_type;
    return "application/octet-stream";
}

static bool send_all(int destination_socket, const char* data, size_t length) {
    while (length > 0) {
        ssize_t bytes_sent = send(destination_socket, data, length, MSG_NOSIGNAL);
        if (bytes_sent <= 0) return false;
        data += bytes_sent;
        length -= (size_t)bytes_sent;
    }
    return true;
}

static bool send_file_range(int destination_socket, int source_file, off_t offset,
                            size_t remaining_bytes) {
    while (remaining_bytes > 0) {
        ssize_t bytes_sent =
            sendfile(destination_socket, source_file, &offset, remaining_bytes);
        if (bytes_sent <= 0) return false;
        remaining_bytes -= (size_t)bytes_sent;
    }
    return true;
}

// Streams `length` virtual bytes formed by the source file repeated
// back-to-back, starting at virtual byte `start`. The stream is larger
// than one file, so virtual offsets wrap across copies.
static bool send_repeated_file(int destination_socket, int source_file,
                               off_t source_size, off_t start, off_t length) {
    off_t position = start;
    while (length > 0) {
        off_t offset_in_copy = position % source_size;
        off_t span =
            std::min(source_size - offset_in_copy, length);
        if (!send_file_range(destination_socket, source_file,
                             offset_in_copy, (size_t)span))
            return false;
        position += span;
        length -= span;
    }
    return true;
}

static std::string executable_directory() {
    char executable_path_buffer[4096];
    ssize_t path_length =
        readlink("/proc/self/exe", executable_path_buffer,
                 sizeof(executable_path_buffer) - 1);
    if (path_length <= 0) return "";
    executable_path_buffer[path_length] = '\0';
    std::string executable_path(executable_path_buffer);
    size_t last_slash = executable_path.find_last_of('/');
    if (last_slash == std::string::npos) return "";
    return executable_path.substr(0, last_slash);
}

static bool directory_exists(const std::string& path) {
    struct stat path_status{};
    return stat(path.c_str(), &path_status) == 0 &&
           S_ISDIR(path_status.st_mode);
}

static std::string to_lower(std::string text) {
    std::transform(text.begin(), text.end(), text.begin(),
                   [](unsigned char c) { return (char)std::tolower(c); });
    return text;
}

static std::string header_value(const std::string& request,
                                const char* header_name) {
    std::string lowercased_name = to_lower(header_name);
    const std::string lowered_request = to_lower(request);
    size_t search_from = 0;
    while (true) {
        search_from = lowered_request.find(lowercased_name, search_from);
        if (search_from == std::string::npos) return "";
        bool at_line_start = (search_from == 0 || request[search_from - 1] == '\n');
        if (at_line_start) {
            size_t colon_pos = request.find(':', search_from);
            size_t line_break_pos = request.find("\r\n", search_from);
            if (colon_pos != std::string::npos &&
                line_break_pos != std::string::npos && colon_pos < line_break_pos) {
                size_t value_start = colon_pos + 1;
                while (value_start < line_break_pos &&
                       request[value_start] == ' ')
                    value_start++;
                return request.substr(value_start,
                                      line_break_pos - value_start);
            }
        }
        search_from += lowercased_name.size();
    }
}

static long query_param(const std::string& query, const char* name,
                        long fallback, long minimum, long maximum) {
    const std::string prefix = std::string(name) + "=";
    size_t search_from = 0;
    while (true) {
        search_from = query.find(prefix, search_from);
        if (search_from == std::string::npos) return fallback;
        bool at_param_start = (search_from == 0 || query[search_from - 1] == '&');
        if (at_param_start) {
            long value =
                strtol(query.c_str() + search_from + prefix.size(), nullptr, 10);
            return std::min(std::max(value, minimum), maximum);
        }
        search_from += prefix.size();
    }
}

struct Request {
    std::string method;
    std::string target;
    bool keep_alive = true;
    bool has_range = false;
    off_t range_start = 0;
    off_t range_end = -1;
    size_t content_length = 0;
    bool body_chunked = false;
};

static bool parse_request(const std::string& raw_request, Request& request) {
    size_t request_line_end = raw_request.find("\r\n");
    if (request_line_end == std::string::npos) return false;
    std::string request_line = raw_request.substr(0, request_line_end);

    size_t method_end = request_line.find(' ');
    size_t target_end = request_line.find(' ', method_end + 1);
    if (method_end == std::string::npos || target_end == std::string::npos)
        return false;
    request.method = request_line.substr(0, method_end);
    request.target = request_line.substr(method_end + 1, target_end - method_end - 1);

    std::string http_version = request_line.substr(target_end + 1);
    bool is_http_1_0 = (http_version.rfind("HTTP/1.0", 0) == 0);
    request.keep_alive = !is_http_1_0;

    std::string connection_header =
        to_lower(header_value(raw_request, "Connection:"));
    if (connection_header.find("close") != std::string::npos)
        request.keep_alive = false;
    else if (connection_header.find("keep-alive") != std::string::npos)
        request.keep_alive = true;

    std::string range_header = header_value(raw_request, "Range:");
    if (!range_header.empty() && range_header.rfind("bytes=", 0) == 0) {
        std::string byte_range_spec = range_header.substr(6);
        size_t dash_pos = byte_range_spec.find('-');
        if (dash_pos != std::string::npos) {
            std::string start_text = byte_range_spec.substr(0, dash_pos);
            std::string end_text = byte_range_spec.substr(dash_pos + 1);
            if (!start_text.empty()) {
                request.range_start = strtoll(start_text.c_str(), nullptr, 10);
                if (!end_text.empty())
                    request.range_end = strtoll(end_text.c_str(), nullptr, 10);
                request.has_range = true;
            }
        }
    }
    std::string content_length_header =
        header_value(raw_request, "Content-Length:");
    if (!content_length_header.empty())
        request.content_length = strtoull(content_length_header.c_str(), nullptr, 10);

    std::string transfer_encoding_header =
        to_lower(header_value(raw_request, "Transfer-Encoding:"));
    request.body_chunked =
        (transfer_encoding_header.find("chunked") != std::string::npos);

    return !request.method.empty() && !request.target.empty();
}

static void respond_error(int client_socket, int status_code,
                          const char* reason_phrase, bool include_body) {
    char error_body[128];
    int body_length = snprintf(error_body, sizeof(error_body), "<h1>%d %s</h1>",
                               status_code, reason_phrase);
    char response_head[256];
    int head_length = snprintf(response_head, sizeof(response_head),
                               "HTTP/1.1 %d %s\r\n"
                               "Content-Type: text/html; charset=utf-8\r\n"
                               "Content-Length: %d\r\n"
                               "Connection: close\r\n\r\n",
                               status_code, reason_phrase, body_length);
    send_all(client_socket, response_head, (size_t)head_length);
    if (include_body) send_all(client_socket, error_body, (size_t)body_length);
}

static void respond_range_not_satisfiable(int client_socket, off_t total_size) {
    char unsatisfiable_head[192];
    int unsatisfiable_head_length = snprintf(
        unsatisfiable_head, sizeof(unsatisfiable_head),
        "HTTP/1.1 416 Range Not Satisfiable\r\n"
        "Content-Range: bytes */%lld\r\n"
        "Content-Length: 0\r\n"
        "Connection: close\r\n\r\n",
        (long long)total_size);
    send_all(client_socket, unsatisfiable_head,
             (size_t)unsatisfiable_head_length);
}

static void handle_client(int client_socket) {
    int enabled = 1;
    setsockopt(client_socket, IPPROTO_TCP, TCP_NODELAY, &enabled, sizeof(enabled));
    timeval receive_timeout{10, 0};
    setsockopt(client_socket, SOL_SOCKET, SO_RCVTIMEO, &receive_timeout,
               sizeof(receive_timeout));

    while (true) {
        std::string raw_request;
        char read_buffer[4096];
        bool headers_complete = false;
        while (raw_request.size() < 8192) {
            ssize_t bytes_read = recv(client_socket, read_buffer, sizeof(read_buffer), 0);
            if (bytes_read <= 0) break;
            raw_request.append(read_buffer, (size_t)bytes_read);
            if (raw_request.find("\r\n\r\n") != std::string::npos) {
                headers_complete = true;
                break;
            }
        }
        if (!headers_complete) break;

        Request request;
        if (!parse_request(raw_request, request)) {
            respond_error(client_socket, 400, "Bad Request", true);
            break;
        }
        if (request.method != "GET" && request.method != "HEAD" &&
            request.method != "POST") {
            respond_error(client_socket, 405, "Method Not Allowed", true);
            break;
        }
        bool client_expects_body = (request.method == "GET");

        std::string url_path = request.target;
        std::string url_query;
        size_t query_start = url_path.find('?');
        if (query_start != std::string::npos) {
            url_query = url_path.substr(query_start + 1);
            url_path.resize(query_start);
        }
        if (url_path == "/") url_path = "/index.html";

        // Upload sink: swallows the request body and answers 204 so
        // upload-progress events map one-to-one onto bytes on the wire.
        bool is_upload_sink = (url_path == "/upload");

        if (request.method == "POST") {
            if (request.body_chunked) {
                respond_error(client_socket, 411, "Length Required", false);
                break;
            }

            size_t header_end = raw_request.find("\r\n\r\n") + 4;
            size_t body_bytes_already_read = raw_request.size() - header_end;
            size_t body_bytes_remaining = request.content_length;
            if (body_bytes_already_read >= body_bytes_remaining)
                body_bytes_remaining = 0;
            else
                body_bytes_remaining -= body_bytes_already_read;

            char discard_buffer[65536];
            bool upload_fully_received = true;
            while (body_bytes_remaining > 0) {
                size_t to_read =
                    std::min(body_bytes_remaining, sizeof(discard_buffer));
                ssize_t bytes_read =
                    recv(client_socket, discard_buffer, to_read, 0);
                if (bytes_read <= 0) {
                    upload_fully_received = false;
                    break;
                }
                body_bytes_remaining -= (size_t)bytes_read;
            }

            char redirect_head[192];
            int head_length;
            if (is_upload_sink) {
                head_length = snprintf(redirect_head, sizeof(redirect_head),
                                       "HTTP/1.1 204 No Content\r\n"
                                       "Access-Control-Allow-Origin: *\r\n"
                                       "Content-Length: 0\r\n"
                                       "Connection: %s\r\n\r\n",
                                       request.keep_alive ? "keep-alive" : "close");
            } else {
                head_length = snprintf(redirect_head, sizeof(redirect_head),
                                       "HTTP/1.1 303 See Other\r\n"
                                       "Location: /null\r\n"
                                       "Content-Length: 0\r\n"
                                       "Connection: %s\r\n\r\n",
                                       request.keep_alive ? "keep-alive" : "close");
            }
            send_all(client_socket, redirect_head, (size_t)head_length);
            if (!upload_fully_received || !request.keep_alive) break;
            continue;
        }

        if (url_path.find("..") != std::string::npos) {
            respond_error(client_socket, 400, "Bad Request", client_expects_body);
            break;
        }

        if (url_path == "/null") {
            char no_content_head[160];
            int head_length = snprintf(no_content_head, sizeof(no_content_head),
                                       "HTTP/1.1 204 No Content\r\n"
                                       "Access-Control-Allow-Origin: *\r\n"
                                       "Connection: %s\r\n\r\n",
                                       request.keep_alive ? "keep-alive" : "close");
            send_all(client_socket, no_content_head, (size_t)head_length);
            if (!request.keep_alive) break;
            continue;
        }

        // Parks the connection for `ms` milliseconds before replying.
        // Used by the client to detect how many connections the browser
        // will open in parallel to one origin: requests beyond that limit
        // queue instead of completing in the first wave.
        if (url_path == "/hold") {
            long delay_ms = query_param(url_query, "ms", 700, 1, 5000);
            std::this_thread::sleep_for(std::chrono::milliseconds(delay_ms));
            char hold_head[128];
            int head_length = snprintf(hold_head, sizeof(hold_head),
                                       "HTTP/1.1 204 No Content\r\n"
                                       "Connection: %s\r\n\r\n",
                                       request.keep_alive ? "keep-alive" : "close");
            send_all(client_socket, hold_head, (size_t)head_length);
            if (!request.keep_alive) break;
            continue;
        }

        if (url_path == "/download") {
            struct stat binary_status{};
            int self_binary = open("/proc/self/exe", O_RDONLY);
            if (self_binary < 0 || fstat(self_binary, &binary_status) != 0 ||
                binary_status.st_size <= 0) {
                if (self_binary >= 0) close(self_binary);
                respond_error(client_socket, 500, "Internal Server Error",
                              client_expects_body);
                break;
            }
            off_t binary_size = binary_status.st_size;

            // Exactly size_mb bytes: whole copies of the binary plus a
            // partial tail (or a truncated single copy if the binary alone
            // exceeds the request).
            long size_mb =
                query_param(url_query, "size", DEFAULT_DOWNLOAD_MB, 1,
                            MAX_DOWNLOAD_MB);
            off_t total_size = (off_t)size_mb * 1000 * 1000;

            off_t range_start_byte = 0;
            off_t range_end_byte = total_size - 1;
            int status_code = 200;
            if (request.has_range) {
                range_start_byte = request.range_start;
                range_end_byte =
                    (request.range_end >= 0 && request.range_end < total_size)
                        ? request.range_end
                        : total_size - 1;
                if (range_start_byte >= total_size ||
                    range_start_byte > range_end_byte) {
                    respond_range_not_satisfiable(client_socket, total_size);
                    close(self_binary);
                    break;
                }
                status_code = 206;
            }

            off_t content_length = range_end_byte - range_start_byte + 1;
            char download_head[512];
            int head_length;
            // Content-Encoding: identity is asserted explicitly: the payload
            // is raw bytes and must not be compressed in transit.
            if (status_code == 206) {
                head_length = snprintf(download_head, sizeof(download_head),
                                       "HTTP/1.1 206 Partial Content\r\n"
                                       "Content-Type: application/octet-stream\r\n"
                                       "Content-Length: %lld\r\n"
                                       "Content-Range: bytes %lld-%lld/%lld\r\n"
                                       "Accept-Ranges: bytes\r\n"
                                       "Content-Encoding: identity\r\n"
                                       "Cache-Control: no-store\r\n"
                                       "Access-Control-Allow-Origin: *\r\n"
                                       "Connection: %s\r\n\r\n",
                                       (long long)content_length,
                                       (long long)range_start_byte,
                                       (long long)range_end_byte,
                                       (long long)total_size,
                                       request.keep_alive ? "keep-alive" : "close");
            } else {
                head_length = snprintf(download_head, sizeof(download_head),
                                       "HTTP/1.1 200 OK\r\n"
                                       "Content-Type: application/octet-stream\r\n"
                                       "Content-Length: %lld\r\n"
                                       "Accept-Ranges: bytes\r\n"
                                       "Content-Encoding: identity\r\n"
                                       "Cache-Control: no-store\r\n"
                                       "Access-Control-Allow-Origin: *\r\n"
                                       "Connection: %s\r\n\r\n",
                                       (long long)content_length,
                                       request.keep_alive ? "keep-alive" : "close");
            }

            bool response_sent =
                send_all(client_socket, download_head, (size_t)head_length);
            if (response_sent && request.method == "GET")
                response_sent =
                    send_repeated_file(client_socket, self_binary, binary_size,
                                       range_start_byte, content_length);
            close(self_binary);
            if (!response_sent || !request.keep_alive) break;
            continue;
        }

        std::string filesystem_path = document_root + url_path;
        int opened_file = open(filesystem_path.c_str(), O_RDONLY);
        if (opened_file < 0) {
            respond_error(client_socket, 404, "Not Found", client_expects_body);
            break;
        }

        struct stat file_status{};
        if (fstat(opened_file, &file_status) != 0 ||
            S_ISDIR(file_status.st_mode)) {
            close(opened_file);
            respond_error(client_socket, 404, "Not Found", client_expects_body);
            break;
        }

        off_t file_size = file_status.st_size;
        off_t range_start_byte = 0;
        off_t range_end_byte = file_size - 1;
        int status_code = 200;

        if (request.has_range) {
            range_start_byte = request.range_start;
            range_end_byte =
                (request.range_end >= 0 && request.range_end < file_size)
                    ? request.range_end
                    : file_size - 1;
            if (range_start_byte >= file_size || range_start_byte > range_end_byte) {
                respond_range_not_satisfiable(client_socket, file_size);
                close(opened_file);
                break;
            }
            status_code = 206;
        }

        off_t content_length = range_end_byte - range_start_byte + 1;
        char response_head[512];
        int head_length;
        if (status_code == 206) {
            head_length = snprintf(response_head, sizeof(response_head),
                                   "HTTP/1.1 206 Partial Content\r\n"
                                   "Content-Type: %s\r\n"
                                   "Content-Length: %lld\r\n"
                                   "Content-Range: bytes %lld-%lld/%lld\r\n"
                                   "Accept-Ranges: bytes\r\n"
                                   "Cache-Control: no-cache\r\n"
                                   "Connection: %s\r\n\r\n",
                                    mime_for(url_path), (long long)content_length,
                                   (long long)range_start_byte,
                                   (long long)range_end_byte, (long long)file_size,
                                   request.keep_alive ? "keep-alive" : "close");
        } else {
            head_length = snprintf(response_head, sizeof(response_head),
                                   "HTTP/1.1 200 OK\r\n"
                                   "Content-Type: %s\r\n"
                                   "Content-Length: %lld\r\n"
                                   "Accept-Ranges: bytes\r\n"
                                   "Cache-Control: no-cache\r\n"
                                   "Connection: %s\r\n\r\n",
                                   mime_for(url_path), (long long)content_length,
                                   request.keep_alive ? "keep-alive" : "close");
        }

        bool response_sent = send_all(client_socket, response_head, (size_t)head_length);
        if (response_sent && request.method == "GET")
            response_sent = send_file_range(client_socket, opened_file, range_start_byte,
                                            (size_t)content_length);
        close(opened_file);
        if (!response_sent || !request.keep_alive) break;
    }
    close(client_socket);
}

int main(int argc, char** argv) {
    signal(SIGPIPE, SIG_IGN);

    int port = argc > 1 ? atoi(argv[1]) : 8080;
    if (argc > 2) {
        document_root = argv[2];
    } else if (!directory_exists(document_root)) {
        std::string exe_directory = executable_directory();
        std::string candidate = exe_directory + "/../src/frontend";
        if (directory_exists(candidate)) document_root = candidate;
    }
    if (!directory_exists(document_root)) {
        fprintf(stderr, "document root not found: %s\n", document_root.c_str());
        return 1;
    }

    int listening_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (listening_socket < 0) {
        perror("socket");
        return 1;
    }
    int enabled = 1;
    setsockopt(listening_socket, SOL_SOCKET, SO_REUSEADDR, &enabled, sizeof(enabled));

    sockaddr_in server_address{};
    server_address.sin_family = AF_INET;
    server_address.sin_addr.s_addr = INADDR_ANY;
    server_address.sin_port = htons((uint16_t)port);

    if (bind(listening_socket, (sockaddr*)&server_address, sizeof(server_address)) < 0) {
        perror("bind");
        return 1;
    }
    if (listen(listening_socket, 64) < 0) {
        perror("listen");
        return 1;
    }

    printf("OpenLightSpeedTest serving %s on http://0.0.0.0:%d\n", document_root.c_str(),
           port);

    while (true) {
        int client_socket = accept(listening_socket, nullptr, nullptr);
        if (client_socket < 0) continue;
        std::thread(handle_client, client_socket).detach();
    }
}