// SPDX-License-Identifier: AGPL-3.0-only
// SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

// The SQL helper: its statements, refusals, the URL, SCRAM, the protocol on a
// stand-in server (FakePg, below). Then what the helper prints over the same
// stand-in, line for line against the goldens (fixture/pipeline.json,
// `runs`), and the broker: the URL stays in the driver, a pass gets an
// address and a token.

#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
import std;
import Crafter.Review;
namespace fs = std::filesystem;
using namespace Crafter::Review;

namespace {
    std::int32_t Failures = 0;

    void Check(std::string_view name, bool cond, std::string_view detail = {}) {
        std::println("{}{}{}", cond ? "  ok   " : "  FAIL ", name, !cond && !detail.empty() ? std::format("  ({})", detail) : std::string());
        if (!cond) ++Failures;
    }

    std::string Be32(std::uint32_t v) {
        return {static_cast<char>(v >> 24), static_cast<char>((v >> 16) & 0xFF), static_cast<char>((v >> 8) & 0xFF), static_cast<char>(v & 0xFF)};
    }

    std::string Msg(char type, const std::string& body = {}) {
        return std::format("{}{}{}", type, Be32(static_cast<std::uint32_t>(body.size() + 4)), body);
    }

    std::string Err(const std::string& code, const std::string& text) {
        return Msg('E', std::format("SERROR{}C{}{}M{}{}{}", std::string(1, '\0'), code, std::string(1, '\0'), text, std::string(1, '\0'), std::string(1, '\0')));
    }

    // A Postgres stand-in on 127.0.0.1 for one connection: cleartext
    // password, simple Query, Parse/Bind/Describe/Execute/Sync. It records
    // what it is sent and answers as the server does where the helper relies
    // on it: two commands in one Parse are an error; COMMIT ends the
    // transaction (status I); a statement after an error fails.
    class FakePg {
    public:
        explicit FakePg(std::string password = "pw") : password(std::move(password)) {
            listener = socket(AF_INET, SOCK_STREAM, 0);
            sockaddr_in address {};
            address.sin_family = AF_INET;
            address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
            bind(listener, reinterpret_cast<sockaddr*>(&address), sizeof(address));
            listen(listener, 1);
            socklen_t length = sizeof(address);
            getsockname(listener, reinterpret_cast<sockaddr*>(&address), &length);
            port = ntohs(address.sin_port);
            thread = std::thread([this] { Serve(); });
        }

        ~FakePg() {
            Served();
        }

        // The log once the connection is over.
        std::vector<std::pair<std::string, std::string>> Served() {
            if (thread.joinable()) thread.join();
            return log;
        }

        std::uint16_t port = 0;

    private:
        std::string password;
        std::int32_t listener = -1;
        std::thread thread;
        std::vector<std::pair<std::string, std::string>> log;

        void Serve() {
            const std::int32_t c = accept(listener, nullptr, nullptr);
            std::string buffer;
            auto read = [&](std::size_t n) {
                while (buffer.size() < n) {
                    std::array<char, 4096> chunk {};
                    const ssize_t got = recv(c, chunk.data(), chunk.size(), 0);
                    if (got <= 0) break;
                    buffer.append(chunk.data(), static_cast<std::size_t>(got));
                }
                const std::string out = buffer.substr(0, std::min(n, buffer.size()));
                buffer.erase(0, out.size());
                return out;
            };
            auto length = [](const std::string& bytes) {
                return (static_cast<std::uint32_t>(static_cast<std::uint8_t>(bytes[0])) << 24) | (static_cast<std::uint32_t>(static_cast<std::uint8_t>(bytes[1])) << 16)
                       | (static_cast<std::uint32_t>(static_cast<std::uint8_t>(bytes[2])) << 8) | static_cast<std::uint8_t>(bytes[3]);
            };
            auto send = [&](const std::string& data) { ::send(c, data.data(), data.size(), MSG_NOSIGNAL); };
            const std::uint32_t startup = length(read(4));
            read(startup - 4);
            send(Msg('R', Be32(3)));
            read(1);
            std::string got = read(length(read(4)) - 4);
            while (!got.empty() && got.back() == '\0') got.pop_back();
            log.emplace_back("password", got);
            if (got != password) {
                send(Err("28P01", "password authentication failed"));
                close(c);
                close(listener);
                return;
            }
            send(Msg('S', std::string("server_version\x00" "17.6\x00", 20)) + Msg('R', Be32(0)) + Msg('Z', "I"));
            std::string status = "I";
            std::string pending;
            while (true) {
                const std::string type = read(1);
                if (type.empty()) break;
                const std::string body = read(length(read(4)) - 4);
                if (type == "Q") {
                    std::string q = body;
                    while (!q.empty() && q.back() == '\0') q.pop_back();
                    log.emplace_back("Q", q);
                    status = q.starts_with("BEGIN") ? "T" : (q == "ROLLBACK" ? "I" : status);
                    send(Msg('C', Dyn::Split(q)[0] + std::string(1, '\0')) + Msg('Z', status));
                } else if (type == "P") {
                    pending = body.substr(1, body.find('\0', 1) - 1);
                    log.emplace_back("P", pending);
                } else if (type == "S") {
                    const std::string low = Dyn::Lower(pending);
                    std::string out;
                    if (status == "E") {
                        out = Err("25P02", "current transaction is aborted");
                    } else if (pending.contains(';')) {
                        out = Err("42601", "cannot insert multiple commands into a prepared statement");
                        status = "E";
                    } else if (low.starts_with("commit")) {
                        out = Msg('C', std::string("COMMIT\0", 7));
                        status = "I";
                    } else if (low.starts_with("select")) {
                        out = Msg('T', std::string("\x00\x01x\x00", 4) + std::string(18, '\0')) + Msg('D', std::string("\x00\x01\x00\x00\x00\x01" "1", 7))
                              + Msg('N', std::string("SNOTICE\0Mhello\0\0", 16)) + Msg('C', std::string("SELECT 1\0", 9));
                    } else if (low.starts_with("boom")) {
                        out = Err("42601", "syntax error at or near \"boom\"");
                        status = "E";
                    } else {
                        out = Msg('C', Dyn::Upper(Dyn::Split(pending)[0]) + std::string(1, '\0'));
                    }
                    send(Msg('1') + Msg('2') + out + Msg('Z', status));
                } else if (type == "X") {
                    log.emplace_back("X", "");
                    break;
                }
            }
            close(c);
            close(listener);
        }
    };

    std::uint16_t DeadPort() {
        const std::int32_t s = socket(AF_INET, SOCK_STREAM, 0);
        sockaddr_in address {};
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        bind(s, reinterpret_cast<sockaddr*>(&address), sizeof(address));
        socklen_t length = sizeof(address);
        getsockname(s, reinterpret_cast<sockaddr*>(&address), &length);
        close(s);
        return ntohs(address.sin_port);
    }

    std::vector<std::string> Of(const std::vector<std::pair<std::string, std::string>>& log, std::string_view kind) {
        std::vector<std::string> out;
        for (const auto& [k, x] : log) {
            if (k == kind) out.push_back(x);
        }
        return out;
    }

    bool Any(const std::vector<std::string>& lines, std::string_view part) {
        return std::ranges::any_of(lines, [&](const std::string& line) { return line.contains(part); });
    }

    std::string Show(const std::vector<std::string>& lines) {
        Json out = Json::array();
        for (const std::string& line : lines) out.push_back(line);
        return Dyn::ReprValue(out);
    }

    // A raw call of the broker: one line in, every line back.
    std::vector<std::string> RawCall(std::uint16_t port, const std::string& line) {
        const std::int32_t s = socket(AF_INET, SOCK_STREAM, 0);
        sockaddr_in address {};
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        address.sin_port = htons(port);
        connect(s, reinterpret_cast<sockaddr*>(&address), sizeof(address));
        ::send(s, line.data(), line.size(), MSG_NOSIGNAL);
        shutdown(s, SHUT_WR);
        std::string all;
        std::array<char, 4096> chunk {};
        ssize_t got = 0;
        while ((got = recv(s, chunk.data(), chunk.size(), 0)) > 0) all.append(chunk.data(), static_cast<std::size_t>(got));
        close(s);
        return Dyn::Split(Dyn::RStrip(all, "\n"), "\n");
    }
}

int main() {
    std::println("the SQL helper (pgwire): statements, refusals, the URL, SCRAM, the protocol on a stand-in server");
    using Pieces = std::vector<std::string>;
    auto sp = [](std::string_view sql) { return PgWire::SplitStatements(sql); };
    Check("split: semicolons in quotes, E-strings, dollar quotes (tagged, and $ inside identifiers), comments (nested, -- to \\r), double quotes and parentheses do not cut; empty statements drop",
          sp("select ';' ; select 2;") == Pieces{"select ';'", "select 2"} && sp("select $x$ a; $y$ ; $x$; commit") == Pieces{"select $x$ a; $y$ ; $x$", "commit"}
          && sp("select 1 as a$b$; commit; select $b$") == Pieces{"select 1 as a$b$", "commit", "select $b$"} && sp("select E'a\\'; commit; '") == Pieces{"select E'a\\'; commit; '"}
          && sp("-- c;\nselect 1; /* x; /* y; */ z; */ select 2") == Pieces{"-- c;\nselect 1", "/* x; /* y; */ z; */ select 2"} && sp("select 1 --x\r; commit") == Pieces{"select 1 --x", "commit"}
          && sp("create rule r as on insert to t do also (insert into a values (1); update b set c = 1); select 3")
                 == Pieces{"create rule r as on insert to t do also (insert into a values (1); update b set c = 1)", "select 3"}
          && sp("select \"a;b\" from t; select 2") == Pieces{"select \"a;b\" from t", "select 2"} && sp("  ;; -- only\n ; ").empty());
    const std::vector<std::string> refuse = {"commit", "COMMIT AND CHAIN", "end", "END WORK", "rollback", "rollback work", "abort", "begin", "start transaction", "prepare transaction 'x'",
                                             "commit prepared 'x'", "rollback prepared 'x'", "/* c */ commit", "-- c\ncommit", "--c\rcommit", "(commit)", "copy t to stdout", "COPY t FROM PROGRAM 'x'"};
    const std::vector<std::string> allow = {"rollback to s1", "ROLLBACK TO SAVEPOINT s", "savepoint s", "release s", "prepare p as select 1", "select 1", "with x as (select 1) select * from x",
                                            "do $$ begin commit; end $$", "call p()", "endless_fn()", "end1()"};
    std::vector<std::string> wrong;
    for (const std::string& s : refuse) {
        if (!PgWire::Refused(s)) wrong.push_back(s);
    }
    for (const std::string& s : allow) {
        if (PgWire::Refused(s)) wrong.push_back(s);
    }
    Check("refused: every statement that ends, restarts or prepares the transaction, and COPY - past comments and parentheses; ROLLBACK TO, SAVEPOINT, DO and CALL go to the server (which refuses COMMIT inside them in a transaction block)",
          wrong.empty(), Show(wrong));
    Check("URL: this machine only; the password masked for a log or a dump", !PgWire::CheckUrl("postgresql://postgres:postgres@127.0.0.1:54322/postgres") && !PgWire::CheckUrl("postgres://u@localhost/x")
          && !PgWire::CheckUrl("postgresql://u@[::1]:5/x") && PgWire::CheckUrl("postgresql://u:p@db.abcdefgh.supabase.co:5432/postgres").value_or("").contains("not this machine")
          && PgWire::CheckUrl("http://127.0.0.1:54321") && PgWire::CheckUrl("") && PgWire::CheckUrl(nullptr)
          && PgWire::Redact("postgresql://postgres:hunter2@127.0.0.1:54322/postgres") == "postgresql://postgres:***@127.0.0.1:54322/postgres"
          && PgWire::Redact("postgresql://127.0.0.1/x") == "postgresql://127.0.0.1/x");
    PgWire::Scram sc("pencil", "user", "rOprNGfwEbeRWgbNEkqO");
    const std::string sf = "r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,s=W22ZaJ0SNY7soEsUEjb6gQ==,i=4096";
    Check("SCRAM-SHA-256: the RFC 7677 test vector - the client-first, the client proof, the server signature verified (and a wrong one refused)",
          sc.ClientFirst() == "n,,n=user,r=rOprNGfwEbeRWgbNEkqO"
          && sc.ClientFinal(sf) == "c=biws,r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,p=dHzbZapWIk4jUhN+Ute9ytag9zjfMHgsqmmiz7AndVQ="
          && sc.Verify("v=6rriTRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4=") && !sc.Verify("v=AAAATRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4="));
    std::vector<std::string> out;
    const PgWire::Out collect = [&](const std::string& line) { out.push_back(line); };
    {
        FakePg pg;
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), "select 1; insert into t values ('a'); boom; select 3", collect);
        const auto log = pg.Served();
        const std::vector<std::string> queries = Of(log, "Q");
        const std::vector<std::string> setup(queries.begin(), queries.begin() + std::min<std::ptrdiff_t>(4, static_cast<std::ptrdiff_t>(queries.size())));
        Check("run: the password, BEGIN and the three SET LOCALs, then one Parse per statement up to the failing one, then ROLLBACK and Terminate; the call says it connected, prints rows, the notice, the error and that it stopped; exit 1",
              !log.empty() && log[0] == std::pair<std::string, std::string>{"password", "pw"}
              && setup == std::vector<std::string>{"BEGIN", "SET LOCAL statement_timeout = '20s'", "SET LOCAL lock_timeout = '3s'", "SET LOCAL idle_in_transaction_session_timeout = '60s'"}
              && Of(log, "P") == std::vector<std::string>{"select 1", "insert into t values ('a')", "boom"} && log.size() >= 2 && log[log.size() - 2] == std::pair<std::string, std::string>{"Q", "ROLLBACK"}
              && log.back() == std::pair<std::string, std::string>{"X", ""} && code == 1 && out.at(0).starts_with(PgWire::Answered) && std::ranges::contains(out, std::string("NOTICE: hello"))
              && Any(out, "syntax error") && Any(out, "statement 3 failed") && out.back() == "-- rolled back: nothing of this call was committed", Show(out));
    }
    {
        FakePg pg;
        std::unique_ptr<PgWire::Connection> cn = PgWire::Connect(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), 1);
        cn->Simple("BEGIN");
        const PgWire::Answer two = cn->Extended("select 1; select 2");
        cn->Simple("ROLLBACK");
        cn->Close();
        const auto log = pg.Served();
        Check("extended: a whole script in one Parse comes back as the server refuses it (two commands), the transaction failed - never run", two.error && two.error->at('C') == "42601"
              && two.rows.empty() && std::ranges::contains(log, std::pair<std::string, std::string>{"P", "select 1; select 2"}));
    }
    {
        out.clear();
        FakePg pg;
        // A statement the refusal list misses: the transaction status is the
        // backstop.
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), "create table x (a int); commit; select 1", collect, {},
                                                 [](std::string_view) { return std::optional<std::string>(); });
        const auto log = pg.Served();
        Check("run: the backstop - after a statement that ended the transaction nothing more is sent, the call fails and says what may have been committed",
              Of(log, "P") == std::vector<std::string>{"create table x (a int)", "commit"} && code == 1 && Any(out, "STOPPED") && out.back().contains("may have been committed"), Show(out));
    }
    {
        out.clear();
        FakePg pg;
        const std::string url = std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port);
        const std::int32_t code = PgWire::RunSql(url, "select 1; commit", collect);
        PgWire::Connect(url, 1)->Close();
        const auto log = pg.Served();
        Check("run: a refused statement stops the call before it connects - nothing is sent, no connected line", code == 1 && out.at(0).starts_with("review-gates sql: refused, nothing ran")
              && out.at(0).contains("statement 2: COMMIT") && !Any(out, PgWire::Answered) && Of(log, "Q").empty() && Of(log, "P").empty());
    }
    {
        out.clear();
        const std::uint16_t dead = DeadPort();
        const PgWire::Connector once = [](std::string_view url) { return PgWire::Connect(url, 1, [](double) {}); };
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", dead), "select 1", collect, once);
        std::vector<std::string> ignored;
        const PgWire::Out drop = [&](const std::string& line) { ignored.push_back(line); };
        Check("run: nothing listening -> NO DATABASE, exit 2, no connected line; a hosted URL never connects; no statement is a usage error", code == 2 && out.at(0).starts_with(PgWire::NoDatabase)
              && PgWire::RunSql("postgresql://u@db.x.supabase.co/x", "select 1", drop) == 2 && PgWire::RunSql(std::format("postgresql://u@127.0.0.1:{}/x", dead), " -- nothing\n", drop) == 64);
        Check("...and the reason: cannot connect to 127.0.0.1:<port> ([Errno 111] Connection refused)",
              out.at(0) == std::format("{} - cannot connect to 127.0.0.1:{} ([Errno 111] Connection refused)", PgWire::NoDatabase, dead), out.at(0));
    }
    {
        out.clear();
        FakePg pg("other");
        const PgWire::Connector once = [](std::string_view url) { return PgWire::Connect(url, 1); };
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), "select 1", collect, once);
        const std::string detail = out.at(0).substr(out.at(0).find(" - ") + 3);
        Check("run: a refused login is NO DATABASE (exit 2), its message without the password", code == 2 && out.at(0).starts_with(PgWire::NoDatabase) && out.at(0).contains("28P01")
              && !detail.contains("pw"), Show(out));
    }
    {
        // A busy server is asked again, but not slept on after the last try.
        std::vector<double> slept;
        try {
            PgWire::Connect(std::format("postgresql://u@127.0.0.1:{}/x", DeadPort()), 3, [&](double s) { slept.push_back(s); });
        } catch (const PgWire::DbError&) {
            // Expected: nothing listens.
        }
        Check("connect: an unreachable server is tried three times, slept on between tries only", slept == std::vector<double>{0.5, 1.0});
    }

    // What the helper prints over the stand-in, line for line against the
    // goldens.
    const Json goldens = Dyn::LoadJsonFile(fs::current_path() / "tests" / "Goldens" / "fixture" / "pipeline.json");
    std::int32_t compared = 0;
    for (const Json& entry : goldens.at("pgwire").at("runs")) {
        if (entry.contains("noServer")) continue;
        out.clear();
        FakePg pg;
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), entry.at("sql").get<std::string>(), collect);
        pg.Served();
        std::vector<std::string> normalized;
        for (const std::string& line : out) normalized.push_back(Dyn::Replace(line, std::format(":{}", pg.port), ":PORT"));
        std::vector<std::string> want;
        for (const Json& line : entry.at("out")) want.push_back(line.get<std::string>());
        if (code == entry.at("code").get<std::int32_t>() && normalized == want) ++compared;
        else Check(std::format("run over the stand-in, as the goldens record it: {}", entry.at("sql").get<std::string>()), false, Show(normalized));
    }
    Check(std::format("run over the stand-in prints what the goldens record, line for line ({} scripts)", compared), compared == 6);

    // The broker.
    std::vector<std::string> ran;
    PgWire::Broker broker(
        [&](const std::string& url, const std::string& sql, const PgWire::Out& say) {
            ran.push_back(std::format("{}|{}", url, sql));
            say(std::format("{} to x", PgWire::Answered));
            say("é line");
            return 1;
        },
        [](const std::string& url) { return std::pair<bool, std::string>{url.contains("good"), "detail"}; });
    broker.Start();
    const std::string url = "postgresql://postgres:hunter2@127.0.0.1:54322/postgres";
    const std::string address = broker.Address(url);
    const std::string token = address.substr(address.rfind('#') + 1);
    Check("broker: the address is 127.0.0.1:<port>#<32 hex>, one token per URL, and the token names its URL", address.starts_with("127.0.0.1:") && token.size() == 32
          && std::ranges::all_of(token, [](char c) { return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f'); }) && broker.Address(url) == address
          && broker.Address("postgresql://other@127.0.0.1/x") != address && broker.UrlOf(token) == url && broker.Address("").empty() && !address.contains("hunter2"));
    out.clear();
    std::int32_t code = PgWire::CallBroker(address, "select 1;\nselect 'é';", false, collect);
    Check("broker: a call runs the SQL on its token's URL, streams the lines back and returns the call's exit code", code == 1 && ran.size() == 1 && ran[0] == std::format("{}|select 1;\nselect 'é';", url)
          && out == std::vector<std::string>{std::format("{} to x", PgWire::Answered), "é line"}, Show(out));
    out.clear();
    code = PgWire::CallBroker(broker.Address("postgresql://good@127.0.0.1/x"), "", true, collect);
    Check("broker: a probe answers with the connected line and exit 0", code == 0 && out == std::vector<std::string>{std::format("{} - detail", PgWire::Answered)}, Show(out));
    out.clear();
    code = PgWire::CallBroker(std::format("{}#{}", address.substr(0, address.rfind('#')), std::string(32, '0')), "select 1", false, collect);
    Check("broker: an unknown token is NO DATABASE, exit 2, and nothing runs", code == 2 && out.size() == 1 && out[0].starts_with(PgWire::NoDatabase) && ran.size() == 1, Show(out));
    const std::uint16_t brokerPort = static_cast<std::uint16_t>(std::stoi(address.substr(10, address.rfind('#') - 10)));
    const std::vector<std::string> notWhole = RawCall(brokerPort, "{\"token\": \"x\"");
    Check("broker: a call that never ends its line did not arrive whole - nothing runs, exit 64", notWhole.size() == 2 && notWhole[0].contains("did not arrive whole")
          && notWhole[1] == std::format("{}64", PgWire::ExitMark), Show(notWhole));
    const std::vector<std::string> listToken = RawCall(brokerPort, "{\"token\": [1], \"sql\": \"select 1\"}\n");
    Check("broker: a token that is no string gets NO DATABASE and the exit mark", listToken.size() == 2
          && listToken[0].starts_with(PgWire::NoDatabase) && listToken[1] == std::format("{}2", PgWire::ExitMark), Show(listToken));
    out.clear();
    Check("call_broker: no address, a bad port, a closed port - NO DATABASE, exit 2, with the reason",
          PgWire::CallBroker("", "x", false, collect) == 2 && PgWire::CallBroker("127.0.0.1:x#t", "x", false, collect) == 2
          && PgWire::CallBroker(std::format("127.0.0.1:{}#t", DeadPort()), "x", false, collect) == 2 && out.size() == 3
          && out[0] == std::format("{} - this pass has no database (the driver gives one to a pass with the supabase environment)", PgWire::NoDatabase)
          && out[1] == std::format("{} - the driver's SQL broker does not answer (invalid literal for int() with base 10: 'x')", PgWire::NoDatabase)
          && out[2].ends_with("does not answer ([Errno 111] Connection refused)"), Show(out));
    broker.Close();
    out.clear();
    Check("broker: closed, its port answers no more", PgWire::CallBroker(address, "select 1", false, collect) == 2 && out.at(0).ends_with("([Errno 111] Connection refused)"), Show(out));
    return Failures ? 1 : 0;
}