const std = @import("std");

const DEFAULT_ARCH = SupportedArch.riscv64;

const SupportedArch = enum {
    aarch64,
    riscv64,

    fn make_target(self: SupportedArch, b: *std.Build) std.Build.ResolvedTarget {
        switch (self) {
            .riscv64 => {
                return b.resolveTargetQuery(.{
                    .cpu_arch = .riscv64,
                    .os_tag = .freestanding,
                    .abi = .none,
                });
            },
            .aarch64 => {
                const T = std.Target.aarch64;

                const add_features = T.featureSet(&.{
                    T.Feature.v8a,
                    T.Feature.strict_align,
                });
                const sub_features = T.featureSet(&.{
                    T.Feature.neon,
                    T.Feature.fp_armv8,
                });

                return b.resolveTargetQuery(.{
                    .cpu_arch = .aarch64,
                    .os_tag = .freestanding,
                    .abi = .none,
                    .cpu_features_add = add_features,
                    .cpu_features_sub = sub_features,
                });
            },
        }
    }

    fn add_target_specific(self: SupportedArch, b: *std.Build, kernel: *std.Build.Step.Compile) anyerror!*std.Build.Step {
        switch (self) {
            .riscv64 => {
                kernel.entry = .{ .symbol_name = "__rv64_entry" };

                kernel.setLinkerScript(b.path("etc/riscv64-unknown-none.ld"));
                kernel.addCSourceFiles(.{
                    .files = &.{"src/arch/riscv64/entry.S"},
                    .flags = &.{},
                });
            },
            .aarch64 => {
                kernel.entry = .{ .symbol_name = "__aa64_entry" };
                kernel.link_z_max_page_size = 0x1000;

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

        b.installArtifact(kernel);

        if (self == .riscv64 or self == .aarch64) {
            const fake_linux_header: *std.Build.Step = try b.allocator.create(std.Build.Step);
            fake_linux_header.* = std.Build.Step.init(.{
                .id = std.Build.Step.Id.custom,
                .name = "insert fake linux header",
                .owner = kernel.step.owner,
                .makeFn = insert_fake_linux_image_header,
            });

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

            fake_linux_header.dependOn(b.getInstallStep());
            elf2bin.step.dependOn(fake_linux_header);

            return &elf2bin.step;
        } else {
            @panic("Unreachable + don't care");
        }
    }
};

fn insert_fake_linux_image_header(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 image_addr_max: 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 > image_addr_max) {
            image_addr_max = 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(&image_addr_max), phdr.p_offset + 16);
            break;
        }
    }

    _ = step;
    _ = opts;
}

fn build_riscv64(b: *std.Build) anyerror!void {
    _ = b;
}

pub fn build(b: *std.Build) anyerror!void {
    const maybe_arch_option = b.option(SupportedArch, "arch", "Architecture to use");

    const arch = maybe_arch_option orelse DEFAULT_ARCH;
    const target = arch.make_target(b);
    const optimize = b.standardOptimizeOption(.{ .preferred_optimize_mode = .ReleaseFast });

    const code_model: std.builtin.CodeModel = switch (arch) {
        .riscv64 => .medium,
        .aarch64 => .small,
    };
    const kernel_module = b.addModule("kernel", .{
        .optimize = optimize,
        .target = target,
        .pic = true,
        .red_zone = false,
        .code_model = code_model,
        .root_source_file = b.path("src/kernel.zig"),
    });
    const kernel = b.addExecutable(.{
        .name = "kernel",
        .root_module = kernel_module,
        .pic = true,
        .use_lld = true,
    });
    kernel.pie = true;

    const install_docs = b.addInstallDirectory(.{
        .source_dir = kernel.getEmittedDocs(),
        .install_dir = .prefix,
        .install_subdir = "docs",
    });

    const docs_step = b.step("docs", "Install documentation");
    docs_step.dependOn(&install_docs.step);

    const kernel_step = try arch.add_target_specific(b, kernel);

    // TODO QEMU binary override
    const qemu_info = switch (target.result.cpu.arch) {
        .riscv64 => .{ "qemu-system-riscv64", "rv64" },
        .aarch64 => .{ "qemu-system-aarch64", "cortex-a72" },
        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" });
    }

    qemu_cmd.step.dependOn(kernel_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);
}