MsvcLibDir

const MsvcLibDir = struct {
    fn findInstancesDirViaSetup(gpa: Allocator, io: Io, registry: *Registry) error{ OutOfMemory, PathNotFound }!Dir {
        const vs_setup_key_path = L("Microsoft\\VisualStudio\\Setup");
        const vs_setup_key = registry.openSoftwareKey(.{ .root = .local_machine }, vs_setup_key_path) catch |err| switch (err) {
            error.KeyNotFound => return error.PathNotFound,
        };
        defer vs_setup_key.close();

        const packages_path = vs_setup_key.getString(gpa, .{ .name = L("CachePath") }, .wtf8) catch |err| switch (err) {
            error.NotAString,
            error.ValueNameNotFound,
            error.StringNotFound,
            => return error.PathNotFound,

            error.OutOfMemory => |e| return e,
        };
        defer gpa.free(packages_path);

        if (!std.fs.path.isAbsolute(packages_path)) return error.PathNotFound;

        const instances_path = try std.fs.path.join(gpa, &.{ packages_path, "_Instances" });
        defer gpa.free(instances_path);

        return Dir.openDirAbsolute(io, instances_path, .{ .iterate = true }) catch return error.PathNotFound;
    }

    fn findInstancesDirViaCLSID(gpa: Allocator, io: Io, registry: *Registry) error{ OutOfMemory, PathNotFound }!Dir {
        const setup_configuration_clsid = "{177f0c4a-1cd3-4de7-a32c-71dbbb9fa36d}";

        // HKEY_CLASSES_ROOT is not a single key but instead a combination of
        // HKCU\Software\Classes and HKLM\Software\Classes with HKCU taking precedent
        // https://learn.microsoft.com/en-us/windows/win32/sysinfo/hkey-classes-root-key
        //
        // Instead of a CLASSES_ROOT abstraction, we emulate the behavior with a more
        // general abstraction, which also means we need to include `Classes` in the path since
        // we're starting from the `Software` keys instead of the "classes root".
        //
        // The advapi32 APIs with `HKEY_CLASSES_ROOT` go through `\REGISTRY\USER\<SID>_Classes`
        // instead of `\REGISTRY\USER\<SID>\Software\Classes`, but we go through the latter
        // because it allows us to take advantage of `RtlOpenCurrentUser` to avoid needing to implement
        // the logic for getting the current user registry path, and it appears that the two keys are
        // effectively equivalent. Further investigation of the relationship of these keys would probably
        // be beneficial, though.
        const setup_config_key = registry.tryOpenSoftwareKeyWithPrecedence(&.{
            .{ .root = .current_user },
            .{ .root = .local_machine },
        }, L("Classes\\CLSID\\" ++ setup_configuration_clsid)) catch |err| switch (err) {
            error.KeyNotFound => return error.PathNotFound,
        };
        defer setup_config_key.close();

        const inproc_server = setup_config_key.open(L("InprocServer32")) catch return error.PathNotFound;
        const dll_path = inproc_server.getString(gpa, .default, .wtf8) catch |err| switch (err) {
            error.NotAString,
            error.ValueNameNotFound,
            error.StringNotFound,
            => return error.PathNotFound,

            error.OutOfMemory => |e| return e,
        };
        defer gpa.free(dll_path);

        if (!std.fs.path.isAbsolute(dll_path)) return error.PathNotFound;

        var path_it = std.fs.path.componentIterator(dll_path);
        // the .dll filename
        _ = path_it.last();
        const root_path = while (path_it.previous()) |dir_component| {
            if (std.ascii.eqlIgnoreCase(dir_component.name, "VisualStudio")) {
                break dir_component.path;
            }
        } else {
            return error.PathNotFound;
        };

        const instances_path = try std.fs.path.join(gpa, &.{ root_path, "Packages", "_Instances" });
        defer gpa.free(instances_path);

        return Dir.openDirAbsolute(io, instances_path, .{ .iterate = true }) catch return error.PathNotFound;
    }

    fn findInstancesDir(
        gpa: Allocator,
        io: Io,
        registry: *Registry,
        environ_map: *const Environ.Map,
    ) error{ OutOfMemory, PathNotFound }!Dir {
        // First, try getting the packages cache path from the registry.
        // This only seems to exist when the path is different from the default.
        method1: {
            return findInstancesDirViaSetup(gpa, io, registry) catch |err| switch (err) {
                error.OutOfMemory => |e| return e,
                error.PathNotFound => break :method1,
            };
        }
        // Otherwise, try to get the path from the .dll that would have been
        // loaded via COM for SetupConfiguration.
        method2: {
            return findInstancesDirViaCLSID(gpa, io, registry) catch |err| switch (err) {
                error.OutOfMemory => |e| return e,
                error.PathNotFound => break :method2,
            };
        }
        // If that can't be found, fall back to manually appending
        // `Microsoft\VisualStudio\Packages\_Instances` to %PROGRAMDATA%
        method3: {
            const program_data = std.zig.EnvVar.PROGRAMDATA.get(environ_map) orelse break :method3;

            if (!std.fs.path.isAbsolute(program_data)) break :method3;

            const instances_path = try Dir.path.join(gpa, &.{
                program_data, "Microsoft", "VisualStudio", "Packages", "_Instances",
            });
            defer gpa.free(instances_path);

            return Dir.openDirAbsolute(io, instances_path, .{ .iterate = true }) catch break :method3;
        }
        return error.PathNotFound;
    }

    /// Intended to be equivalent to `ISetupHelper.ParseVersion`
    /// Example: 17.4.33205.214 -> 0x0011000481b500d6
    fn parseVersionQuad(version: []const u8) error{InvalidVersion}!u64 {
        var it = std.mem.splitScalar(u8, version, '.');
        const a = it.first();
        const b = it.next() orelse return error.InvalidVersion;
        const c = it.next() orelse return error.InvalidVersion;
        const d = it.next() orelse return error.InvalidVersion;
        if (it.next()) |_| return error.InvalidVersion;
        var result: u64 = undefined;
        var result_bytes = std.mem.asBytes(&result);

        std.mem.writeInt(
            u16,
            result_bytes[0..2],
            std.fmt.parseUnsigned(u16, d, 10) catch return error.InvalidVersion,
            .little,
        );
        std.mem.writeInt(
            u16,
            result_bytes[2..4],
            std.fmt.parseUnsigned(u16, c, 10) catch return error.InvalidVersion,
            .little,
        );
        std.mem.writeInt(
            u16,
            result_bytes[4..6],
            std.fmt.parseUnsigned(u16, b, 10) catch return error.InvalidVersion,
            .little,
        );
        std.mem.writeInt(
            u16,
            result_bytes[6..8],
            std.fmt.parseUnsigned(u16, a, 10) catch return error.InvalidVersion,
            .little,
        );

        return result;
    }

    /// Intended to be equivalent to ISetupConfiguration.EnumInstances:
    /// https://learn.microsoft.com/en-us/dotnet/api/microsoft.visualstudio.setup.configuration
    /// but without the use of COM in order to avoid a dependency on ole32.dll
    ///
    /// The logic in this function is intended to match what ISetupConfiguration does
    /// under-the-hood, as verified using Procmon.
    fn findViaCOM(
        gpa: Allocator,
        io: Io,
        registry: *Registry,
        arch: std.Target.Cpu.Arch,
        environ_map: *const Environ.Map,
    ) error{ OutOfMemory, PathNotFound }![]const u8 {
        // Typically `%PROGRAMDATA%\Microsoft\VisualStudio\Packages\_Instances`
        // This will contain directories with names of instance IDs like 80a758ca,
        // which will contain `state.json` files that have the version and
        // installation directory.
        var instances_dir = try findInstancesDir(gpa, io, registry, environ_map);
        defer instances_dir.close(io);

        var state_subpath_buf: [Dir.max_name_bytes + 32]u8 = undefined;
        var latest_version_lib_dir: std.ArrayList(u8) = .empty;
        errdefer latest_version_lib_dir.deinit(gpa);

        var latest_version: u64 = 0;
        var instances_dir_it = instances_dir.iterateAssumeFirstIteration();
        while (instances_dir_it.next(io) catch return error.PathNotFound) |entry| {
            if (entry.kind != .directory) continue;

            var writer: Writer = .fixed(&state_subpath_buf);

            writer.writeAll(entry.name) catch unreachable;
            writer.writeByte(Dir.path.sep) catch unreachable;
            writer.writeAll("state.json") catch unreachable;

            const json_contents = instances_dir.readFileAlloc(io, writer.buffered(), gpa, .limited(std.math.maxInt(usize))) catch continue;
            defer gpa.free(json_contents);

            var parsed = std.json.parseFromSlice(std.json.Value, gpa, json_contents, .{}) catch continue;
            defer parsed.deinit();

            if (parsed.value != .object) continue;
            const catalog_info = parsed.value.object.get("catalogInfo") orelse continue;
            if (catalog_info != .object) continue;
            const product_version_value = catalog_info.object.get("buildVersion") orelse continue;
            if (product_version_value != .string) continue;
            const product_version_text = product_version_value.string;
            const parsed_version = parseVersionQuad(product_version_text) catch continue;

            // We want to end up with the most recent version installed
            if (parsed_version <= latest_version) continue;

            const installation_path = parsed.value.object.get("installationPath") orelse continue;
            if (installation_path != .string) continue;

            const lib_dir_path = libDirFromInstallationPath(gpa, io, installation_path.string, arch) catch |err| switch (err) {
                error.OutOfMemory => |e| return e,
                error.PathNotFound => continue,
            };
            defer gpa.free(lib_dir_path);

            latest_version_lib_dir.clearRetainingCapacity();
            try latest_version_lib_dir.appendSlice(gpa, lib_dir_path);
            latest_version = parsed_version;
        }

        if (latest_version_lib_dir.items.len == 0) return error.PathNotFound;
        return latest_version_lib_dir.toOwnedSlice(gpa);
    }

    fn libDirFromInstallationPath(
        gpa: Allocator,
        io: Io,
        installation_path: []const u8,
        arch: std.Target.Cpu.Arch,
    ) error{ OutOfMemory, PathNotFound }![]const u8 {
        var lib_dir_buf = try std.array_list.Managed(u8).initCapacity(gpa, installation_path.len + 64);
        errdefer lib_dir_buf.deinit();

        lib_dir_buf.appendSliceAssumeCapacity(installation_path);

        if (!Dir.path.isSep(lib_dir_buf.last().?)) {
            try lib_dir_buf.append('\\');
        }
        const installation_path_with_trailing_sep_len = lib_dir_buf.items.len;

        try lib_dir_buf.appendSlice("VC\\Auxiliary\\Build\\Microsoft.VCToolsVersion.default.txt");
        var default_tools_version_buf: [512]u8 = undefined;
        const default_tools_version_contents = Dir.cwd().readFile(io, lib_dir_buf.items, &default_tools_version_buf) catch {
            return error.PathNotFound;
        };
        var tokenizer = std.mem.tokenizeAny(u8, default_tools_version_contents, " \r\n");
        const default_tools_version = tokenizer.next() orelse return error.PathNotFound;

        lib_dir_buf.shrinkRetainingCapacity(installation_path_with_trailing_sep_len);
        try lib_dir_buf.appendSlice("VC\\Tools\\MSVC\\");
        try lib_dir_buf.appendSlice(default_tools_version);
        try lib_dir_buf.appendSlice("\\Lib\\");
        try lib_dir_buf.appendSlice(switch (arch) {
            .thumb => "arm",
            .aarch64 => "arm64",
            .x86 => "x86",
            .x86_64 => "x64",
            else => unreachable,
        });

        if (!verifyLibDir(io, lib_dir_buf.items)) {
            return error.PathNotFound;
        }

        return lib_dir_buf.toOwnedSlice();
    }

    // https://learn.microsoft.com/en-us/visualstudio/install/tools-for-managing-visual-studio-instances?view=vs-2022#editing-the-registry-for-a-visual-studio-instance
    fn findViaRegistry(
        gpa: Allocator,
        io: Io,
        arch: std.Target.Cpu.Arch,
        environ_map: *const Environ.Map,
    ) error{ OutOfMemory, PathNotFound }![]const u8 {
        // %localappdata%\Microsoft\VisualStudio\
        // %appdata%\Local\Microsoft\VisualStudio\
        const local_app_data_path = std.zig.EnvVar.LOCALAPPDATA.get(environ_map) orelse return error.PathNotFound;
        const visualstudio_folder_path = try Dir.path.join(gpa, &.{
            local_app_data_path, "Microsoft\\VisualStudio\\",
        });
        defer gpa.free(visualstudio_folder_path);

        if (!Dir.path.isAbsolute(visualstudio_folder_path)) return error.PathNotFound;
        // To make things easier later on, we open the VisualStudio directory here which
        // allows us to pass relative paths to NtLoadKeyEx in order to avoid dealing with
        // conversion to NT namespace paths.
        var visualstudio_folder = Dir.openDirAbsolute(io, visualstudio_folder_path, .{
            .iterate = true,
        }) catch return error.PathNotFound;
        defer visualstudio_folder.close(io);

        const vs_versions: []const []const u8 = vs_versions: {
            // enumerate folders that contain `privateregistry.bin`, looking for all versions
            // f.i. %localappdata%\Microsoft\VisualStudio\17.0_9e9cbb98\
            var iterator = visualstudio_folder.iterate();
            break :vs_versions try iterateAndFilterByVersion(&iterator, gpa, io, "");
        };
        defer {
            for (vs_versions) |vs_version| gpa.free(vs_version);
            gpa.free(vs_versions);
        }
        var key_path_buf: [windows.NAME_MAX * 2]u16 = undefined;
        var sub_path_buf: [windows.NAME_MAX * 2]u16 = undefined;
        const source_directories: []const u8 = source_directories: for (vs_versions) |vs_version| {
            const sub_path = blk: {
                var buf: std.ArrayList(u16) = .initBuffer(&sub_path_buf);
                buf.items.len += std.unicode.wtf8ToWtf16Le(buf.unusedCapacitySlice(), vs_version) catch unreachable;
                buf.appendSliceAssumeCapacity(L("\\privateregistry.bin"));
                break :blk buf.items;
            };

            // The goal is to emulate advapi32.RegLoadAppKeyW with a direct call
            // to NtLoadKeyEx instead.
            //
            // RegLoadAppKeyW loads the hive into a registry key of the format:
            // \REGISTRY\A\{fdb2baa5-8ca8-ef03-78d0-3b1f868fd2a9}
            // where `\REGISTRY\A` is a special unenumerable location intended for
            // per-app hives, and the GUID is randomly generated (in testing, it
            // was different for each run of the program).
            //
            // The OS is responsible for cleaning up `\REGISTRY\A` whenever all handles
            // to one of its keys are closed, so we don't have to do anything special
            // with regards to that.

            const temp_key_path = blk: {
                var guid: windows.GUID = undefined;
                io.random(std.mem.asBytes(&guid));

                var guid_buf: [38]u8 = undefined;
                const guid_str = std.mem.print(&guid_buf, "{f}", .{guid}) catch unreachable;

                var buf: std.ArrayList(u16) = .initBuffer(&key_path_buf);
                buf.appendSliceAssumeCapacity(L("\\REGISTRY\\A\\"));
                buf.items.len += std.unicode.wtf8ToWtf16Le(buf.unusedCapacitySlice(), guid_str) catch unreachable;
                break :blk buf.items;
            };

            const target_key: windows.OBJECT.ATTRIBUTES = .{
                .RootDirectory = null,
                .Attributes = .{},
                .ObjectName = @constCast(&windows.UNICODE_STRING.init(temp_key_path)),
                .SecurityDescriptor = null,
            };
            const source_file: windows.OBJECT.ATTRIBUTES = .{
                .RootDirectory = visualstudio_folder.handle,
                .Attributes = .{},
                .ObjectName = @constCast(&windows.UNICODE_STRING.init(sub_path)),
                .SecurityDescriptor = null,
            };
            var root_key: Registry.Key = undefined;
            const rc = windows.ntdll.NtLoadKeyEx(
                &target_key,
                &source_file,
                .{
                    .APP_HIVE = true,
                    // This wasn't set by RegLoadAppKeyW, but it seems relevant
                    // since we aren't intending to do any modifcation of the hive.
                    .OPEN_READ_ONLY = true,
                },
                null,
                null,
                .{ .SPECIFIC = .{
                    .KEY = .{
                        .QUERY_VALUE = true,
                        .ENUMERATE_SUB_KEYS = true,
                    },
                } },
                &root_key.handle,
                null,
            );
            switch (rc) {
                .SUCCESS => {},
                else => continue,
            }
            defer root_key.close();

            const config_path = blk: {
                var buf: std.ArrayList(u16) = .initBuffer(&key_path_buf);
                buf.appendSliceAssumeCapacity(L("Software\\Microsoft\\VisualStudio\\"));
                buf.items.len += std.unicode.wtf8ToWtf16Le(buf.unusedCapacitySlice(), vs_version) catch unreachable;
                buf.appendSliceAssumeCapacity(L("_Config"));
                break :blk buf.items;
            };
            const config_key = root_key.open(config_path) catch continue;

            const source_directories_value = config_key.getString(gpa, .{ .name = L("Source Directories") }, .wtf8) catch |err| switch (err) {
                error.OutOfMemory => |e| return e,
                else => continue,
            };

            break :source_directories source_directories_value;
        } else return error.PathNotFound;
        defer gpa.free(source_directories);

        var source_directories_split = std.mem.splitScalar(u8, source_directories, ';');

        const msvc_dir: []const u8 = msvc_dir: {
            const msvc_include_dir_maybe_with_trailing_slash = try gpa.dupe(u8, source_directories_split.first());

            if (msvc_include_dir_maybe_with_trailing_slash.len > Dir.max_path_bytes or !Dir.path.isAbsolute(msvc_include_dir_maybe_with_trailing_slash)) {
                gpa.free(msvc_include_dir_maybe_with_trailing_slash);
                return error.PathNotFound;
            }

            var msvc_dir = std.array_list.Managed(u8).fromOwnedSlice(gpa, msvc_include_dir_maybe_with_trailing_slash);
            errdefer msvc_dir.deinit();

            // String might contain trailing slash, so trim it here
            if (msvc_dir.items.len > "C:\\".len and msvc_dir.last().? == '\\') _ = msvc_dir.pop();

            // Remove `\include` at the end of path
            if (std.mem.endsWith(u8, msvc_dir.items, "\\include")) {
                msvc_dir.shrinkRetainingCapacity(msvc_dir.items.len - "\\include".len);
            }

            try msvc_dir.appendSlice("\\Lib\\");
            try msvc_dir.appendSlice(switch (arch) {
                .thumb => "arm",
                .aarch64 => "arm64",
                .x86 => "x86",
                .x86_64 => "x64",
                else => unreachable,
            });
            const msvc_dir_with_arch = try msvc_dir.toOwnedSlice();
            break :msvc_dir msvc_dir_with_arch;
        };
        errdefer gpa.free(msvc_dir);

        if (!verifyLibDir(io, msvc_dir)) {
            return error.PathNotFound;
        }

        return msvc_dir;
    }

    fn findViaVs7Key(
        gpa: Allocator,
        io: Io,
        registry: *Registry,
        arch: std.Target.Cpu.Arch,
        environ_map: *const Environ.Map,
    ) error{ OutOfMemory, PathNotFound }![]const u8 {
        var base_path: std.array_list.Managed(u8) = base_path: {
            try_env: {
                if (environ_map.get("VS140COMNTOOLS")) |VS140COMNTOOLS| {
                    if (VS140COMNTOOLS.len < "C:\\Common7\\Tools".len) break :try_env;
                    if (!Dir.path.isAbsolute(VS140COMNTOOLS)) break :try_env;
                    var list = std.array_list.Managed(u8).init(gpa);
                    errdefer list.deinit();

                    try list.appendSlice(VS140COMNTOOLS); // C:\Program Files (x86)\Microsoft Visual Studio 14.0\Common7\Tools
                    // String might contain trailing slash, so trim it here
                    if (list.items.len > "C:\\".len and list.last().? == '\\') _ = list.pop();
                    list.shrinkRetainingCapacity(list.items.len - "\\Common7\\Tools".len); // C:\Program Files (x86)\Microsoft Visual Studio 14.0
                    break :base_path list;
                }
            }

            const vs7_key = registry.openSoftwareKey(.{ .root = .local_machine, .wow64 = .wow64_32 }, L("Microsoft\\VisualStudio\\SxS\\VS7")) catch return error.PathNotFound;
            defer vs7_key.close();
            try_vs7_key: {
                const path_maybe_with_trailing_slash = vs7_key.getString(gpa, .{ .name = L("14.0") }, .wtf8) catch |err| switch (err) {
                    error.OutOfMemory => |e| return e,
                    else => break :try_vs7_key,
                };

                if (path_maybe_with_trailing_slash.len > Dir.max_path_bytes or !Dir.path.isAbsolute(path_maybe_with_trailing_slash)) {
                    gpa.free(path_maybe_with_trailing_slash);
                    break :try_vs7_key;
                }

                var path = std.array_list.Managed(u8).fromOwnedSlice(gpa, path_maybe_with_trailing_slash);
                errdefer path.deinit();

                // String might contain trailing slash, so trim it here
                if (path.items.len > "C:\\".len and path.last().? == '\\') _ = path.pop();
                break :base_path path;
            }
            return error.PathNotFound;
        };
        errdefer base_path.deinit();

        try base_path.appendSlice("\\VC\\lib\\");
        try base_path.appendSlice(switch (arch) {
            .thumb => "arm",
            .aarch64 => "arm64",
            .x86 => "", //x86 is in the root of the Lib folder
            .x86_64 => "amd64",
            else => unreachable,
        });

        if (!verifyLibDir(io, base_path.items)) {
            return error.PathNotFound;
        }

        const full_path = try base_path.toOwnedSlice();
        return full_path;
    }

    fn verifyLibDir(io: Io, lib_dir_path: []const u8) bool {
        std.debug.assert(Dir.path.isAbsolute(lib_dir_path)); // should be already handled in `findVia*`

        var dir = Dir.openDirAbsolute(io, lib_dir_path, .{}) catch return false;
        defer dir.close(io);

        const stat = dir.statFile(io, "vcruntime.lib", .{}) catch return false;
        if (stat.kind != .file)
            return false;

        return true;
    }

    /// Find path to MSVC's `lib/` directory.
    /// Caller owns the result.
    pub fn find(
        gpa: Allocator,
        io: Io,
        registry: *Registry,
        arch: std.Target.Cpu.Arch,
        environ_map: *const Environ.Map,
    ) error{ OutOfMemory, MsvcLibDirNotFound }![]const u8 {
        const full_path = MsvcLibDir.findViaCOM(gpa, io, registry, arch, environ_map) catch |err1| switch (err1) {
            error.OutOfMemory => |e| return e,
            error.PathNotFound => MsvcLibDir.findViaRegistry(gpa, io, arch, environ_map) catch |err2| switch (err2) {
                error.OutOfMemory => |e| return e,
                error.PathNotFound => MsvcLibDir.findViaVs7Key(gpa, io, registry, arch, environ_map) catch |err3| switch (err3) {
                    error.OutOfMemory => |e| return e,
                    error.PathNotFound => return error.MsvcLibDirNotFound,
                },
            },
        };
        errdefer gpa.free(full_path);

        return full_path;
    }
}