const std = @import("std");
const fmt = std.fmt;
const math = std.math;
const sort = std.sort;
const unicode = std.unicode;
const Allocator = std.mem.Allocator;
const ArrayListUnmanaged = std.ArrayListUnmanaged;
const ArrayList = std.ArrayList;
const MultiArrayList = std.MultiArrayList;
const StringHashMap = std.StringHashMap;
const AutoHashMap = std.AutoHashMap;
const BufSet = std.BufSet;

const pad = @import("padding.zig");
const compress = @import("compress.zig");
const shellImport = @import("shell.zig");
const userImport = @import("user.zig");
const groupImport = @import("group.zig");
const cmph = @import("cmph.zig");
const bdz = @import("bdz.zig");
const User = userImport.User;
const Group = groupImport.Group;
const ShellSections = shellImport.ShellWriter.ShellSections;

const Corpus = struct {
    arena: std.heap.ArenaAllocator,

    // sorted by name, by unicode codepoint
    users: []User,
    // sorted by gid
    groups: []Group,

    // convenience users and groups by column
    usersMulti: MultiArrayList(User),
    groupsMulti: MultiArrayList(Group),

    // pointing to `users` and `groups` slices above.
    name2user: StringHashMap(*const User),
    uid2user: AutoHashMap(u32, *const User),
    name2group: StringHashMap(*const Group),
    gid2group: AutoHashMap(u32, *const Group),
    groupname2users: StringHashMap([]*const User),
    username2groups: StringHashMap([]*const Group),

    pub fn init(
        baseAllocator: Allocator,
        usersConst: []const User,
        groupsConst: []const Group,
    ) error{ OutOfMemory, InvalidUtf8, Duplicate, NotFound }!Corpus {
        var arena = std.heap.ArenaAllocator.init(baseAllocator);
        var allocator = arena.allocator();
        errdefer arena.deinit();

        var users = try allocator.alloc(User, usersConst.len);
        var groups = try allocator.alloc(Group, groupsConst.len);
        for (usersConst) |*user, i|
            users[i] = try user.clone(allocator);
        for (groupsConst) |*group, i|
            groups[i] = try group.clone(allocator);

        sort.sort(User, users, {}, cmpUser);
        sort.sort(Group, groups, {}, cmpGroup);

        var usersMulti = MultiArrayList(User){};
        try usersMulti.ensureTotalCapacity(allocator, users.len);
        for (users) |user|
            usersMulti.appendAssumeCapacity(user);
        var groupsMulti = MultiArrayList(Group){};
        try groupsMulti.ensureTotalCapacity(allocator, groups.len);
        for (groups) |group|
            groupsMulti.appendAssumeCapacity(group);

        var name2user = StringHashMap(*const User).init(allocator);
        var uid2user = AutoHashMap(u32, *const User).init(allocator);
        var name2group = StringHashMap(*const Group).init(allocator);
        var gid2group = AutoHashMap(u32, *const Group).init(allocator);
        for (users) |*user| {
            var res1 = try name2user.getOrPut(user.name);
            if (res1.found_existing)
                return error.Duplicate;
            res1.value_ptr.* = user;

            var res2 = try uid2user.getOrPut(user.uid);
            if (res2.found_existing)
                return error.Duplicate;

            res2.value_ptr.* = user;
        }

        for (groups) |*group| {
            var res1 = try name2group.getOrPut(group.name);
            if (res1.found_existing)
                return error.Duplicate;
            res1.value_ptr.* = group;
            var res2 = try gid2group.getOrPut(group.gid);
            if (res2.found_existing)
                return error.Duplicate;
            res2.value_ptr.* = group;
        }

        var groupname2users = StringHashMap([]*const User).init(allocator);

        // uses baseAllocator, because it will be freed before
        // returning from this function. This keeps the arena clean.
        var username2groups = StringHashMap(
            ArrayListUnmanaged(*const Group),
        ).init(baseAllocator);
        defer username2groups.deinit();

        for (groups) |*group| {
            var members = try allocator.alloc(*const User, group.members.count());
            members.len = 0;

            var it = group.members.iterator();
            while (it.next()) |memberName| {
                if (name2user.get(memberName.*)) |user| {
                    members.len += 1;
                    members[members.len - 1] = user;
                } else {
                    return error.NotFound;
                }

                var groupsOfMember = try username2groups.getOrPut(memberName.*);
                if (!groupsOfMember.found_existing)
                    groupsOfMember.value_ptr.* = ArrayListUnmanaged(*const Group){};
                try groupsOfMember.value_ptr.*.append(allocator, group);
            }

            var result = try groupname2users.getOrPut(group.name);
            if (result.found_existing)
                return error.Duplicate;
            result.value_ptr.* = members;
        }

        var it1 = groupname2users.valueIterator();
        while (it1.next()) |groupUsers| {
            sort.sort(*const User, groupUsers.*, {}, cmpUserPtr);
        }

        var it2 = username2groups.valueIterator();
        while (it2.next()) |userGroups|
            sort.sort(*const Group, userGroups.items, {}, cmpGroupPtr);

        var username2groups_final = StringHashMap([]*const Group).init(allocator);
        var it = username2groups.iterator();
        while (it.next()) |elem| {
            const username = elem.key_ptr.*;
            const usergroups = elem.value_ptr.*.toOwnedSlice(allocator);
            try username2groups_final.put(username, usergroups);
        }

        return Corpus{
            .arena = arena,
            .users = users,
            .groups = groups,
            .usersMulti = usersMulti,
            .groupsMulti = groupsMulti,
            .name2user = name2user,
            .uid2user = uid2user,
            .name2group = name2group,
            .gid2group = gid2group,
            .groupname2users = groupname2users,
            .username2groups = username2groups_final,
        };
    }

    pub fn deinit(self: *Corpus) void {
        self.arena.deinit();
        self.* = undefined;
    }
};

pub fn bdzGid(allocator: Allocator, corpus: *const Corpus) cmph.Error![]const u8 {
    return try cmph.pack_u32(allocator, corpus.groupsMulti.items(.gid));
}

pub fn bdzGroupname(allocator: Allocator, corpus: *const Corpus) cmph.Error![]const u8 {
    return try cmph.pack_str(allocator, corpus.groupsMulti.items(.name));
}

pub fn bdzUid(allocator: Allocator, corpus: *const Corpus) cmph.Error![]const u8 {
    return try cmph.pack_u32(allocator, corpus.usersMulti.items(.uid));
}

pub fn bdzUsername(allocator: Allocator, corpus: *const Corpus) cmph.Error![]const u8 {
    return try cmph.pack_str(allocator, corpus.usersMulti.items(.name));
}

// TODO(motiejus) there are a few problems:
// - memory management for shell sections is a mess. Make it easier by ...
// - shell module should accept a list of shells and spit out two slices
// (allocated with a given allocator). There is too much dancing around
// here.
pub fn shellSections(
    allocator: Allocator,
    corpus: *const Corpus,
) error{ OutOfMemory, Overflow }!ShellSections {
    var popcon = shellImport.ShellWriter.init(allocator);
    for (corpus.usersMulti.items(.shell)) |shell| {
        try popcon.put(shell);
    }
    return popcon.toOwnedSections(shellImport.max_shells);
}

pub const UserGids = struct {
    // username -> offset in blob
    name2offset: StringHashMap(u32),
    // compressed user gids blob. A blob contains N <= users.len items,
    // an item is:
    //   len: varint
    //   gid: [varint]varint,
    // ... and the gid list is delta-compressed.
    blob: []u8,

    pub fn deinit(self: *UserGids, allocator: Allocator) void {
        self.name2offset.deinit();
        allocator.free(self.blob);
        self.* = undefined;
    }
};

const userGidsPaddingBits = 3;

pub fn userGids(
    allocator: Allocator,
    corpus: *const Corpus,
) error{ OutOfMemory, Overflow }!UserGids {
    var blob = ArrayList(u8).init(allocator);
    errdefer blob.deinit();
    var name2offset = StringHashMap(u32).init(allocator);
    errdefer name2offset.deinit();

    // zero'th entry is empty, so groupless users can refer to it.
    try compress.appendUvarint(&blob, 0);
    try pad.arrayList(&blob, userGidsPaddingBits);

    var scratch = try allocator.alloc(u32, 256);
    defer allocator.free(scratch);
    for (corpus.users) |user| {
        if (corpus.username2groups.get(user.name)) |usergroups| {
            try name2offset.putNoClobber(user.name, try math.cast(u32, blob.items.len));
            scratch = try allocator.realloc(scratch, usergroups.len);
            scratch.len = usergroups.len;
            for (usergroups) |group, i|
                scratch[i] = group.gid;
            compress.deltaCompress(u32, scratch) catch |err| switch (err) {
                error.NotSorted => unreachable,
            };
            try compress.appendUvarint(&blob, usergroups.len);
            for (scratch) |gid|
                try compress.appendUvarint(&blob, gid);
            try pad.arrayList(&blob, userGidsPaddingBits);
        } else {
            try name2offset.putNoClobber(user.name, 0);
        }
    }

    return UserGids{
        .name2offset = name2offset,
        .blob = blob.toOwnedSlice(),
    };
}

pub fn usersSection(
    allocator: Allocator,
    corpus: *const Corpus,
    gids: *const UserGids,
    shells: *const ShellSections,
) error{ OutOfMemory, Overflow, InvalidRecord }![]const u8 {
    // as of writing each user takes 15 bytes + strings + padding, padded to
    // 8 bytes. 24 is an optimistic lower bound for an average record size.
    var buf = try ArrayList(u8).initCapacity(allocator, 24 * corpus.users.len);
    for (corpus.users) |user| {
        const offset = gids.name2offset.get(user.name).?;
        std.debug.assert(offset & 7 == 0);
        try userImport.PackedUserHash.packTo(
            &buf,
            user,
            @truncate(u29, @shrExact(offset, 3)),
            shells.indices,
        );
    }
    return buf.toOwnedSlice();
}

pub fn groupMembers(
    allocator: Allocator,
    corpus: *const Corpus,
) error{OutOfMemory}!void {
    var buf: [compress.maxVarintLen64]u8 = undefined;
    var offsets = ArrayListUnmanaged(usize).initCapacity(
        allocator,
        corpus.groups.len,
    );
    var bytes = ArrayList(u8).init(allocator);
    var offset: usize = 0;
    for (corpus.groups) |group, i| {
        offsets[i] = offset;
        const users = corpus.groupname2users.get(group.name).?;
        const len = compress.putVarint(&buf, users.len);
        offset += len;
        try bytes.appendSlice(buf[0..len]);
        for (users) |user| {
            // TODO: offset into the User's record
            _ = user;
        }
    }
}

// cmpUser compares two users for sorting. By username's utf8 codepoints, ascending.
fn cmpUser(_: void, a: User, b: User) bool {
    var utf8_a = (unicode.Utf8View.init(a.name) catch unreachable).iterator();
    var utf8_b = (unicode.Utf8View.init(b.name) catch unreachable).iterator();
    while (utf8_a.nextCodepoint()) |codepoint_a| {
        if (utf8_b.nextCodepoint()) |codepoint_b| {
            if (codepoint_a == codepoint_b) {
                continue;
            } else {
                return codepoint_a < codepoint_b;
            }
        }

        // a is a prefix of b. It is thus shorter.
        return false;
    }
    // b is a prefix of a
    return true;
}

fn cmpUserPtr(context: void, a: *const User, b: *const User) bool {
    return cmpUser(context, a.*, b.*);
}

fn cmpGroup(_: void, a: Group, b: Group) bool {
    return a.gid < b.gid;
}

fn cmpGroupPtr(context: void, a: *const Group, b: *const Group) bool {
    return cmpGroup(context, a.*, b.*);
}

const testing = std.testing;

fn testCorpus(allocator: Allocator) !Corpus {
    const users = [_]User{ User{
        .uid = 128,
        .gid = 128,
        .name = "vidmantas",
        .gecos = "Vidmantas Kaminskas",
        .home = "/home/vidmantas",
        .shell = "/bin/bash",
    }, 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 = "/",
    }, User{
        .uid = 65534,
        .gid = 65534,
        .name = "nobody",
        .gecos = "nobody",
        .home = "/nonexistent",
        .shell = "/usr/sbin/nologin",
    } };

    var members1 = try groupImport.someMembers(
        allocator,
        &[_][]const u8{"vidmantas"},
    );
    defer members1.deinit();

    var members2 = try groupImport.someMembers(
        allocator,
        &[_][]const u8{ "svc-bar", "vidmantas" },
    );
    defer members2.deinit();

    var members3 = try groupImport.someMembers(
        allocator,
        &[_][]const u8{ "svc-bar", "Name" ** 8, "vidmantas" },
    );
    defer members3.deinit();

    const groups = [_]Group{ Group{
        .gid = 128,
        .name = "vidmantas",
        .members = members1,
    }, Group{
        .gid = 9999,
        .name = "all",
        .members = members3,
    }, Group{
        .gid = 0,
        .name = "service-account",
        .members = members2,
    } };

    return try Corpus.init(allocator, users[0..], groups[0..]);
}

test "test corpus" {
    var corpus = try testCorpus(testing.allocator);
    defer corpus.deinit();

    try testing.expectEqualStrings(corpus.users[0].name, "Name" ** 8);
    try testing.expectEqualStrings(corpus.users[1].name, "nobody");
    try testing.expectEqualStrings(corpus.users[2].name, "svc-bar");
    try testing.expectEqualStrings(corpus.users[3].name, "vidmantas");

    try testing.expectEqual(corpus.name2user.get("404"), null);
    try testing.expectEqual(corpus.name2user.get("vidmantas").?.uid, 128);
    try testing.expectEqual(corpus.uid2user.get(42), null);
    try testing.expectEqual(corpus.uid2user.get(128).?.gid, 128);
    try testing.expectEqual(corpus.name2group.get("404"), null);
    try testing.expectEqual(corpus.name2group.get("vidmantas").?.gid, 128);
    try testing.expectEqual(corpus.gid2group.get(42), null);
    try testing.expectEqual(corpus.gid2group.get(128).?.gid, 128);

    const membersOfAll = corpus.groupname2users.get("all").?;
    try testing.expectEqualStrings(membersOfAll[0].name, "Name" ** 8);
    try testing.expectEqualStrings(membersOfAll[1].name, "svc-bar");
    try testing.expectEqualStrings(membersOfAll[2].name, "vidmantas");
    try testing.expectEqual(corpus.groupname2users.get("404"), null);

    const groupsOfVidmantas = corpus.username2groups.get("vidmantas").?;
    try testing.expectEqual(groupsOfVidmantas[0].gid, 0);
    try testing.expectEqual(groupsOfVidmantas[1].gid, 128);
    try testing.expectEqual(groupsOfVidmantas[2].gid, 9999);
    try testing.expectEqual(corpus.username2groups.get("nobody"), null);
    try testing.expectEqual(corpus.username2groups.get("doesnotexist"), null);
}

test "test sections" {
    const allocator = testing.allocator;
    var corpus = try testCorpus(allocator);
    defer corpus.deinit();

    const bdz_gid = try bdzGid(allocator, &corpus);
    defer allocator.free(bdz_gid);

    const bdz_groupname = try bdzGroupname(allocator, &corpus);
    defer allocator.free(bdz_groupname);

    const bdz_uid = try bdzUid(allocator, &corpus);
    defer allocator.free(bdz_uid);

    const bdz_username = try bdzUsername(allocator, &corpus);
    defer allocator.free(bdz_username);

    var shell_sections = try shellSections(allocator, &corpus);
    defer shell_sections.deinit();

    var user_gids = try userGids(allocator, &corpus);
    defer user_gids.deinit(allocator);

    var users_section = try usersSection(
        allocator,
        &corpus,
        &user_gids,
        &shell_sections,
    );
    defer allocator.free(users_section);
}

test "userGids" {
    const allocator = testing.allocator;
    var corpus = try testCorpus(allocator);
    defer corpus.deinit();

    var user_gids = try userGids(allocator, &corpus);
    defer user_gids.deinit(allocator);

    for (corpus.users) |user| {
        const groups = corpus.username2groups.get(user.name);
        const offset = user_gids.name2offset.get(user.name);
        if (groups == null) {
            try testing.expect(offset.? == 0);
            continue;
        }
        var vit = try compress.VarintSliceIterator(user_gids.blob[offset.?..]);
        var it = compress.DeltaDecompressionIterator(&vit);
        try testing.expectEqual(it.remaining(), groups.?.len);
        var i: usize = 0;
        while (try it.next()) |gid| : (i += 1) {
            try testing.expectEqual(gid, groups.?[i].gid);
        }
    }
}

test "pack gids" {
    const allocator = testing.allocator;
    var corpus = try testCorpus(allocator);
    defer corpus.deinit();

    const cmph_gid = try cmph.pack_u32(allocator, corpus.groupsMulti.items(.gid));
    defer allocator.free(cmph_gid);

    const k1 = bdz.search_u32(cmph_gid, 0);
    const k2 = bdz.search_u32(cmph_gid, 128);
    const k3 = bdz.search_u32(cmph_gid, 9999);
    var hashes = &[_]u32{ k1, k2, k3 };
    sort.sort(u32, hashes, {}, comptime sort.asc(u32));
    for (hashes) |hash, i|
        try testing.expectEqual(i, hash);
}

fn testUser(name: []const u8) User {
    var result = std.mem.zeroes(User);
    result.name = name;
    return result;
}

test "users compare function" {
    const a = testUser("a");
    const b = testUser("b");
    const bb = testUser("bb");
    try testing.expect(cmpUser({}, a, b));
    try testing.expect(!cmpUser({}, b, a));
    try testing.expect(cmpUser({}, a, bb));
    try testing.expect(!cmpUser({}, bb, a));
    try testing.expect(cmpUser({}, b, bb));
    try testing.expect(!cmpUser({}, bb, b));
}