// Compiles Packwand's Luau stdlib into a C++ source that embeds each module's bytecode (packluau
// SPEC.md section 8). Run by build.zig at build time:
//
//   stdlib_compile <out.cpp> <name>=<file.luau>...
//
// Each module's `--!pw` header must parse, and each module must pass the determinism lint: no
// `pairs` or `next`, and no `for ... in t` over a bare table, since those follow hash order.
// AI-authored (Claude).

#include "Luau/Ast.h"
#include "Luau/Compiler.h"
#include "Luau/Parser.h"
#include "header.hpp"

#include <algorithm>
#include <array>
#include <cstdio>
#include <cstdlib>
#include <exception>
#include <memory>
#include <print>
#include <span>
#include <string>
#include <string_view>
#include <vector>

namespace {

struct Source {
    std::string name;
    std::string path;
    std::string text;
};

struct FileCloser {
    void operator()(std::FILE *file) const noexcept {
        // The tool only reads through this handle, so a failed close loses nothing.
        static_cast<void>(std::fclose(file)); // NOLINT(cppcoreguidelines-owning-memory)
    }
};
using File = std::unique_ptr<std::FILE, FileCloser>;

[[noreturn]] void die(std::string_view message) {
    std::print(stderr, "stdlib_compile: {}\n", message);
    std::exit(1); // NOLINT(concurrency-mt-unsafe): a single-threaded build tool.
}

std::string read_file(const std::string &path) {
    const File file(std::fopen(path.c_str(), "rb")); // NOLINT(cppcoreguidelines-owning-memory)
    if (!file) {
        die("cannot open " + path);
    }
    std::string text;
    std::array<char, 4096> buffer{};
    while (std::feof(file.get()) == 0 && std::ferror(file.get()) == 0) {
        const std::size_t n = std::fread(buffer.data(), 1, buffer.size(), file.get());
        text.append(buffer.data(), n);
    }
    if (std::ferror(file.get()) != 0) {
        die("cannot read " + path);
    }
    return text;
}

/// The determinism lint, over Luau's syntax tree.
class Lint : public Luau::AstVisitor {
  public:
    explicit Lint(const Source &source) : source_(&source) {}

    bool visit(Luau::AstExprGlobal *node) override {
        const std::string_view name = node->name.value;
        if (name == "pairs" || name == "next") {
            report(node->location,
                   std::string(name) +
                       " iterates in hash order; use pw.sorted or a sorted iterator");
        }
        return true;
    }

    bool visit(Luau::AstStatForIn *node) override {
        if (node->values.size == 1 && (*node->values.begin())->as<Luau::AstExprCall>() == nullptr) {
            report(node->location, "for ... in over a table iterates in hash order; use pw.sorted");
        }
        return true;
    }

    [[nodiscard]] bool clean() const { return clean_; }

  private:
    void report(const Luau::Location &location, std::string_view message) {
        std::print(stderr, "{}:{}: {}\n", source_->path, location.begin.line + 1, message);
        clean_ = false;
    }

    const Source *source_;
    bool clean_ = true;
};

/// The `--!pw` lines of a module, as the runtime re-reads them.
std::string header_lines(std::string_view text) {
    std::string out;
    std::size_t at = 0;
    while (at < text.size()) {
        const std::size_t end = text.find('\n', at);
        const std::string_view line =
            text.substr(at, end == std::string_view::npos ? std::string_view::npos : end - at);
        if (line.starts_with("--!pw")) {
            out.append(line);
            out += '\n';
        }
        at = end == std::string_view::npos ? text.size() : end + 1;
    }
    return out;
}

std::string escaped(std::string_view text) {
    std::string out;
    for (const char c : text) {
        if (c == '\\' || c == '"') {
            out += '\\';
            out += c;
        } else if (c == '\n') {
            out += "\\n";
        } else {
            out += c;
        }
    }
    return out;
}

/// Checks and compiles one module, appending its bytecode array to `out`; false on any error.
bool compile(const Source &source, std::size_t index, std::string &out) {
    bool ok = true;
    if (auto header = packluau::parse_header(source.text); !header) {
        std::print(stderr, "{}:{}: {}\n", source.path, header.error().line, header.error().message);
        ok = false;
    }
    if (header_lines(source.text).empty()) {
        std::print(stderr, "{}: a stdlib module starts with a --!pw caps= header\n", source.path);
        ok = false;
    }
    Luau::Allocator allocator;
    Luau::AstNameTable names(allocator);
    const Luau::ParseResult parsed =
        Luau::Parser::parse(source.text.data(), source.text.size(), names, allocator);
    for (const Luau::ParseError &error : parsed.errors) {
        std::print(stderr, "{}:{}: {}\n", source.path, error.getLocation().begin.line + 1,
                   error.getMessage());
        ok = false;
    }
    if (parsed.errors.empty()) {
        Lint lint(source);
        parsed.root->visit(&lint);
        ok = ok && lint.clean();
    }
    Luau::CompileOptions options;
    options.optimizationLevel = 2;
    options.debugLevel = 1;
    const std::string bytecode = Luau::compile(source.text, options);
    if (bytecode.empty() || bytecode.front() == 0) {
        std::print(stderr, "{}: {}\n", source.path,
                   bytecode.empty() ? std::string_view("no bytecode")
                                    : std::string_view(bytecode).substr(1));
        return false;
    }
    out += "constexpr unsigned char module_" + std::to_string(index) + "[] = {";
    for (std::size_t k = 0; k < bytecode.size(); ++k) {
        out += k % 24 == 0 ? "\n    " : " ";
        out += std::to_string(static_cast<unsigned char>(bytecode.at(k)));
        out += ',';
    }
    out += "\n};\n\n";
    return ok;
}

} // namespace

namespace {

int run(std::span<char *> arguments) {
    if (arguments.size() < 2) {
        die("usage: stdlib_compile <out.cpp> <name>=<file.luau>...");
    }
    const std::string output = arguments.subspan(1, 1).front();
    std::vector<Source> sources;
    for (const char *arg : arguments.subspan(2)) {
        const std::string_view spec(arg);
        const std::size_t equals = spec.find('=');
        if (equals == std::string_view::npos) {
            die("expected <name>=<file.luau>, got " + std::string(spec));
        }
        Source source{.name = std::string(spec.substr(0, equals)),
                      .path = std::string(spec.substr(equals + 1)),
                      .text = {}};
        source.text = read_file(source.path);
        sources.push_back(std::move(source));
    }
    std::ranges::sort(sources, {}, &Source::name);

    bool ok = true;
    std::string out =
        "// Generated by apps/packluau/tools/stdlib_compile.cpp from the Luau stdlib.\n"
        "// Do not edit.\n\n#include \"stdlib.hpp\"\n\n#include <array>\n\n"
        "namespace packluau {\n\nnamespace {\n\n";
    for (std::size_t i = 0; i < sources.size(); ++i) {
        ok = compile(sources.at(i), i, out) && ok;
    }
    if (!ok) {
        return 1;
    }
    out += "} // namespace\n\nstd::span<const StdlibModule> stdlib_modules() noexcept {\n"
           "    static constexpr std::array<StdlibModule, " +
           std::to_string(sources.size()) + "> modules = {{\n";
    for (std::size_t i = 0; i < sources.size(); ++i) {
        const Source &source = sources.at(i);
        out += "        {.name = \"" + escaped(source.name) + "\", .header = \"" +
               escaped(header_lines(source.text)) + "\", .bytecode = module_" + std::to_string(i) +
               "},\n";
    }
    out += "    }};\n    return modules;\n}\n\n} // namespace packluau\n";

    const File file(std::fopen(output.c_str(), "wb")); // NOLINT(cppcoreguidelines-owning-memory)
    if (!file || std::fwrite(out.data(), 1, out.size(), file.get()) != out.size() ||
        std::fflush(file.get()) != 0) {
        die("cannot write " + output);
    }
    return 0;
}

} // namespace

int main(int argc, char **argv) {
    try {
        return run(std::span<char *>(argv, static_cast<std::size_t>(argc)));
    } catch (const std::exception &e) {
        static_cast<void>(std::fputs("stdlib_compile: ", stderr));
        static_cast<void>(std::fputs(e.what(), stderr));
        static_cast<void>(std::fputs("\n", stderr));
        return 1;
    }
}