const std = @import("std");
const pad = @import("padding.zig");

const assert = std.debug.assert;
const Allocator = std.mem.Allocator;
const ArrayList = std.ArrayList;
const math = std.math;
const mem = std.mem;

// ShellIndexFn is a function prototype that, given a shell's index (in
// global shell section), will return a shell string. Matches ShelReader.get.
const ShellIndexFn = fn (u6) []const u8;

// User is a convenient public struct for record construction and
// serialization. Iterator can help retrieve these records.
pub const User = struct {
    uid: u32,
    gid: u32,
    name: []const u8,
    gecos: []const u8,
    home: []const u8,
    shell: []const u8,
};

pub const PackedUser = struct {
    const AlignmentBits = 3;
    const shellIndexFnType = fn ([]const u8) ?u6;

    const InnerSize = @divExact(@bitSizeOf(Inner), 8);
    const Inner = packed struct {
        uid: u32,
        gid: u32,
        additional_gids_offset: u29,
        shell_here: bool,
        shell_len_or_idx: u6,
        home_len: u6,
        name_is_a_suffix: bool,
        name_len: u5,
        gecos_len: u8,

        fn homeLen(self: *const Inner) usize {
            return @as(u32, self.home_len) + 1;
        }

        fn nameStart(self: *const Inner) usize {
            const name_len = self.nameLen();
            if (self.name_is_a_suffix) {
                return self.homeLen() - name_len;
            } else {
                return self.homeLen();
            }
        }

        fn nameLen(self: *const Inner) usize {
            return @as(u32, self.name_len) + 1;
        }

        fn gecosStart(self: *const Inner) usize {
            if (self.name_is_a_suffix) {
                return self.homeLen();
            } else {
                return self.homeLen() + self.nameLen();
            }
        }

        fn gecosLen(self: *const Inner) usize {
            return self.gecos_len;
        }

        fn maybeShellStart(self: *const Inner) usize {
            assert(self.shell_here);
            return self.gecosStart() + self.gecosLen();
        }

        fn shellLen(self: *const Inner) usize {
            return @as(u32, self.shell_len_or_idx) + 1;
        }

        // blobLength returns the length of the blob storing string values.
        fn blobLength(self: *const Inner) usize {
            var result: usize = self.homeLen() + self.gecosLen();
            if (!self.name_is_a_suffix) {
                result += self.nameLen();
            }
            if (self.shell_here) {
                result += self.shellLen();
            }
            return result;
        }
    };

    // PackedUser does not allocate; it re-interprets the "bytes" blob field.
    // Both of those fields are pointers to "our representation" of that field.
    inner: *const Inner,
    userdata: []const u8,

    pub const Entry = struct {
        user: PackedUser,
        next: ?[]const u8,
    };

    pub fn fromBytes(bytes: []const u8) Entry {
        const inner = std.mem.bytesAsValue(
            Inner,
            // Should use InnerSize instead of sizeOf, see
            // https://github.com/ziglang/zig/issues/10958
            bytes[0..@sizeOf(Inner)],
        );

        const startBlob = InnerSize;
        const endBlob = startBlob + inner.blobLength();

        const nextStart = pad.roundUp(usize, AlignmentBits, endBlob);
        var next: ?[]const u8 = null;
        if (nextStart < bytes.len) {
            next = bytes[nextStart..];
        }

        return Entry{
            .user = PackedUser{
                .inner = inner,
                .userdata = bytes[startBlob..endBlob],
            },
            .next = next,
        };
    }

    fn downCast(comptime T: type, n: u64) error{InvalidRecord}!T {
        return std.math.cast(T, n) catch |err| switch (err) {
            error.Overflow => {
                return error.InvalidRecord;
            },
        };
    }

    fn validateUtf8(s: []const u8) error{InvalidRecord}!void {
        if (!std.unicode.utf8ValidateSlice(s)) {
            return error.InvalidRecord;
        }
    }

    // packTo packs the User record and copies it to the given byte slice. The
    // slice must have at least maxRecordSize() bytes available.
    // The slice is passed as a pointer, so it can be mutated.
    pub fn packTo(
        buf: *[]u8,
        user: User,
        shellIndexFn: shellIndexFnType,
    ) error{InvalidRecord}!void {
        // function arguments are consts. We need to mutate the underlying
        // slice, so passing it via pointer instead.
        const home_len = try downCast(u6, user.home.len - 1);
        const name_len = try downCast(u5, user.name.len - 1);
        const shell_len = try downCast(u6, user.shell.len - 1);
        const gecos_len = try downCast(u8, user.gecos.len);

        try validateUtf8(user.home);
        try validateUtf8(user.name);
        try validateUtf8(user.shell);
        try validateUtf8(user.gecos);

        const inner = Inner{
            .uid = user.uid,
            .gid = user.gid,
            .additional_gids_offset = std.math.maxInt(u29),
            .shell_here = shellIndexFn(user.shell) == null,
            .shell_len_or_idx = shellIndexFn(user.shell) orelse shell_len,
            .home_len = home_len,
            .name_is_a_suffix = std.mem.endsWith(u8, user.home, user.name),
            .name_len = name_len,
            .gecos_len = gecos_len,
        };
        const innerBytes = mem.asBytes(&inner);
        var pos: usize = buf.*.len;
        buf.*.len += InnerSize +
            user.home.len +
            user.gecos.len;

        // innerBytes.len is longer than InnerSize. We want to copy
        // only the InnerSize-number of bytes.
        mem.copy(u8, buf.*[pos .. pos + InnerSize], innerBytes[0..InnerSize]);
        pos += InnerSize;
        mem.copy(u8, buf.*[pos .. pos + user.home.len], user.home);
        pos += user.home.len;

        if (!inner.name_is_a_suffix) {
            buf.*.len += user.name.len;
            mem.copy(u8, buf.*[pos .. pos + user.name.len], user.name);
            pos += user.name.len;
        }
        mem.copy(u8, buf.*[pos .. pos + user.gecos.len], user.gecos);
        pos += user.gecos.len;

        if (inner.shell_here) {
            buf.*.len += user.shell.len;
            mem.copy(u8, buf.*[pos .. pos + user.shell.len], user.shell);
            pos += user.shell.len;
        }

        const padding = pad.roundUpPadding(u64, AlignmentBits, pos);
        buf.*.len += padding;
        mem.set(u8, buf.*[pos .. pos + padding], 0);
    }

    // maxSize is the maximum number of records a PackedUser can take
    // (struct + userdata).
    pub fn maxSize() usize {
        comptime {
            const unpadded = InnerSize +
                std.math.maxInt(u6) + // home
                std.math.maxInt(u5) + // name
                std.math.maxInt(u6) + // shell
                std.math.maxInt(u8); // gecos
            return pad.roundUp(u64, AlignmentBits, unpadded);
        }
    }

    pub fn uid(self: *const PackedUser) u32 {
        return self.inner.uid;
    }

    pub fn gid(self: *const PackedUser) u32 {
        return self.inner.gid;
    }

    pub fn home(self: *const PackedUser) []const u8 {
        return self.userdata[0..self.inner.homeLen()];
    }

    pub fn name(self: *const PackedUser) []const u8 {
        const name_pos = self.inner.nameStart();
        const name_len = self.inner.nameLen();
        return self.userdata[name_pos .. name_pos + name_len];
    }

    pub fn gecos(self: *const PackedUser) []const u8 {
        const gecos_pos = self.inner.gecosStart();
        const gecos_len = self.inner.gecosLen();
        return self.userdata[gecos_pos .. gecos_pos + gecos_len];
    }

    pub fn shell(self: *const PackedUser, idxFn: ShellIndexFn) []const u8 {
        if (self.inner.shell_here) {
            const shell_pos = self.inner.maybeShellStart();
            const shell_len = self.inner.shellLen();
            return self.userdata[shell_pos .. shell_pos + shell_len];
        }
        return idxFn(self.inner.shell_len_or_idx);
    }
};

pub fn userIterator(section: []const u8, idxFn: ShellIndexFn) Iterator {
    return Iterator{
        .section = section,
        .shellIndex = idxFn,
    };
}

pub const Iterator = struct {
    section: ?[]const u8,
    shellIndex: ShellIndexFn,

    pub fn next(it: *Iterator) ?PackedUser {
        if (it.section) |section| {
            const entry = PackedUser.fromBytes(section);
            it.section = entry.next;
            return entry.user;
        }
        return null;
    }
};

const testing = std.testing;

test "PackedUser internal and external alignment" {
    // External padding (AlignmentBits) must be higher or equal to
    // the "internal" PackedUser alignment. By aligning PackedUser we are also
    // working around https://github.com/ziglang/zig/issues/10958 ; PackedUser
    // cannot be converted from/to [@bitSizeOf(PackedUser)/8]u8;
    // asBytes/bytesAsValue use @sizeOf, which is larger. Now we are putting no
    // more than 1, but it probably could be higher.
    try testing.expect(@sizeOf(PackedUser) * 8 - @bitSizeOf(PackedUser) <= 8);
}

fn testShellIndex(shell: []const u8) ?u6 {
    if (std.mem.eql(u8, shell, "/bin/bash")) {
        return 0;
    } else if (std.mem.eql(u8, shell, "/bin/zsh")) {
        return 1;
    }
    return null;
}

fn testShell(index: u6) []const u8 {
    return switch (index) {
        0 => "/bin/bash",
        1 => "/bin/zsh",
        else => unreachable,
    };
}

test "construct PackedUser section" {
    var buf = ArrayList(u8).init(testing.allocator);
    defer buf.deinit();

    const users = [_]User{ User{
        .uid = 1000,
        .gid = 1000,
        .name = "vidmantas",
        .gecos = "Vidmantas Kaminskas",
        .home = "/home/vidmantas",
        .shell = "/bin/bash",
    }, User{
        .uid = 1001,
        .gid = 1001,
        .name = "svc-foo",
        .gecos = "Service Account",
        .home = "/home/service1",
        .shell = "/usr/bin/nologin",
    }, User{
        .uid = 0,
        .gid = math.maxInt(u32),
        .name = "Name" ** 8,
        .gecos = "Gecos" ** 51,
        .home = "Home" ** 16,
        .shell = "She.l..l" ** 8,
    }, User{
        .uid = 1002,
        .gid = 1002,
        .name = "svc-bar",
        .gecos = "",
        .home = "/",
        .shell = "/",
    } };
    for (users) |user| {
        try buf.ensureUnusedCapacity(PackedUser.maxSize());
        try PackedUser.packTo(&buf.items, user, testShellIndex);
    }

    var it = userIterator(buf.items, testShell);
    var i: u32 = 0;
    while (it.next()) |user| : (i += 1) {
        try testing.expectEqual(users[i].uid, user.uid());
        try testing.expectEqual(users[i].gid, user.gid());
        try testing.expectEqualStrings(users[i].name, user.name());
        try testing.expectEqualStrings(users[i].gecos, user.gecos());
        try testing.expectEqualStrings(users[i].home, user.home());
        try testing.expectEqualStrings(users[i].shell, user.shell(testShell));
    }
    try testing.expectEqual(users.len, i);
}