const std = @import("std");
const uefi = std.os.uefi;
const log = std.log;

const Path = @import("../../Config.zig").Path;
const luefi = @import("../../luefi.zig");

const Initrd = @This();

protocol: LoadFile2 = .{},

initrd_buf: []align(4096) u8,
initrd_pages: usize,

handle: uefi.Handle,

pub fn load(paths: []const Path) !*Initrd {
    const File = struct { file: *uefi.protocol.File, size: usize };

    var files: std.ArrayList(File) = try .initCapacity(uefi.pool_allocator, paths.len);
    defer files.deinit(uefi.pool_allocator);

    defer for (files.items) |f| f.file.close() catch unreachable;

    var total_size: usize = 0;
    for (paths) |path| {
        log.info("loading initrd {f}", .{path});
        const file = try luefi.fileOpen(path);
        errdefer file.close() catch unreachable;

        const size = try luefi.fileSize(file);
        try files.append(uefi.pool_allocator, .{ .file = file, .size = size });

        total_size += size;
    }

    const pages_count = try std.math.divCeil(usize, total_size, 4096);

    const boot_svc = uefi.system_table.boot_services.?;

    const initrd_pages = try boot_svc.allocatePages(.any, .loader_data, pages_count);
    errdefer boot_svc.freePages(initrd_pages) catch unreachable;

    const initrd_buf_ptr: [*]align(4096) u8 = @ptrCast(initrd_pages);
    const initrd_buf = initrd_buf_ptr[0..total_size];

    var start_pos: usize = 0;
    for (files.items) |file| {
        const slice = initrd_buf[start_pos .. start_pos + file.size];
        try luefi.fileReadSlice(file.file, slice);
        start_pos += file.size;
    }

    const self = try uefi.pool_allocator.create(Initrd);

    self.* = .{
        .initrd_buf = initrd_buf,
        .initrd_pages = pages_count,
        .handle = undefined,
    };

    self.handle = try luefi.installProtocols(.{ &device_path, &self.protocol });

    return self;
}

pub fn deinit(self: *Initrd) void {
    luefi.uninstallProtocols(self.handle, .{ &device_path, &self.protocol });
    luefi.freePages(self.initrd_buf, self.initrd_pages);
    uefi.pool_allocator.destroy(self);
}

pub const LoadFile2 = extern struct {
    loadFile: *const @TypeOf(impl) = impl,

    pub fn impl(
        protocol: *LoadFile2,
        _: *uefi.protocol.DevicePath,
        _: bool,
        buf_size: *usize,
        buf_out: ?[*]u8,
    ) callconv(uefi.cc) uefi.Status {
        const initrd: *Initrd = @fieldParentPtr("protocol", protocol);

        if (buf_out == null or buf_size.* < initrd.initrd_buf.len) {
            buf_size.* = initrd.initrd_buf.len;
            return .buffer_too_small;
        }

        @memcpy(buf_out.?, initrd.initrd_buf);

        return .success;
    }

    pub const guid: uefi.Guid = .{
        .time_low = 0x4006c0c1,
        .time_mid = 0xfcb3,
        .time_high_and_version = 0x403e,
        .clock_seq_high_and_reserved = 0x99,
        .clock_seq_low = 0x6d,
        .node = [_]u8{ 0x4a, 0x6c, 0x87, 0x24, 0xe0, 0x6d },
    };
};

pub var device_path = struct {
    vendor: uefi.DevicePath.Media.VendorDevicePath,
    end: uefi.DevicePath.End.EndEntireDevicePath,

    pub const guid = uefi.protocol.DevicePath.guid;
}{
    .vendor = uefi.DevicePath.Media.VendorDevicePath{
        .type = .media,
        .subtype = .vendor,
        .length = @sizeOf(uefi.DevicePath.Media.VendorDevicePath),
        .guid = .{
            .time_low = 0x5568e427,
            .time_mid = 0x68fc,
            .time_high_and_version = 0x4f3d,
            .clock_seq_high_and_reserved = 0xac,
            .clock_seq_low = 0x74,
            .node = [_]u8{ 0xca, 0x55, 0x52, 0x31, 0xcc, 0x68 },
        },
    },
    .end = uefi.DevicePath.End.EndEntireDevicePath{
        .type = .end,
        .subtype = .end_entire,
        .length = @sizeOf(uefi.DevicePath.End.EndEntireDevicePath),
    },
};