#include "src/lib/types.cpp"
#include "src/lib/dvalue.cpp"
#include "src/lib/functionlib.cpp"
#include "src/lib/fcgi_forward.h"
#include <sys/wait.h>
#include <cassert>
#include <chrono>

static String record(unsigned char type, const String& content, int padding = 0)
{
    String out;
    unsigned char header[8] = {1, type, 0, 1, (unsigned char)(content.size() >> 8), (unsigned char)content.size(), (unsigned char)padding, 0};
    out.append((const char*)header, 8);
    return out + content + String(padding, 'p');
}

static void send_all(int fd, const String& bytes, bool fragment = false)
{
    for(size_t pos = 0; pos < bytes.size();)
    {
        ssize_t n = send(fd, bytes.data() + pos, fragment ? 1 : bytes.size() - pos, MSG_NOSIGNAL);
        if(n <= 0) return;
        pos += n;
    }
}

static FcgiRequestResult run(int family, int mode, u64 limit = 8 * 1024 * 1024, bool adapter = false)
{
    FcgiRequestOptions options;
    options.timeout_ms = mode == 7 || mode == 9 ? 150 : 2000;
    options.max_response_bytes = limit;
    int listener = socket(family, SOCK_STREAM, 0);
    assert(listener >= 0);
    String path = "/tmp/bearer-fcgi-native-" + std::to_string(getpid()) + ".sock";
    sockaddr_storage address{};
    socklen_t length;
    if(family == AF_UNIX)
    {
        options.socket_path = path;
        auto& a = (sockaddr_un&)address;
        a.sun_family = family;
        memcpy(a.sun_path, path.c_str(), path.size() + 1);
        length = sizeof(a);
        unlink(path.c_str());
    }
    else if(family == AF_INET)
    {
        options.host = "127.0.0.1";
        auto& a = (sockaddr_in&)address;
        a.sin_family = family;
        inet_pton(family, options.host.c_str(), &a.sin_addr);
        length = sizeof(a);
    }
    else
    {
        options.host = "::1";
        auto& a = (sockaddr_in6&)address;
        a.sin6_family = family;
        inet_pton(family, options.host.c_str(), &a.sin6_addr);
        length = sizeof(a);
    }
    if(bind(listener, (sockaddr*)&address, length) != 0)
    {
        assert(family == AF_INET6 && (errno == EADDRNOTAVAIL || errno == EAFNOSUPPORT));
        close(listener);
        return {};
    }
    assert(listen(listener, 2) == 0);
    assert(getsockname(listener, (sockaddr*)&address, &length) == 0);
    if(family == AF_INET) options.port = ntohs(((sockaddr_in&)address).sin_port);
    if(family == AF_INET6) options.port = ntohs(((sockaddr_in6&)address).sin6_port);
    StringMap env;
    env[String(140, 'K')] = String(70000, 'V');
    String body(mode == 9 ? 4 * 1024 * 1024 : 300000, 'b'); body[17] = 0;
    String expected = fcgi_build_request(env, body);
    pid_t child = fork();
    assert(child >= 0);
    if(child == 0)
    {
        alarm(4);
        int fd = accept(listener, nullptr, nullptr);
        assert(fd >= 0);
        if(mode == 9)
        {
            usleep(300000);
            close(fd); close(listener); _exit(0);
        }
        if(mode == 8)
        {
            // Fill the output before consuming the large request.
            String output;
            fcgi_forward_stream(output, 6, String(300000, 'o'));
            send_all(fd, output);
        }
        String input;
        char buffer[65536];
        while(input.size() < expected.size())
        {
            ssize_t n = read(fd, buffer, sizeof(buffer));
            if(n <= 0) _exit(2);
            input.append(buffer, n);
        }
        assert(input == expected);
        String end(8, 0);
        if(mode == 1) { end[3] = 42; end[4] = 2; }
        String output = record(6, String("X: a\r\nX: b\r\n\r\na\0b", 17), 7) + record(7, "diagnostic", 3);
        if(mode == 8) output = record(7, "diagnostic");
        if(mode == 13) output = record(6, "Status: 201 Created\r\nContent-Type: text/plain\r\n\r\ncreated");
        if(mode == 14) output = record(6, "unframed worker output");
        if(mode == 2) output[0] = 2;
        if(mode == 3) output[3] = 2;
        if(mode == 10) output = record(6, "") + record(6, "after close");
        if(mode == 11) { output[6] = 255; output.resize(12); }
        if(mode == 12) output[1] = 4;
        if(mode == 4) output += record(3, String(7, 0));
        else if(mode == 5) { end[4] = 4; output += record(3, end); }
        else if(mode == 6) output += record(3, end).substr(0, 12);
        else if(mode == 7)
        {
            for(int i = 0; i < 500; ++i)
            {
                send_all(fd, record(7, "x"));
                usleep(1000);
            }
        }
        else output += record(3, end);
        if(mode != 7) send_all(fd, output, mode == 0);
        // A second connection would repeat a side effect. Reject retries.
        pollfd item{listener, POLLIN, 0};
        assert(poll(&item, 1, 30) == 0);
        close(fd); close(listener); _exit(0);
    }
    close(listener);
    auto start = std::chrono::steady_clock::now();
    FcgiRequestResult result;
    if(adapter)
    {
        FcgiForwardResult forwarded = fcgi_forward_request(path, env, body, 2, limit);
        assert(forwarded.ok == (mode == 0 || mode == 13 || mode == 14));
        if(mode == 0)
        {
            assert(forwarded.status == 200 && forwarded.headers["X"] == "b");
            assert(forwarded.body == String("a\0b", 3));
        }
        if(mode == 13)
        {
            assert(forwarded.status == 201 && forwarded.headers["Content-Type"] == "text/plain");
            assert(forwarded.body == "created");
        }
        if(mode == 14)
        {
            assert(forwarded.status == 200 && forwarded.headers.empty());
            assert(forwarded.body == "unframed worker output");
        }
    }
    else result = fcgi_request(options, env, body);
    auto elapsed = std::chrono::steady_clock::now() - start;
    assert(elapsed < std::chrono::seconds(2));
    int status; assert(waitpid(child, &status, 0) == child);
    assert(WIFEXITED(status) && WEXITSTATUS(status) == 0);
    unlink(path.c_str());
    if(adapter) return result;
    if(mode == 0 || mode == 1)
    {
        if(limit < 27) assert(!result.ok && result.error_code == "output_limit");
        else
        {
            assert(result.ok && result.stderr_data == "diagnostic");
            assert(result.stdout_data == String("X: a\r\nX: b\r\n\r\na\0b", 17));
            assert(result.app_status == (mode == 1 ? 42u : 0u));
            assert(result.protocol_status == (mode == 1 ? 2 : 0));
        }
    }
    else if(mode == 8) assert(result.ok && result.stdout_data == String(300000, 'o'));
    else assert(!result.ok);
    if(mode >= 2 && mode <= 5) assert(result.error_code == "protocol");
    if(mode == 6 || mode == 11) assert(result.error_code == "truncated");
    if(mode == 7 || mode == 9) assert(result.error_code == "timeout");
    if(mode == 10 || mode == 12) assert(result.error_code == "protocol");
    return result;
}

int main()
{
    signal(SIGPIPE, SIG_IGN);
    auto handshake = [](String path, String cookie, String origin) {
        StringMap params = split_http_headers("GET " + path + " HTTP/1.1\r\n"
            "Host: socket.example.test\r\nCookie: discarded=1\r\ncOoKiE: " + cookie + "\r\n"
            "Origin: https://discarded.example.test\r\noRiGiN: " + origin + "\r\n"
            "Authorization: Bearer test-token\r\nX-Forwarded-Proto: https\r\n"
            "Bearer-CLI: 1\r\nContent-Length: 4\r\nContent-Type: text/plain\r\n\r\n");
        params["SCRIPT_FILENAME"] = "/srv/site" + params["DOCUMENT_URI"];
        params["DOCUMENT_ROOT"] = "/srv/site";
        params["REMOTE_ADDR"] = "192.0.2.1";
        params["HTTPS"] = "on";
        params["BEARER_CLI"] = "1";
        params["BEARER_WS_DISPATCH"] = "1";
        params["BEARER_WS_STATE"] = "stale-worker-state";
        params["ROUTE_PAGE"] = "stale-route";
        return(params);
    };
    const StringMap alice = handshake("/alice.capy?room=one&%2Fprivate", "session=alice; csrf=a", "https://alice.example.test");
    const StringMap bob = handshake("/bob.capy?room=two", "session=bob; csrf=b", "https://bob.example.test");
    StringMap alice_frame = fcgi_websocket_request_context(alice);
    StringMap bob_frame = fcgi_websocket_request_context(bob);
    assert(alice_frame.at("HTTP_COOKIE") == "session=alice; csrf=a");
    assert(alice_frame.at("HTTP_ORIGIN") == "https://alice.example.test");
    assert(bob_frame.at("HTTP_COOKIE") == "session=bob; csrf=b");
    assert(bob_frame.at("HTTP_ORIGIN") == "https://bob.example.test");
    for(const StringMap* source : {&alice, &bob})
    {
        StringMap frame = fcgi_websocket_request_context(*source);
        for(const String key : {"REQUEST_URI", "DOCUMENT_URI", "SCRIPT_FILENAME", "QUERY_STRING", "DOCUMENT_ROOT", "REMOTE_ADDR", "HTTPS", "HTTP_AUTHORIZATION", "HTTP_X_FORWARDED_PROTO"})
            assert(frame.at(key) == source->at(key));
        for(const String key : {"CONTENT_LENGTH", "CONTENT_TYPE", "REQUEST_METHOD", "BEARER_CLI", "BEARER_WS_DISPATCH", "BEARER_WS_STATE", "ROUTE_PAGE"})
            assert(!frame.contains(key));
        assert(frame.at("HTTP_BEARER_CLI") == "1");
    }
    alice_frame["HTTP_COOKIE"] = "session=forged";
    alice_frame["QUERY_STRING"] = "room=other";
    assert(fcgi_websocket_request_context(alice).at("HTTP_COOKIE") == "session=alice; csrf=a");
    assert(fcgi_websocket_request_context(alice).at("QUERY_STRING") == "room=one&%2Fprivate");
    assert(bob_frame.at("HTTP_COOKIE") == "session=bob; csrf=b");
    assert(fcgi_websocket_request_context(split_http_headers("GET /anonymous.capy HTTP/1.1\r\nHost: socket.example.test\r\n\r\n")).count("HTTP_COOKIE") == 0);
    String lengths;
    fcgi_forward_put_len(lengths, 127);
    fcgi_forward_put_len(lengths, 128);
    fcgi_forward_put_len(lengths, 70000);
    assert(lengths == String("\x7f\x80\x00\x00\x80\x80\x01\x11\x70", 9));
    for(int mode = 0; mode <= 12; ++mode) run(AF_UNIX, mode);
    run(AF_UNIX, 0, 25);
    run(AF_UNIX, 0, 1024, true);
    run(AF_UNIX, 1, 1024, true);
    run(AF_UNIX, 13, 1024, true);
    run(AF_UNIX, 14, 1024, true);
    run(AF_INET, 0);
    run(AF_INET6, 0);
    FcgiRequestOptions options;
    assert(fcgi_request(options, {}, "").error_code == "invalid_options");
    options.socket_path = "/unused"; options.host = "127.0.0.1";
    assert(fcgi_request(options, {}, "").error_code == "invalid_options");
    options.socket_path = ""; options.host = "localhost"; options.port = 9000;
    assert(fcgi_request(options, {}, "").error_code == "invalid_options");
    options.host = "127.0.0.1"; options.timeout_ms = 0;
    assert(fcgi_request(options, {}, "").error_code == "invalid_options");
    return 0;
}