const std = @import("../../../std.zig");
const mem = std.mem;
const uefi = std.os.uefi;
const Allocator = mem.Allocator;
const Guid = uefi.Guid;
const assert = std.debug.assert;

// All Device Path Nodes are byte-packed and may appear on any byte boundary.
// All code references to device path nodes must assume all fields are unaligned.

pub const DevicePath = extern struct {
    type: uefi.DevicePath.Type,
    subtype: u8,
    length: u16 align(1),

    pub const CreateFileDevicePathError = Allocator.Error;

    pub const guid align(8) = Guid{
        .time_low = 0x09576e91,
        .time_mid = 0x6d3f,
        .time_high_and_version = 0x11d2,
        .clock_seq_high_and_reserved = 0x8e,
        .clock_seq_low = 0x39,
        .node = [_]u8{ 0x00, 0xa0, 0xc9, 0x69, 0x72, 0x3b },
    };

    /// Returns the next DevicePath node in the sequence, if any.
    pub fn next(self: *const DevicePath) ?*const DevicePath {
        const subtype: uefi.DevicePath.End.Subtype = @enumFromInt(self.subtype);
        if (self.type == .end and subtype == .end_entire) return null;
        const bytes: [*]const u8 = @ptrCast(self);
        return @ptrCast(bytes + self.length);
    }

    /// Calculates the total length of the device path structure in bytes, including the end of device path node.
    pub fn size(self: *const DevicePath) usize {
        var node = self;

        while (node.next()) |next_node| {
            node = next_node;
        }

        return (@intFromPtr(node) + node.length) - @intFromPtr(self);
    }

    /// Creates a file device path from the existing device path and a file path.
    pub fn createFileDevicePath(
        self: *const DevicePath,
        allocator: Allocator,
        path: []const u16,
    ) CreateFileDevicePathError!*const DevicePath {
        const path_size = self.size();

        // 2 * (path.len + 1) for the path and its null terminator, which are u16s
        // DevicePath for the extra node before the end
        var buf = try allocator.alloc(u8, path_size + 2 * (path.len + 1) + @sizeOf(DevicePath));

        @memcpy(buf[0..path_size], @as([*]const u8, @ptrCast(self))[0..path_size]);

        // Pointer to the copy of the end node of the current chain, which is - 4 from the buffer
        // as the end node itself is 4 bytes (type: u8 + subtype: u8 + length: u16).
        var new = @as(*uefi.DevicePath.Media.FilePathDevicePath, @ptrCast(buf.ptr + path_size - 4));

        new.type = .media;
        new.subtype = .file_path;
        new.length = @sizeOf(uefi.DevicePath.Media.FilePathDevicePath) + 2 * (@as(u16, @intCast(path.len)) + 1);

        // The same as new.getPath(), but not const as we're filling it in.
        var ptr = @as([*:0]align(1) u16, @ptrCast(@as([*]u8, @ptrCast(new)) + @sizeOf(uefi.DevicePath.Media.FilePathDevicePath)));

        for (path, 0..) |s, i|
            ptr[i] = s;

        ptr[path.len] = 0;

        var end = @as(*uefi.DevicePath.End.EndEntireDevicePath, @ptrCast(@constCast(@as(*DevicePath, @ptrCast(new)).next().?)));
        end.type = .end;
        end.subtype = .end_entire;
        end.length = @sizeOf(uefi.DevicePath.End.EndEntireDevicePath);

        return @as(*DevicePath, @ptrCast(buf.ptr));
    }

    pub fn getDevicePath(self: *const DevicePath) ?uefi.DevicePath {
        const u_info = @typeInfo(uefi.DevicePath).@"union";
        inline for (u_info.field_names, u_info.field_types) |ufield_name, ufield_type| {
            const enum_value = std.meta.stringToEnum(uefi.DevicePath.Type, ufield_name);

            // Got the associated union type for self.type, now
            // we need to initialize it and its subtype
            if (self.type == enum_value) {
                const subtype = self.initSubtype(ufield_type);
                if (subtype) |sb| {
                    // e.g. return .{ .hardware = .{ .pci = @ptrCast(...) } }
                    return @unionInit(uefi.DevicePath, ufield_name, sb);
                }
            }
        }

        return null;
    }

    pub fn initSubtype(self: *const DevicePath, comptime TUnion: type) ?TUnion {
        const type_info = @typeInfo(TUnion).@"union";
        const TTag = type_info.tag_type.?;

        inline for (type_info.field_names, type_info.field_types) |subtype_name, subtype_type| {
            // The tag names match the union names, so just grab that off the enum
            const tag_val: u8 = @intFromEnum(@field(TTag, subtype_name));

            if (self.subtype == tag_val) {
                // e.g. expr = .{ .pci = @ptrCast(...) }
                return @unionInit(TUnion, subtype_name, @as(subtype_type, @ptrCast(self)));
            }
        }

        return null;
    }
};

comptime {
    assert(4 == @sizeOf(DevicePath));
    assert(1 == @alignOf(DevicePath));

    assert(0 == @offsetOf(DevicePath, "type"));
    assert(1 == @offsetOf(DevicePath, "subtype"));
    assert(2 == @offsetOf(DevicePath, "length"));
}
