const std = @import("std");
const builtin = @import("builtin");

pub fn build(b: *std.Build) !void {
    const arch =
        b.option(std.Target.Cpu.Arch, "arch", "The target kernel architecture")
        orelse builtin.cpu.arch;

    var query: std.Target.Query = .{
        .cpu_arch = arch,
        .os_tag = .freestanding,
        .abi = .none,
    };
    var code_model: std.builtin.CodeModel = .default;

    switch (arch) {
        .aarch64 => {
            const Feature = std.Target.aarch64.Feature;

            query.cpu_features_sub.addFeature(@intFromEnum(Feature.fp_armv8));
            query.cpu_features_sub.addFeature(@intFromEnum(Feature.crypto));
            query.cpu_features_sub.addFeature(@intFromEnum(Feature.neon));
        },
        .loongarch64 => {
            const Feature = std.Target.loongarch.Feature;

            query.cpu_features_add.addFeature(@intFromEnum(Feature.f));
            query.cpu_features_add.addFeature(@intFromEnum(Feature.d));
        },
        .riscv64 => {
            const Feature = std.Target.riscv.Feature;

            query.cpu_features_add.addFeature(@intFromEnum(Feature.f));
            query.cpu_features_sub.addFeature(@intFromEnum(Feature.d));
        },
        .x86_64 => {
            const Feature = std.Target.x86.Feature;

            query.cpu_features_add.addFeature(@intFromEnum(Feature.soft_float));
            query.cpu_features_sub.addFeature(@intFromEnum(Feature.mmx));
            query.cpu_features_sub.addFeature(@intFromEnum(Feature.sse));
            query.cpu_features_sub.addFeature(@intFromEnum(Feature.sse2));
            query.cpu_features_sub.addFeature(@intFromEnum(Feature.avx));
            query.cpu_features_sub.addFeature(@intFromEnum(Feature.avx2));
            code_model = .kernel;
        },
        else => std.debug.panic("Unsupported architecture: {s}", .{@tagName(arch)}),
    }

    // Standard target options allows the person running `zig build` to choose
    // what target to build for. Here we do not override the defaults, which
    // means any target is allowed, and the default is native. Other options
    // for restricting supported target set are available.
    const target = b.resolveTargetQuery(query);

    // Standard optimization options allow the person running `zig build` to select
    // between Debug, ReleaseSafe, ReleaseFast, and ReleaseSmall. Here we do not
    // set a preferred release mode, allowing the user to decide how to optimize.
    const optimize = b.standardOptimizeOption(.{});

    const kernel_mod = b.createModule(.{
        .root_source_file = b.path("src/main.zig"),
        .target = target,
        .optimize = optimize,
        .code_model = code_model,
        .red_zone = false,
    });

    const kernel = b.addExecutable(.{
        .name = "kernel",
        // We will also create a module for our other entry point, 'main.zig'.
        .root_module = kernel_mod,
    });

    // Disable LTO. This prevents Limine requests from being optimized away.
    kernel.want_lto = false;

    // Add kernel dependencies.
    try resolveDependencies(b, kernel_mod);

    // Set the linker script.
    const linker_script = try generateLinkerScript(b, arch);
    kernel.setLinkerScript(linker_script);

    // This declares intent for the executable to be installed into the
    // standard location when the user invokes the "install" step (the default
    // step when running `zig build`).
    b.installArtifact(kernel);

    const img = buildBootableImage(b, kernel.getEmittedBin(), arch);
    const run_cmd = generateRunCommand(b, img, arch);
    
    // By making the run step depend on the install step, it will be run from the
    // installation directory rather than directly from within the cache directory.
    // This is not necessary, however, if the application depends on other installed
    // files, this ensures they will be present and in the expected location.
    run_cmd.step.dependOn(b.getInstallStep());

    // This allows the user to pass arguments to the application in the build
    // command itself, like this: `zig build run -- arg1 arg2 etc`
    if (b.args) |args| {
        run_cmd.addArgs(args);
    }

    // This creates a build step. It will be visible in the `zig build --help` menu,
    // and can be selected like this: `zig build run`
    // This will evaluate the `run` step rather than the default, which is "install".
    const run_step = b.step("run", "Run the kernel");
    run_step.dependOn(&run_cmd.step);

    // Creates a step for unit testing. This only builds the test executable
    // but does not run it.
    const exe_unit_tests = b.addTest(.{
        .root_source_file = b.path("src/main.zig"),
        .target = target,
        .optimize = optimize,
    });

    const run_exe_unit_tests = b.addRunArtifact(exe_unit_tests);

    // Similar to creating the run step earlier, this exposes a `test` step to
    // the `zig build --help` menu, providing a way for the user to request
    // running the unit tests.
    const test_step = b.step("test", "Run unit tests");
    test_step.dependOn(&run_exe_unit_tests.step);
}

fn buildBootableImage(b: *std.Build, kernel: std.Build.LazyPath, arch: std.Target.Cpu.Arch) std.Build.LazyPath {
    // Add Limine to the dependency tree.
    const limine = b.dependency("limine", .{});
    const rootfs = b.addNamedWriteFiles("rootfs");

    // Build the bootloader from source.
    const exe = b.addExecutable(.{
        .name = "limine",
        .target = b.resolveTargetQuery(.{}),
        .optimize = .ReleaseSafe,
    });
    exe.addCSourceFile(.{
        .file = limine.path("limine.c"),
        .flags = &.{"-std=c99"},
    });
    exe.linkLibC();

    _ = rootfs.addCopyFile(b.path("lib/limine/limine.conf"), "boot/limine/limine.conf");
    _ = rootfs.addCopyFile(kernel, "boot/kernel");

    const run_cmd = b.addSystemCommand(switch (arch) {
        .aarch64 => block: {
            _ = rootfs.addCopyFile(limine.path("limine-uefi-cd.bin"), "boot/limine/limine-uefi-cd.bin");
            _ = rootfs.addCopyFile(limine.path("BOOTAA64.EFI"), "EFI/BOOT/BOOTAA64.EFI");

            break :block &.{
                "xorriso", "-as", "mkisofs",
                    "-R", "-r", "-J", 
                    "-hfsplus",
                    "-apm-block-size", "2048",
                    "--efi-boot", "boot/limine/limine-uefi-cd.bin",
                    "-efi-boot-part",
                    "--efi-boot-image",
                    "--protective-msdos-label",
            };
        },
        .loongarch64 => block: {
            _ = rootfs.addCopyFile(limine.path("limine-uefi-cd.bin"), "boot/limine/limine-uefi-cd.bin");
            _ = rootfs.addCopyFile(limine.path("BOOTLOONGARCH64.EFI"), "EFI/BOOT/BOOTLOONGARCH64.EFI");

            break :block &.{
                "xorriso", "-as", "mkisofs",
                    "-R", "-r", "-J",
                    "-hfsplus",
                    "-apm-block-size", "2048",
                    "--efi-boot", "boot/limine/limine-uefi-cd.bin",
                    "-efi-boot-part",
                    "--efi-boot-image",
                    "--protective-msdos-label",
            };
        },
        .riscv64 => block: {
            _ = rootfs.addCopyFile(limine.path("limine-uefi-cd.bin"), "boot/limine/limine-uefi-cd.bin");
            _ = rootfs.addCopyFile(limine.path("BOOTRISCV64.EFI"), "EFI/BOOT/BOOTRISCV64.EFI");

            break :block &.{
                "xorriso", "-as", "mkisofs",
                    "-R", "-r", "-J", 
                    "-hfsplus",
                    "-apm-block-size", "2048",
                    "--efi-boot", "boot/limine/limine-uefi-cd.bin",
                    "-efi-boot-part",
                    "--efi-boot-image",
                    "--protective-msdos-label",
            };
        },
        .x86_64 => block: {
            _ = rootfs.addCopyFile(limine.path("limine-bios.sys"), "boot/limine/limine-bios.sys");
            _ = rootfs.addCopyFile(limine.path("limine-bios-cd.bin"), "boot/limine/limine-bios-cd.bin");
            _ = rootfs.addCopyFile(limine.path("limine-uefi-cd.bin"), "boot/limine/limine-uefi-cd.bin");
            _ = rootfs.addCopyFile(limine.path("BOOTX64.EFI"), "EFI/BOOT/BOOTX64.EFI");
            _ = rootfs.addCopyFile(limine.path("BOOTIA32.EFI"), "EFI/BOOT/BOOTIA32.EFI");

            break :block &.{
                "xorriso", "-as", "mkisofs",
                    "-R", "-r", "-J", 
                    "-b", "boot/limine/limine-bios-cd.bin",
                    "-no-emul-boot",
                    "-boot-load-size", "4",
                    "-boot-info-table",
                    "-hfsplus",
                    "-apm-block-size", "2048",
                    "--efi-boot", "boot/limine/limine-uefi-cd.bin",
                    "-efi-boot-part",
                    "--efi-boot-image",
                    "--protective-msdos-label",
            };
        },
        else => unreachable,
    });

    run_cmd.step.dependOn(&rootfs.step);
    run_cmd.addDirectoryArg(rootfs.getDirectory());
    run_cmd.addArg("-o");
    // Retrieve the bootable image from command.
    const image = run_cmd.addOutputFileArg("bootable.iso");

    // Install Limine stage 1 and 2 for legacy BIOS boot.
    // This is only required on x86_64, but we are deploying for all anyway.
    const deploy = b.addRunArtifact(exe);
    deploy.step.dependOn(&run_cmd.step);
    deploy.addArg("bios-install");
    deploy.addFileArg(image);

    const step = b.step("image", "Generate a bootable ISO image");
    step.dependOn(&deploy.step);

    return image;
}

fn generateLinkerScript(b: *std.Build, arch: std.Target.Cpu.Arch) !std.Build.LazyPath {
    const path = "linker.ld";
    const data = try std.mem.join(b.allocator, "\n", &.{
        switch (arch) {
            .aarch64     => "OUTPUT_FORMAT(elf64-littleaarch64)",
            .loongarch64 => "OUTPUT_FORMAT(elf64-loongarch)",
            .riscv64     => "OUTPUT_FORMAT(elf64-littleriscv)",
            .x86_64      => "OUTPUT_FORMAT(elf64-x86-64)",
            else         => unreachable,
        },
        \\ENTRY(_start)
        \\
        \\/* Define the program headers we want so the bootloader gives us the right */
        \\/* MMU permissions; this also allows us to exert more control over the linking */
        \\/* process. */
        \\PHDRS
        \\{
        \\    text    PT_LOAD;
        \\    rodata  PT_LOAD;
        \\    data    PT_LOAD;
        \\}
        \\
        \\SECTIONS
        \\{
        \\    /* We want to be placed in the topmost 2GiB of the address space, for optimisations */
        \\    /* and because that is what the Limine spec mandates. */
        \\    /* Any address in this region will do, but often 0xffffffff80000000 is chosen as */
        \\    /* that is the beginning of the region. */
        \\    . = 0xffffffff80000000;
        \\
        \\    .text : {
        \\        *(.text .text.*)
        \\    } :text
        \\
        \\    /* Move to the next memory page for .rodata */
        \\    . = ALIGN(CONSTANT(MAXPAGESIZE));
        \\
        \\    .rodata : {
        \\        *(.rodata .rodata.*)
        \\    } :rodata
        \\
        \\    /* Move to the next memory page for .data */
        \\    . = ALIGN(CONSTANT(MAXPAGESIZE));
        \\
        \\    .data : {
        \\        *(.data .data.*)
        \\
        \\        /* Place the sections that contain the Limine requests as part of the .data */
        \\        /* output section. */
        \\        KEEP(*(.requests_start_marker))
        \\        KEEP(*(.requests))
        \\        KEEP(*(.requests_end_marker))
        \\    } :data
        \\
        \\    /* NOTE: .bss needs to be the last thing mapped to :data, otherwise lots of */
        \\    /* unnecessary zeros will be written to the binary. */
        \\    /* If you need, for example, .init_array and .fini_array, those should be placed */
        \\    /* above this. */
        \\    .bss : {
        \\        *(.bss .bss.*)
        \\        *(COMMON)
        \\    } :data
        \\
        \\    /* Discard .note.* and .eh_frame* since they may cause issues on some hosts. */
        \\    /DISCARD/ : {
        \\        *(.eh_frame*)
        \\        *(.note .note.*)
        \\    }
        \\}
    });

    const script = b.addWriteFile(path, data);
    b.getInstallStep().dependOn(&script.step);

    return script.getDirectory().path(b, path);
}

fn generateRunCommand(b: *std.Build, image: std.Build.LazyPath, arch: std.Target.Cpu.Arch) *std.Build.Step.Run {
    // Add edk2_ovmf to the dependency tree
    const edk2 = b.dependency("edk2", .{});

    const fs = b.addWriteFiles();
    var code: std.Build.LazyPath = undefined;
    var vars: std.Build.LazyPath = undefined;

    // This *creates* a Run step in the build graph, to be executed when another
    // step is evaluated that depends on it. The next line below will establish
    // such a dependency.
    const run_cmd: *std.Build.Step.Run = switch (arch) {
        .aarch64 => block: {
            code = fs.addCopyFile(edk2.path("bin/RELEASEAARCH64_QEMU_EFI.fd"), "code.fd");
            vars = fs.addCopyFile(edk2.path("bin/RELEASEAARCH64_QEMU_VARS.fd"), "vars.fd");

            const dd1_cmd = b.addSystemCommand(&.{"dd", "if=/dev/zero", "bs=1", "count=0", "seek=67108864"});
            dd1_cmd.addPrefixedFileArg("of=", code);
            dd1_cmd.step.dependOn(&fs.step);

            const dd2_cmd = b.addSystemCommand(&.{"dd", "if=/dev/zero", "bs=1", "count=0", "seek=67108864"});
            dd2_cmd.addPrefixedFileArg("of=", vars);
            dd2_cmd.step.dependOn(&fs.step);

            const run_cmd = b.addSystemCommand(&.{
                "qemu-system-aarch64",
                    "-M", "virt",
                    "-cpu", "cortex-a72",
                    "-device", "ramfb",
                    "-device", "qemu-xhci",
                    "-device", "usb-kbd",
                    "-device", "usb-mouse",
                    "-serial", "stdio",
            });
            run_cmd.step.dependOn(&dd1_cmd.step);
            run_cmd.step.dependOn(&dd2_cmd.step);

            break :block run_cmd;
        },
        .loongarch64 => block: {
            code = fs.addCopyFile(edk2.path("bin/RELEASELOONGARCH64_QEMU_EFI.fd"), "code.fd");
            vars = fs.addCopyFile(edk2.path("bin/RELEASELOONGARCH64_QEMU_VARS.fd"), "vars.fd");

            const dd1_cmd = b.addSystemCommand(&.{"dd", "if=/dev/zero", "bs=1", "count=0", "seek=5242880"});
            dd1_cmd.addPrefixedFileArg("of=", code);
            dd1_cmd.step.dependOn(&fs.step);

            const dd2_cmd = b.addSystemCommand(&.{"dd", "if=/dev/zero", "bs=1", "count=0", "seek=5242880"});
            dd2_cmd.addPrefixedFileArg("of=", vars);
            dd2_cmd.step.dependOn(&fs.step);

            const run_cmd = b.addSystemCommand(&.{
                "qemu-system-loongarch64",
                    "-M", "virt",
                    "-cpu", "la464",
                    "-device", "ramfb",
                    "-device", "qemu-xhci",
                    "-device", "usb-kbd",
                    "-device", "usb-mouse",
                    "-serial", "stdio",
            });
            run_cmd.step.dependOn(&dd1_cmd.step);
            run_cmd.step.dependOn(&dd2_cmd.step);

            break :block run_cmd;
        },
        .riscv64 => block: {
            code = fs.addCopyFile(edk2.path("bin/RELEASERISCV64_VIRT_CODE.fd"), "code.fd");
            vars = fs.addCopyFile(edk2.path("bin/RELEASERISCV64_VIRT_VARS.fd"), "vars.fd");

            const dd1_cmd = b.addSystemCommand(&.{"dd", "if=/dev/zero", "bs=1", "count=0", "seek=33554432"});
            dd1_cmd.addPrefixedFileArg("of=", code);
            dd1_cmd.step.dependOn(&fs.step);

            const dd2_cmd = b.addSystemCommand(&.{"dd", "if=/dev/zero", "bs=1", "count=0", "seek=33554432"});
            dd2_cmd.addPrefixedFileArg("of=", vars);
            dd2_cmd.step.dependOn(&fs.step);

            const run_cmd = b.addSystemCommand(&.{
                "qemu-system-riscv64",
                    "-M", "virt",
                    "-cpu", "rv64",
                    "-device", "ramfb",
                    "-device", "qemu-xhci",
                    "-device", "usb-kbd",
                    "-device", "usb-mouse",
                    "-serial", "stdio",
            });
            run_cmd.step.dependOn(&dd1_cmd.step);
            run_cmd.step.dependOn(&dd2_cmd.step);

            break :block run_cmd;
        },
        .x86_64 => block: {
            code = edk2.path("bin/RELEASEX64_OVMF_CODE.fd");
            vars = edk2.path("bin/RELEASEX64_OVMF_VARS.fd");

            break :block b.addSystemCommand(&.{
                "qemu-system-x86_64",
                    "-M", "q35",
                    "-device", "isa-debug-exit,iobase=0xf4,iosize=0x04",
                    "-serial", "stdio",
                    "-rtc", "clock=vm",
            });
        },
        else => unreachable,
    };

    run_cmd.addArg("-drive");
    run_cmd.addPrefixedFileArg("if=pflash,unit=0,format=raw,readonly=on,file=", code);

    run_cmd.addArg("-drive");
    run_cmd.addPrefixedFileArg("if=pflash,unit=1,format=raw,file=", vars);

    run_cmd.addArg("-cdrom");
    run_cmd.addFileArg(image);

    return run_cmd;
}

fn resolveDependencies(b: *std.Build, kernel: *std.Build.Module) !void {
    const libk = block: {
        const file = b.path("lib/kernel/kernel.zig");
        const mod = b.addModule("kernel", .{
            .root_source_file = file,
            .optimize = kernel.optimize,
            .target = kernel.resolved_target
        });

        break :block mod;
    };

    const limine = block: {
        const file = b.path("lib/limine/limine.zig");
        const mod = b.createModule(.{ .root_source_file = file });

        break :block mod;
    };

    const uacpi = block: {
        const dependency = b.dependency("uacpi", .{});
        const file = b.path("lib/uacpi/uacpi.zig");
        const mod = b.createModule(.{ .root_source_file = file });

        var flags: std.ArrayList([]const u8) = .init(b.allocator);
        for (b.debug_log_scopes) |scope| {
            if (std.mem.eql(u8, scope, "uacpi")) {
                try flags.append("-DUACPI_DEFAULT_LOG_LEVEL=UACPI_LOG_TRACE");
                break;
            }
        }

        mod.addIncludePath(dependency.path("include"));
        mod.addCSourceFiles(.{
            .root = dependency.path("source"),
            .files = &.{
                "default_handlers.c",
                "event.c",
                "interpreter.c",
                "io.c",
                "mutex.c",
                "namespace.c",
                "notify.c",
                "opcodes.c",
                "opregion.c",
                "osi.c",
                "registers.c",
                "resources.c",
                "shareable.c",
                "sleep.c",
                "stdlib.c",
                "tables.c",
                "types.c",
                "uacpi.c",
                "utilities.c",
            },
            .flags = try flags.toOwnedSlice(),
        });

        break :block mod;
    };

    kernel.addImport("libk", libk);
    kernel.addImport("uacpi", uacpi);
    kernel.addImport("limine", limine);
}