const std = @import("std");
const Step = std.Build.Step;
const OptimizeMode = std.builtin.OptimizeMode;

fn getTarget(b: *std.Build) std.Build.ResolvedTarget {
    var query: std.Target.Query = .{
        .cpu_arch = .x86_64,
        .os_tag = .freestanding,
        .abi = .none,
    };
    const Features = std.Target.x86.Feature;
    query.cpu_features_sub.addFeature(@intFromEnum(Features.mmx));
    query.cpu_features_sub.addFeature(@intFromEnum(Features.sse));
    query.cpu_features_sub.addFeature(@intFromEnum(Features.sse2));
    query.cpu_features_sub.addFeature(@intFromEnum(Features.sse3));
    query.cpu_features_sub.addFeature(@intFromEnum(Features.ssse3));
    query.cpu_features_sub.addFeature(@intFromEnum(Features.sse4_1));
    query.cpu_features_sub.addFeature(@intFromEnum(Features.sse4_2));
    query.cpu_features_sub.addFeature(@intFromEnum(Features.avx));
    query.cpu_features_sub.addFeature(@intFromEnum(Features.avx2));
    query.cpu_features_add.addFeature(@intFromEnum(Features.soft_float));
    const target = b.resolveTargetQuery(query);
    return target;
}
var tool_extract_symbol: *Step.InstallArtifact = undefined;
var tool_generate_image: *Step.InstallArtifact = undefined;
var tool_compress_kernel: *Step.InstallArtifact = undefined;

fn buildTool(b: *std.Build) void {
    const mvzr = b.dependency("mvzr", .{}).module("mvzr");
    const elfy = b.dependency("elfy", .{}).module("elfy");
    const compile_extract_symbol = b.addExecutable(.{
        .name = "symbol_extract",
        .root_module = b.addModule("symbol_extract", .{
            .root_source_file = b.path("arch/x86_64/boot/build_tool/symbol_extract.zig"),
            .target = b.resolveTargetQuery(.{
                .cpu_arch = @import("builtin").target.cpu.arch,
            }),
            .optimize = OptimizeMode.ReleaseSafe,
        }),
    });
    compile_extract_symbol.root_module.addImport("mvzr", mvzr);
    compile_extract_symbol.root_module.addImport("elfy", elfy);
    tool_extract_symbol = b.addInstallArtifact(compile_extract_symbol, .{
        .dest_dir = .{ .override = .{ .custom = "build_tool" } }
    });
    const compile_generate_image = b.addExecutable(.{
        .name = "generate_image",
        .root_module = b.addModule("generate_image", .{
            .root_source_file = b.path("arch/x86_64/boot/build_tool/generate_image.zig"),
            .target = b.resolveTargetQuery(.{
                .cpu_arch = @import("builtin").target.cpu.arch,
            }),
            .optimize = OptimizeMode.ReleaseSafe,
        }),
    });
    tool_generate_image = b.addInstallArtifact(compile_generate_image, .{
        .dest_dir = .{ .override = .{ .custom = "build_tool" } }
    });
    const compile_compress_kernel = b.addExecutable(.{
        .name = "compress_kernel",
        .root_module = b.addModule("compress_kernel", .{
            .root_source_file = b.path("arch/x86_64/boot/build_tool/compress_kernel.zig"),
            .target = b.resolveTargetQuery(.{
                .cpu_arch = @import("builtin").target.cpu.arch,
            }),
            .optimize = OptimizeMode.ReleaseSafe,
        }),
    });
    tool_compress_kernel = b.addInstallArtifact(compile_compress_kernel, .{
        .dest_dir = .{ .override = .{ .custom = "build_tool" } }
    });
} 
fn buildHeader(b: *std.Build, optimize: OptimizeMode, kcapsule: *Step.Compile) *Step.ObjCopy {
    const target = getTarget(b);
    const copy_header = b.addWriteFiles();
    const include_dir = copy_header.addCopyDirectory(b.path("arch/x86_64/boot/include"), "include", .{});
    const symbol_extract = b.addRunArtifact(tool_extract_symbol.artifact);
    symbol_extract.step.dependOn(&tool_extract_symbol.step);
    symbol_extract.addFileArg(kcapsule.getEmittedBin());
    symbol_extract.addArgs(&[_][] const u8{
        "^(_end|_edata|startup_32|startup_64)$",
        "CAP_"
    });
    symbol_extract.addFileArg(include_dir.join(b.allocator, "kcapsule.h") catch @panic("OOM"));
    symbol_extract.step.dependOn(&copy_header.step);
    symbol_extract.step.dependOn(&kcapsule.step);
    const header_elf = b.addExecutable(.{
        .name = "header.bin",
        .root_module = b.addModule("header", .{
            .root_source_file = null,
            .target = target,
            .optimize = optimize,
        }),
    });
    header_elf.root_module.addAssemblyFile(b.path("arch/x86_64/boot/header.S"));
    header_elf.root_module.addIncludePath(include_dir);
    header_elf.setLinkerScript(b.path("arch/x86_64/boot/header.ld"));
    header_elf.step.dependOn(&symbol_extract.step);
    const header = b.addObjCopy(header_elf.getEmittedBin(), .{
        .format = .bin,
    });
    return header;
}
fn buildKcapsule(b: *std.Build, 
    optimize: OptimizeMode,
    kernel: *Step.Compile) *Step.Compile {
    const target = getTarget(b);
    const compress = b.addRunArtifact(tool_compress_kernel.artifact);
    compress.step.dependOn(&kernel.step);
    compress.addFileArg(kernel.getEmittedBin());
    const kernelz_path = kernel.getEmittedBinDirectory().join(b.allocator, "kernelz") catch @panic("OOM");
    const payload_asm_path = kernel.getEmittedBinDirectory().join(b.allocator, "payload.S") catch @panic("OOM");
    compress.addFileArg(kernelz_path);
    compress.addFileArg(payload_asm_path);

    const arch_module = b.addModule("arch", .{
        .root_source_file = b.path("arch/x86_64/arch.zig"),
        .target = target,
        .optimize = optimize,
        .pic = true,
    });
    const loader_module = b.addModule("loader", .{
        .root_source_file = b.path("arch/x86_64/boot/extractor.zig"),
        .target = target,
        .optimize = optimize,
        .pic = true,
    });
    loader_module.addImport("arch", arch_module);
    loader_module.addIncludePath(b.path("arch/x86_64/boot/include"));
    loader_module.addAssemblyFile(payload_asm_path);
    const kcapsule = b.addExecutable(.{
        .name = "kcapsule",
        .use_lld = true,
        .use_llvm = true,
        .root_module = loader_module,
    });
    kcapsule.pie = true;
    kcapsule.setLinkerScript(b.path("arch/x86_64/boot/kcapsule.ld"));
    kcapsule.root_module.addAssemblyFile(b.path("arch/x86_64/boot/head.S"));
    kcapsule.step.dependOn(&compress.step);
    return kcapsule;
}
pub fn buildBootImage(b: *std.Build, kernel: *Step.Compile) void {
    buildTool(b);
    const optimize = OptimizeMode.Debug;
    const kcapsule = buildKcapsule(b, optimize, kernel);
    const header = buildHeader(b, optimize, kcapsule);
    const kcapsule_install = b.addInstallArtifact(kcapsule, .{});
    const kcapsule_bin = b.addObjCopy(kcapsule.getEmittedBin(), .{
        .format = .bin,
    });
    const image = b.addRunArtifact(tool_generate_image.artifact);
    image.addFileArg(header.getOutput());
    image.addFileArg(kcapsule_bin.getOutput());
    image.addArg(std.fs.path.join(b.allocator, &[_][]const u8{ b.install_path, "yukiImage" }) catch @panic("OOM"));
    b.getInstallStep().dependOn(&kcapsule_install.step);
    b.getInstallStep().dependOn(&image.step);
}