//! This tool generates SPIR-V features from the grammar files in the SPIRV-Headers
//! (https://github.com/KhronosGroup/SPIRV-Headers/) and SPIRV-Registry (https://github.com/KhronosGroup/SPIRV-Registry/)
//! repositories. Currently it only generates a basic feature set definition consisting of versions, extensions and capabilities.
//! There is a lot left to be desired, as currently dependencies of extensions and dependencies on extensions aren't generated.
//! This is because there are some peculiarities in the SPIR-V registries:
//! - Capabilities may depend on multiple extensions, which cannot be modelled yet by std.Target.
//! - Extension dependencies are not documented in a machine-readable manner.
//! - Note that the grammar spec also contains definitions from extensions which aren't actually official. Most of these seem to be
//!   from an intel project (https://github.com/intel/llvm/, https://github.com/intel/llvm/tree/sycl/sycl/doc/extensions/SPIRV),
//!   and so ONLY extensions in the SPIRV-Registry should be included.

const std = @import("std");
const fs = std.fs;
const Allocator = std.mem.Allocator;
const g = @import("spirv/grammar.zig");

const Version = struct {
    major: u32,
    minor: u32,

    fn parse(str: []const u8) !Version {
        var it = std.mem.splitScalar(u8, str, '.');

        const major = it.first();
        const minor = it.next() orelse return error.InvalidVersion;

        if (it.next() != null) return error.InvalidVersion;

        return Version{
            .major = std.fmt.parseInt(u32, major, 10) catch return error.InvalidVersion,
            .minor = std.fmt.parseInt(u32, minor, 10) catch return error.InvalidVersion,
        };
    }

    fn eql(a: Version, b: Version) bool {
        return a.major == b.major and a.minor == b.minor;
    }

    fn lessThan(ctx: void, a: Version, b: Version) bool {
        _ = ctx;
        return if (a.major == b.major)
            a.minor < b.minor
        else
            a.major < b.major;
    }
};

pub fn main() !void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const args = try std.process.argsAlloc(allocator);

    if (args.len <= 1) {
        usageAndExit(std.io.getStdErr(), args[0], 1);
    }
    if (std.mem.eql(u8, args[1], "--help")) {
        usageAndExit(std.io.getStdErr(), args[0], 0);
    }
    if (args.len != 3) {
        usageAndExit(std.io.getStdErr(), args[0], 1);
    }

    const spirv_headers_root = args[1];
    const spirv_registry_root = args[2];

    if (std.mem.startsWith(u8, spirv_headers_root, "-") or std.mem.startsWith(u8, spirv_registry_root, "-")) {
        usageAndExit(std.io.getStdErr(), args[0], 1);
    }

    // Required for json parsing.
    @setEvalBranchQuota(10000);

    const registry_path = try fs.path.join(allocator, &.{ spirv_headers_root, "include", "spirv", "unified1", "spirv.core.grammar.json" });
    const registry_json = try std.fs.cwd().readFileAlloc(allocator, registry_path, std.math.maxInt(usize));
    var scanner = std.json.Scanner.initCompleteInput(allocator, registry_json);
    var diagnostics = std.json.Diagnostics{};
    scanner.enableDiagnostics(&diagnostics);
    const registry = std.json.parseFromTokenSourceLeaky(g.CoreRegistry, allocator, &scanner, .{}) catch |err| {
        std.debug.print("line,col: {},{}\n", .{ diagnostics.getLine(), diagnostics.getColumn() });
        return err;
    };

    const capabilities = for (registry.operand_kinds) |opkind| {
        if (std.mem.eql(u8, opkind.kind, "Capability"))
            break opkind.enumerants orelse return error.InvalidRegistry;
    } else return error.InvalidRegistry;

    const extensions = try gather_extensions(allocator, spirv_registry_root);
    const versions = try gatherVersions(allocator, registry);

    var bw = std.io.bufferedWriter(std.io.getStdOut().writer());
    const w = bw.writer();

    try w.writeAll(
        \\//! This file is auto-generated by tools/update_spirv_features.zig.
        \\//! TODO: Dependencies of capabilities on extensions.
        \\//! TODO: Dependencies of extensions on extensions.
        \\//! TODO: Dependencies of extensions on versions.
        \\
        \\const std = @import("../std.zig");
        \\const CpuFeature = std.Target.Cpu.Feature;
        \\const CpuModel = std.Target.Cpu.Model;
        \\
        \\pub const Feature = enum {
        \\
    );

    for (versions) |ver| {
        try w.print("    v{}_{},\n", .{ ver.major, ver.minor });
    }

    for (extensions) |ext| {
        try w.print("    {p},\n", .{std.zig.fmtId(ext)});
    }

    for (capabilities) |cap| {
        try w.print("    {p},\n", .{std.zig.fmtId(cap.enumerant)});
    }

    try w.writeAll(
        \\};
        \\
        \\pub const featureSet = CpuFeature.FeatureSetFns(Feature).featureSet;
        \\pub const featureSetHas = CpuFeature.FeatureSetFns(Feature).featureSetHas;
        \\pub const featureSetHasAny = CpuFeature.FeatureSetFns(Feature).featureSetHasAny;
        \\pub const featureSetHasAll = CpuFeature.FeatureSetFns(Feature).featureSetHasAll;
        \\
        \\pub const all_features = blk: {
        \\    @setEvalBranchQuota(2000);
        \\    const len = @typeInfo(Feature).@"enum".fields.len;
        \\    std.debug.assert(len <= CpuFeature.Set.needed_bit_count);
        \\    var result: [len]CpuFeature = undefined;
        \\
    );

    for (versions, 0..) |ver, i| {
        try w.print(
            \\    result[@intFromEnum(Feature.v{0}_{1})] = .{{
            \\        .llvm_name = null,
            \\        .description = "SPIR-V version {0}.{1}",
            \\
        , .{ ver.major, ver.minor });

        if (i == 0) {
            try w.writeAll(
                \\        .dependencies = featureSet(&[_]Feature{}),
                \\    };
                \\
            );
        } else {
            try w.print(
                \\        .dependencies = featureSet(&[_]Feature{{
                \\            .v{}_{},
                \\        }}),
                \\    }};
                \\
            , .{ versions[i - 1].major, versions[i - 1].minor });
        }
    }

    // TODO: Extension dependencies.
    for (extensions) |ext| {
        try w.print(
            \\    result[@intFromEnum(Feature.{p_})] = .{{
            \\        .llvm_name = null,
            \\        .description = "SPIR-V extension {s}",
            \\        .dependencies = featureSet(&[_]Feature{{}}),
            \\    }};
            \\
        , .{
            std.zig.fmtId(ext),
            ext,
        });
    }

    // TODO: Capability extension dependencies.
    for (capabilities) |cap| {
        try w.print(
            \\    result[@intFromEnum(Feature.{p_})] = .{{
            \\        .llvm_name = null,
            \\        .description = "Enable SPIR-V capability {s}",
            \\        .dependencies = featureSet(&[_]Feature{{
            \\
        , .{
            std.zig.fmtId(cap.enumerant),
            cap.enumerant,
        });

        if (cap.version) |ver_str| {
            if (!std.mem.eql(u8, ver_str, "None")) {
                const ver = try Version.parse(ver_str);
                try w.print("            .v{}_{},\n", .{ ver.major, ver.minor });
            }
        }

        for (cap.capabilities) |cap_dep| {
            try w.print("            .{p_},\n", .{std.zig.fmtId(cap_dep)});
        }

        try w.writeAll(
            \\        }),
            \\    };
            \\
        );
    }

    try w.writeAll(
        \\    const ti = @typeInfo(Feature);
        \\    for (&result, 0..) |*elem, i| {
        \\        elem.index = i;
        \\        elem.name = ti.@"enum".fields[i].name;
        \\    }
        \\    break :blk result;
        \\};
        \\
    );

    try bw.flush();
}

/// SPIRV-Registry should hold all extensions currently registered for SPIR-V.
/// The *.grammar.json in SPIRV-Headers should have most of these as well, but with this we're sure to get only the actually
/// registered ones.
/// TODO: Unfortunately, neither repository contains a machine-readable list of extension dependencies.
fn gather_extensions(allocator: Allocator, spirv_registry_root: []const u8) ![]const []const u8 {
    const extensions_path = try fs.path.join(allocator, &.{ spirv_registry_root, "extensions" });
    var extensions_dir = try fs.cwd().openDir(extensions_path, .{ .iterate = true });
    defer extensions_dir.close();

    var extensions = std.ArrayList([]const u8).init(allocator);

    var vendor_it = extensions_dir.iterate();
    while (try vendor_it.next()) |vendor_entry| {
        std.debug.assert(vendor_entry.kind == .directory); // If this fails, the structure of SPIRV-Registry has changed.

        const vendor_dir = try extensions_dir.openDir(vendor_entry.name, .{ .iterate = true });
        var ext_it = vendor_dir.iterate();
        while (try ext_it.next()) |ext_entry| {
            // There is both a HTML and asciidoc version of every spec (as well as some other directories),
            // we need just the name, but to avoid duplicates here we will just skip anything thats not asciidoc.
            if (!std.mem.endsWith(u8, ext_entry.name, ".asciidoc"))
                continue;

            // Unfortunately, some extension filenames are incorrect, so we need to look for the string in the 'Name Strings' section.
            // This has the following format:
            // ```
            // Name Strings
            // ------------
            //
            // SPV_EXT_name
            // ```
            // OR
            // ```
            // == Name Strings
            //
            // SPV_EXT_name
            // ```

            const ext_spec = try vendor_dir.readFileAlloc(allocator, ext_entry.name, std.math.maxInt(usize));
            const name_strings = "Name Strings";

            const name_strings_offset = std.mem.indexOf(u8, ext_spec, name_strings) orelse return error.InvalidRegistry;

            // As the specs are inconsistent on this next part, just skip any newlines/minuses
            var ext_start = name_strings_offset + name_strings.len + 1;
            while (ext_spec[ext_start] == '\n' or ext_spec[ext_start] == '-') {
                ext_start += 1;
            }

            const ext_end = std.mem.indexOfScalarPos(u8, ext_spec, ext_start, '\n') orelse return error.InvalidRegistry;
            const ext = ext_spec[ext_start..ext_end];

            std.debug.assert(std.mem.startsWith(u8, ext, "SPV_")); // Sanity check, all extensions should have a name like SPV_VENDOR_extension.

            try extensions.append(try allocator.dupe(u8, ext));
        }
    }

    return extensions.items;
}

fn insertVersion(versions: *std.ArrayList(Version), version: ?[]const u8) !void {
    const ver_str = version orelse return;
    if (std.mem.eql(u8, ver_str, "None"))
        return;

    const ver = try Version.parse(ver_str);
    for (versions.items) |existing_ver| {
        if (ver.eql(existing_ver)) return;
    }

    try versions.append(ver);
}

fn gatherVersions(allocator: Allocator, registry: g.CoreRegistry) ![]const Version {
    // Expected number of versions is small
    var versions = std.ArrayList(Version).init(allocator);

    for (registry.instructions) |inst| {
        try insertVersion(&versions, inst.version);
    }

    for (registry.operand_kinds) |opkind| {
        const enumerants = opkind.enumerants orelse continue;
        for (enumerants) |enumerant| {
            try insertVersion(&versions, enumerant.version);
        }
    }

    std.mem.sort(Version, versions.items, {}, Version.lessThan);

    return versions.items;
}

fn usageAndExit(file: fs.File, arg0: []const u8, code: u8) noreturn {
    file.writer().print(
        \\Usage: {s} /path/git/SPIRV-Headers /path/git/SPIRV-Registry
        \\
        \\Prints to stdout Zig code which can be used to replace the file lib/std/target/spirv.zig.
        \\
        \\SPIRV-Headers can be cloned from https://github.com/KhronosGroup/SPIRV-Headers,
        \\SPIRV-Registry can be cloned from https://github.com/KhronosGroup/SPIRV-Registry.
        \\
    , .{arg0}) catch std.process.exit(1);
    std.process.exit(code);
}