const std = @import("std");

fn insertFakeLinuxImageHeader(step: *std.Build.Step, opts: std.Build.Step.MakeOptions) anyerror!void {
    const RISCV_MAGIC1 = "RISCV\x00\x00\x00";
    const RISCV_MAGIC2 = "RSC\x05";

    const elf = std.elf;

    var ehdr: elf.Ehdr = undefined;

    var file = try std.fs.cwd().openFile("zig-out/bin/kernel", .{ .mode = .read_write });
    _ = try file.readAll(std.mem.asBytes(&ehdr));

    // Figure out total image size
    var imageAddrMax: u64 = 0;
    for (0..ehdr.e_phnum) |i| {
        var phdr: elf.Phdr = undefined;
        _ = try file.preadAll(std.mem.asBytes(&phdr), ehdr.e_phoff + i * ehdr.e_phentsize);

        const end = (phdr.p_vaddr + phdr.p_memsz + 0xFFF) & ~@as(u64, 0xFFF);

        if (phdr.p_type == elf.PT_LOAD and end > imageAddrMax) {
            imageAddrMax = end;
        }
    }

    for (0..ehdr.e_phnum) |i| {
        var phdr: elf.Phdr = undefined;
        var data: [64]u8 = undefined;

        _ = try file.preadAll(std.mem.asBytes(&phdr), ehdr.e_phoff + i * ehdr.e_phentsize);

        if (phdr.p_type != elf.PT_LOAD) {
            continue;
        }

        _ = try file.preadAll(&data, phdr.p_offset);

        if (std.mem.eql(u8, RISCV_MAGIC1, data[48..56]) and std.mem.eql(u8, RISCV_MAGIC2, data[56..60])) {
            try file.pwriteAll(std.mem.asBytes(&imageAddrMax), phdr.p_offset + 16);
            break;
        }
    }

    _ = step;
    _ = opts;
}

pub fn build(b: *std.Build) anyerror!void {
    const optimize = b.standardOptimizeOption(.{ .preferred_optimize_mode = .ReleaseFast });
    const target = b.standardTargetOptions(.{ .default_target = .{
        .cpu_arch = .riscv64,
        .os_tag = .freestanding,
        .abi = .none,
    } });

    const kernel_module = b.addModule("kernel", .{ .optimize = optimize, .target = target, .pic = true, .red_zone = false, .code_model = .medium, .root_source_file = b.path("src/kernel.zig") });

    const kernel = b.addExecutable(.{ .name = "kernel", .root_module = kernel_module, .pic = true });
    kernel.pie = true;

    kernel.entry = .{ .symbol_name = "__rv64_entry" };

    kernel.setLinkerScript(b.path("etc/riscv64-unknown-none.ld"));
    kernel.addCSourceFiles(.{
        .files = &.{"src/arch/riscv64/entry.S"},
        .flags = &.{},
    });
    b.installArtifact(kernel);

    const fakeLinuxHeader: *std.Build.Step = try b.allocator.create(std.Build.Step);
    fakeLinuxHeader.* = std.Build.Step.init(.{
        .id = std.Build.Step.Id.custom,
        .name = "insert fake linux header",
        .owner = kernel.step.owner,
        .makeFn = insertFakeLinuxImageHeader,
    });

    const elf2bin = b.addSystemCommand(&.{ "llvm-objcopy", "-O", "binary", "zig-out/bin/kernel", "zig-out/bin/kernel.bin" });

    // TODO QEMU binary override
    const qemu_info = switch (target.result.cpu.arch) {
        .riscv64 => .{ "qemu-system-riscv64", "rv64" },
        else => unreachable,
    };

    const qemu_cmd = b.addSystemCommand(&.{ qemu_info[0], "-M", "virt", "-kernel", "zig-out/bin/kernel.bin", "-m", "256M", "-cpu", qemu_info[1], "-serial", "mon:stdio", "-display", "none" });

    if (target.result.cpu.arch == .riscv64) {
        qemu_cmd.addArgs(&.{ "-bios", "etc/boot/rv64_fw_jump.bin" });
    }

    fakeLinuxHeader.dependOn(b.getInstallStep());
    elf2bin.step.dependOn(fakeLinuxHeader);
    qemu_cmd.step.dependOn(&elf2bin.step);
    if (b.args) |args| qemu_cmd.addArgs(args);
    const run_step = b.step("run", "Start the OS in qemu");
    run_step.dependOn(&qemu_cmd.step);
}