const std = @import("std");

const pad = @import("padding.zig");
const validate = @import("validate.zig");
const compress = @import("compress.zig");
const InvalidRecord = validate.InvalidRecord;

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

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

// User is a convenient public struct for record construction and
// serialization.
pub const User = struct {
    uid: u32,
    gid: u32,
    name: []const u8,
    gecos: []const u8,
    home: []const u8,
    shell: []const u8,

    // deep-clones a User record with a given Allocator.
    pub fn clone(
        self: *const User,
        allocator: Allocator,
    ) Allocator.Error!User {
        const stringdata = try allocator.alloc(u8, self.strlen());
        const gecos_start = self.name.len;
        const home_start = gecos_start + self.gecos.len;
        const shell_start = home_start + self.home.len;
        mem.copy(u8, stringdata[0..self.name.len], self.name);
        mem.copy(u8, stringdata[gecos_start..], self.gecos);
        mem.copy(u8, stringdata[home_start..], self.home);
        mem.copy(u8, stringdata[shell_start..], self.shell);

        return User{
            .uid = self.uid,
            .gid = self.gid,
            .name = stringdata[0..self.name.len],
            .gecos = stringdata[gecos_start .. gecos_start + self.gecos.len],
            .home = stringdata[home_start .. home_start + self.home.len],
            .shell = stringdata[shell_start .. shell_start + self.shell.len],
        };
    }

    fn strlen(self: *const User) usize {
        return self.name.len +
            self.gecos.len +
            self.home.len +
            self.shell.len;
    }

    pub fn deinit(self: *User, allocator: Allocator) void {
        const slice = self.home.ptr[0..self.strlen()];
        allocator.free(slice);
        self.* = undefined;
    }
};

pub fn Shell2Index(T: type) type {
    return struct {
        const Self = @This();
        data: T,

        pub fn init(data: T) Self {
            return Self{ .data = data };
        }

        pub fn get(self: *const Self, str: []const u8) ?u6 {
            return self.data.get(str);
        }
    };
}

pub const PackedUser = struct {
    const Self = @This();

    pub const alignment_bits = 3;

    const Inner = packed struct {
        uid: u32,
        gid: u32,
        padding: u2 = 0,
        shell_len_or_idx: u6,
        shell_here: bool,
        name_is_a_suffix: bool,
        home_len: u6,
        name_len: u5,
        gecos_len: u11,

        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;
        }

        // stringLength returns the length of the blob storing string values.
        fn stringLength(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,
    bytes: []const u8,
    additional_gids_offset: u64,

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

    // TODO(motiejus) provide a way to return an entry without decoding the
    // additional_gids_offset:
    // - will not return the 'next' slice.
    // - cannot throw an Overflow error.
    pub fn fromBytes(bytes: []const u8) error{Overflow}!Entry {
        const inner = mem.bytesAsValue(Inner, bytes[0..@sizeOf(Inner)]);
        const start_blob = @sizeOf(Inner);
        const end_strings = start_blob + inner.stringLength();
        const gids_offset = try compress.uvarint(bytes[end_strings..]);
        const end_blob = end_strings + gids_offset.bytes_read;

        const nextStart = pad.roundUp(usize, alignment_bits, end_blob);
        var next: ?[]const u8 = null;
        if (nextStart < bytes.len)
            next = bytes[nextStart..];

        return Entry{
            .user = Self{
                .inner = inner,
                .bytes = bytes[start_blob..end_blob],
                .additional_gids_offset = gids_offset.value,
            },
            .next = next,
        };
    }

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

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

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

    // 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(
        arr: *ArrayList(u8),
        user: User,
        additional_gids_offset: u64,
        idxFn: StringHashMap(u6),
    ) error{ InvalidRecord, OutOfMemory }!void {
        std.debug.assert(arr.items.len & 7 == 0);
        // function arguments are consts. We need to mutate the underlying
        // slice, so passing it via pointer instead.
        const home_len = try validate.downCast(u6, user.home.len - 1);
        const name_len = try validate.downCast(u5, user.name.len - 1);
        const shell_len = try validate.downCast(u6, user.shell.len - 1);
        const gecos_len = try validate.downCast(u8, user.gecos.len);

        try validate.utf8(user.home);
        try validate.utf8(user.name);
        try validate.utf8(user.shell);
        try validate.utf8(user.gecos);

        const inner = Inner{
            .uid = user.uid,
            .gid = user.gid,
            .shell_here = idxFn.get(user.shell) == null,
            .shell_len_or_idx = idxFn.get(user.shell) orelse shell_len,
            .home_len = home_len,
            .name_is_a_suffix = mem.endsWith(u8, user.home, user.name),
            .name_len = name_len,
            .gecos_len = gecos_len,
        };
        const innerBytes = mem.asBytes(&inner);

        try arr.*.appendSlice(innerBytes[0..@sizeOf(Inner)]);
        try arr.*.appendSlice(user.home);

        if (!inner.name_is_a_suffix)
            try arr.*.appendSlice(user.name);
        try arr.*.appendSlice(user.gecos);
        if (inner.shell_here)
            try arr.*.appendSlice(user.shell);
        try compress.appendUvarint(arr, additional_gids_offset);
    }

    pub fn uid(self: Self) u32 {
        return self.inner.uid;
    }

    pub fn gid(self: Self) u32 {
        return self.inner.gid;
    }

    pub fn additionalGidsOffset(self: Self) u64 {
        return self.additional_gids_offset;
    }

    pub fn home(self: Self) []const u8 {
        return self.bytes[0..self.inner.homeLen()];
    }

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

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

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

const testing = std.testing;

test "PackedUser internal and external alignment" {
    try testing.expectEqual(
        @sizeOf(PackedUser.Inner) * 8,
        @bitSizeOf(PackedUser.Inner),
    );
}

fn testShellIndex(allocator: Allocator) StringHashMap(u6) {
    var result = StringHashMap(u6).init(allocator);
    result.put("/bin/bash", 0) catch unreachable;
    result.put("/bin/zsh", 1) catch unreachable;
    return result;
}

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.LllL" ** 8,
    }, User{
        .uid = 1002,
        .gid = 1002,
        .name = "svc-bar",
        .gecos = "",
        .home = "/",
        .shell = "/",
    } };
    var shellIndex = testShellIndex(testing.allocator);
    defer shellIndex.deinit();
    for (users) |user| {
        try PackedUser.packTo(&buf, user, math.maxInt(u64), shellIndex);
        try pad.arrayList(&buf, PackedUser.alignment_bits);
    }

    var i: u29 = 0;
    var it1 = PackedUser.iterator(buf.items, testShell);
    while (try it1.next()) |user| : (i += 1) {
        try testing.expectEqual(users[i].uid, user.uid());
        try testing.expectEqual(users[i].gid, user.gid());
        try testing.expectEqual(
            @as(u64, math.maxInt(u64)),
            user.additionalGidsOffset(),
        );
        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);
}

test "User.clone" {
    var allocator = testing.allocator;
    const user = User{
        .uid = 1000,
        .gid = 1000,
        .name = "vidmantas",
        .gecos = "Vidmantas Kaminskas",
        .home = "/home/vidmantas",
        .shell = "/bin/bash",
    };
    var user2 = try user.clone(allocator);
    defer user2.deinit(allocator);

    try testing.expectEqualStrings(user.shell, "/bin/bash");
}