const std = @import("std");
const os = std.os;
const fmt = std.fmt;
const mem = std.mem;
const math = std.math;
const heap = std.heap;
const once = std.once;
const Mutex = std.Thread.Mutex;
const Allocator = mem.Allocator;

const DB = @import("DB.zig");
const File = @import("File.zig");
const ErrCtx = @import("ErrCtx.zig");
const compress = @import("compress.zig");
const CGroup = @import("Group.zig").CGroup;
const PackedGroup = @import("PackedGroup.zig");
const CUser = @import("User.zig").CUser;
const PackedUser = @import("PackedUser.zig");
const ShellReader = @import("shell.zig").ShellReader;

const c = @cImport({
    @cInclude("nss.h");
});

// cannot use os.getenv due to ziglang/zig#4524
extern fn getenv([*:0]const u8) ?[*:0]const u8;

const ENV_DB = "TURBONSS_DB";
const ENV_VERBOSE = "TURBONSS_VERBOSE";
const ENV_OMIT_MEMBERS = "TURBONSS_OMIT_MEMBERS";

export var turbonss_db_path: [*:0]const u8 = "/etc/turbonss/db.turbo";

// State is a type of the global variable holding the process state:
// the DB handle and all the iterators.
const State = struct {
    file: File,
    omit_members: bool,
    v: bool = false,

    getpwent_iterator_mu: Mutex = Mutex{},
    getpwent_iterator: ?PackedUser.Iterator = null,

    getgrent_iterator_mu: Mutex = Mutex{},
    getgrent_iterator: ?PackedGroup.Iterator = null,

    // allocator for
    initgroups_dyn_allocator: Allocator,
};

// global_state is initialized on first call to an nss function
var global_state: ?State = null;
var global_init = once(init);

// assigns State from environment variables et al
fn init() void {
    const verbose = blk: {
        if (getenv(ENV_VERBOSE)) |env| {
            const got = mem.sliceTo(env, 0);
            if (mem.eql(u8, got, "0")) {
                break :blk false;
            } else if (mem.eql(u8, got, "1")) {
                break :blk true;
            } else {
                std.debug.print(
                    "warning: unrecognized {s}={s}. Expected between 0 or 1\n",
                    .{ ENV_VERBOSE, got },
                );
            }
        }
        break :blk false;
    };

    const omit_members = blk: {
        if (getenv(ENV_OMIT_MEMBERS)) |env| {
            const got = mem.sliceTo(env, 0);
            if (mem.eql(u8, got, "1")) {
                break :blk true;
            } else if (mem.eql(u8, got, "0")) {
                break :blk false;
            } else if (mem.eql(u8, got, "auto")) {
                // not set, do autodiscover
            } else {
                std.debug.print(
                    "warning: unrecognized {s}={s}. Expected 0, 1 or auto\n",
                    .{ ENV_OMIT_MEMBERS, got },
                );
            }
        }

        // argv does not exist because
        // https://github.com/ziglang/zig/issues/4524#issuecomment-1184748756
        // so reading /proc/self/cmdline
        const fd = os.openZ("/proc/self/cmdline", os.O.RDONLY, 0) catch break :blk false;
        defer os.close(fd);
        break :blk isId(fd);
    };

    const fname = if (getenv(ENV_DB)) |env|
        mem.sliceTo(env, 0)
    else
        mem.sliceTo(turbonss_db_path, 0);

    const file = File.open(fname) catch |err| {
        if (verbose)
            std.debug.print("open '{s}': {s}\n", .{ fname, @errorName(err) });
        return;
    };

    if (verbose) {
        std.debug.print("turbonss database '{s}' opened\n", .{fname});
        std.debug.print("omitting members from getgr* calls: {any}\n", .{omit_members});
    }

    global_state = State{
        .file = file,
        .v = verbose,
        .omit_members = omit_members,
        .initgroups_dyn_allocator = heap.raw_c_allocator,
    };
}

export fn _nss_turbo_getpwuid_r(
    uid: c_uint,
    passwd: *CUser,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    const db = getDBErrno(errnop) orelse return c.NSS_STATUS_UNAVAIL;
    return getpwuid_r(db, uid, passwd, buffer, buflen, errnop);
}

fn getpwuid_r(
    db: *const DB,
    uid: c_uint,
    passwd: *CUser,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    var cuser = db.getpwuid(uid, buffer[0..buflen]) catch |err| switch (err) {
        error.BufferTooSmall => {
            errnop.* = @enumToInt(os.E.RANGE);
            return c.NSS_STATUS_TRYAGAIN;
        },
    };

    const got_cuser = cuser orelse {
        errnop.* = @enumToInt(os.E.NOENT);
        return c.NSS_STATUS_NOTFOUND;
    };

    passwd.* = got_cuser;
    return c.NSS_STATUS_SUCCESS;
}

export fn _nss_turbo_getpwnam_r(
    name: [*:0]const u8,
    passwd: *CUser,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    const db = getDBErrno(errnop) orelse return c.NSS_STATUS_UNAVAIL;
    return getpwnam_r(db, name, passwd, buffer, buflen, errnop);
}

fn getpwnam_r(
    db: *const DB,
    name: [*:0]const u8,
    passwd: *CUser,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    const nameSlice = mem.sliceTo(name, 0);

    var buf = buffer[0..buflen];
    const cuser = db.getpwnam(nameSlice, buf) catch |err| switch (err) {
        error.BufferTooSmall => {
            errnop.* = @enumToInt(os.E.RANGE);
            return c.NSS_STATUS_TRYAGAIN;
        },
    };

    const got_cuser = cuser orelse {
        errnop.* = @enumToInt(os.E.NOENT);
        return c.NSS_STATUS_NOTFOUND;
    };

    passwd.* = got_cuser;
    return c.NSS_STATUS_SUCCESS;
}

export fn _nss_turbo_getgrgid_r(
    gid: c_uint,
    gr: *CGroup,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    global_init.call();
    if (global_state) |*state| {
        return getgrgid_r(state, gid, gr, buffer, buflen, errnop);
    } else {
        errnop.* = @enumToInt(os.E.AGAIN);
        return c.NSS_STATUS_UNAVAIL;
    }
}

fn getgrgid_r(
    state: *const State,
    gid: c_uint,
    gr: *CGroup,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    const db = state.file.db;
    const omit_members = state.omit_members;

    var buf = buffer[0..buflen];
    var cgroup = db.getgrgid(gid, buf, omit_members) catch |err| switch (err) {
        error.BufferTooSmall => {
            errnop.* = @enumToInt(os.E.RANGE);
            return c.NSS_STATUS_TRYAGAIN;
        },
    };

    const got_cgroup = cgroup orelse {
        errnop.* = @enumToInt(os.E.NOENT);
        return c.NSS_STATUS_NOTFOUND;
    };

    gr.* = got_cgroup;
    return c.NSS_STATUS_SUCCESS;
}

export fn _nss_turbo_getgrnam_r(
    name: [*:0]const u8,
    group: *CGroup,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    global_init.call();
    if (global_state) |*state| {
        return getgrnam_r(state, name, group, buffer, buflen, errnop);
    } else {
        errnop.* = @enumToInt(os.E.AGAIN);
        return c.NSS_STATUS_UNAVAIL;
    }
}

fn getgrnam_r(
    state: *const State,
    name: [*:0]const u8,
    group: *CGroup,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    const db = state.file.db;
    const omit_members = state.omit_members;
    const nameSlice = mem.sliceTo(name, 0);
    var buf = buffer[0..buflen];
    var cgroup = db.getgrnam(nameSlice, buf, omit_members) catch |err| switch (err) {
        error.BufferTooSmall => {
            errnop.* = @enumToInt(os.E.RANGE);
            return c.NSS_STATUS_TRYAGAIN;
        },
    };
    const got_cgroup = cgroup orelse {
        errnop.* = @enumToInt(os.E.NOENT);
        return c.NSS_STATUS_NOTFOUND;
    };
    group.* = got_cgroup;
    return c.NSS_STATUS_SUCCESS;
}

export fn _nss_turbo_setpwent(_: c_int) c.enum_nss_status {
    global_init.call();
    if (global_state) |*state|
        return setpwent(state)
    else
        return c.NSS_STATUS_UNAVAIL;
}

fn setpwent(state: *State) c.enum_nss_status {
    state.getpwent_iterator_mu.lock();
    defer state.getpwent_iterator_mu.unlock();

    const db = state.file.db;
    state.getpwent_iterator = PackedUser.iterator(
        db.users,
        db.header.num_users,
        db.shellReader(),
    );
    return c.NSS_STATUS_SUCCESS;
}

export fn _nss_turbo_endpwent() c.enum_nss_status {
    global_init.call();
    if (global_state) |*state|
        return endpwent(state)
    else
        return c.NSS_STATUS_UNAVAIL;
}

fn endpwent(state: *State) c.enum_nss_status {
    state.getpwent_iterator_mu.lock();
    state.getpwent_iterator = null;
    state.getpwent_iterator_mu.unlock();
    return c.NSS_STATUS_SUCCESS;
}

export fn _nss_turbo_setgrent(_: c_int) c.enum_nss_status {
    global_init.call();
    if (global_state) |*state|
        return setgrent(state)
    else
        return c.NSS_STATUS_UNAVAIL;
}

fn setgrent(state: *State) c.enum_nss_status {
    state.getgrent_iterator_mu.lock();
    defer state.getgrent_iterator_mu.unlock();
    state.getgrent_iterator = PackedGroup.iterator(
        state.file.db.groups,
        state.file.db.header.num_groups,
    );
    return c.NSS_STATUS_SUCCESS;
}

export fn _nss_turbo_endgrent() c.enum_nss_status {
    global_init.call();
    if (global_state) |*state|
        return endgrent(state)
    else
        return c.NSS_STATUS_UNAVAIL;
}

fn endgrent(state: *State) c.enum_nss_status {
    state.getgrent_iterator_mu.lock();
    state.getgrent_iterator = null;
    state.getgrent_iterator_mu.unlock();
    return c.NSS_STATUS_SUCCESS;
}

export fn _nss_turbo_getgrent_r(
    result: *CGroup,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    global_init.call();
    if (global_state) |*state| {
        return getgrent_r(state, result, buffer, buflen, errnop);
    } else {
        errnop.* = @enumToInt(os.E.AGAIN);
        return c.NSS_STATUS_UNAVAIL;
    }
}

fn getgrent_r(
    state: *State,
    result: *CGroup,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    state.getgrent_iterator_mu.lock();
    defer state.getgrent_iterator_mu.unlock();

    if (state.getgrent_iterator) |*it| {
        const group = it.next() orelse {
            errnop.* = 0;
            return c.NSS_STATUS_NOTFOUND;
        };

        const cgroup1 = if (state.omit_members)
            DB.packCGroupNoMembers(&group, buffer[0..buflen])
        else
            state.file.db.packCGroup(&group, buffer[0..buflen]);

        if (cgroup1) |cgroup| {
            result.* = cgroup;
            return c.NSS_STATUS_SUCCESS;
        } else |err| switch (err) {
            error.BufferTooSmall => {
                it.rollback();
                errnop.* = @enumToInt(os.E.RANGE);
                return c.NSS_STATUS_TRYAGAIN;
            },
        }
    } else {
        // logic from _nss_systemd_getgrent_r
        errnop.* = @enumToInt(os.E.HOSTDOWN);
        return c.NSS_STATUS_UNAVAIL;
    }
}

export fn _nss_turbo_getpwent_r(
    result: *CUser,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    global_init.call();
    if (global_state) |*state| {
        return getpwent_r(state, result, buffer, buflen, errnop);
    } else {
        errnop.* = @enumToInt(os.E.AGAIN);
        return c.NSS_STATUS_UNAVAIL;
    }
}

fn getpwent_r(
    state: *State,
    result: *CUser,
    buffer: [*]u8,
    buflen: usize,
    errnop: *c_int,
) c.enum_nss_status {
    state.getpwent_iterator_mu.lock();
    defer state.getpwent_iterator_mu.unlock();

    if (state.getpwent_iterator) |*it| {
        const user = it.next() orelse {
            errnop.* = 0;
            return c.NSS_STATUS_NOTFOUND;
        };

        const buf = buffer[0..buflen];
        const cuser = state.file.db.writeUser(user, buf) catch |err| switch (err) {
            error.BufferTooSmall => {
                it.rollback();
                errnop.* = @enumToInt(os.E.RANGE);
                return c.NSS_STATUS_TRYAGAIN;
            },
        };

        result.* = cuser;
        return c.NSS_STATUS_SUCCESS;
    } else {
        // logic from _nss_systemd_getgrent_r
        errnop.* = @enumToInt(os.E.HOSTDOWN);
        return c.NSS_STATUS_UNAVAIL;
    }
}

export fn _nss_turbo_initgroups_dyn(
    user_name: [*:0]const u8,
    gid: u32,
    start: *c_long,
    size: *c_long,
    groupsp: *[*]u32,
    limit: c_long,
    errnop: *c_int,
) c.enum_nss_status {
    global_init.call();
    if (global_state) |*state| {
        return initgroups_dyn(state, user_name, gid, start, size, groupsp, limit, errnop);
    } else {
        errnop.* = @enumToInt(os.E.AGAIN);
        return c.NSS_STATUS_UNAVAIL;
    }
}

fn initgroups_dyn(
    state: *const State,
    user_name: [*:0]const u8,
    _: u32,
    start: *c_long,
    size: *c_long,
    groupsp: *[*]u32,
    limit: c_long,
    errnop: *c_int,
) c.enum_nss_status {
    const db = state.file.db;
    const user = db.getUserByName(mem.sliceTo(user_name, 0)) orelse {
        errnop.* = @enumToInt(os.E.NOENT);
        return c.NSS_STATUS_NOTFOUND;
    };

    // TODO: use db.userGids()
    const offset = user.additional_gids_offset;
    var vit = compress.varintSliceIteratorMust(db.additional_gids[offset..]);
    var gids = compress.deltaDecompressionIterator(&vit);
    const remaining = vit.remaining;

    // the implementation below is ported from glibc's db-initgroups.c
    // even though we know the size of the groups upfront, I found it too difficult
    // to preallocate and juggle size, start and limit while keeping glibc happy.
    var any: bool = false;
    while (gids.nextMust()) |gid| {
        if (start.* == size.*) {
            if (limit > 0 and size.* == limit)
                return c.NSS_STATUS_SUCCESS;

            const oldsize = @intCast(usize, size.*);
            const newsize_want = blk: {
                // double the oldsize until it is above or equal 'remaining'
                var res = oldsize;
                while (res < remaining) res *= 2;
                break :blk res;
            };

            const newsize: usize = if (limit <= 0)
                newsize_want
            else
                math.min(@intCast(usize, limit), newsize_want);

            var buf = groupsp.*[0..oldsize];
            const new_groups = state.initgroups_dyn_allocator.realloc(
                buf,
                newsize * @sizeOf(u32),
            );

            if (new_groups) |newgroups| {
                groupsp.* = newgroups.ptr;
                size.* = @intCast(c_long, newsize);
            } else |err| switch (err) {
                error.OutOfMemory => {
                    errnop.* = @enumToInt(os.E.NOMEM);
                    return c.NSS_STATUS_TRYAGAIN;
                },
            }
        }

        any = true;
        groupsp.*[@intCast(usize, start.*)] = @intCast(u32, gid);
        start.* += 1;
    }

    return if (any) c.NSS_STATUS_SUCCESS else c.NSS_STATUS_NOTFOUND;
}

fn getDBErrno(errnop: *c_int) ?*const DB {
    global_init.call();
    if (global_state) |*state| {
        return &state.file.db;
    } else {
        errnop.* = @enumToInt(os.E.AGAIN);
        return null;
    }
}

// isId tells if this command is "id". Reads the cmdline
// from the given fd. Returns false on any error.
fn isId(fd: os.fd_t) bool {
    var cmdline: [256]u8 = undefined;
    const nb = os.read(fd, cmdline[0..]) catch return false;
    const arg0 = mem.sliceTo(cmdline[0..nb], 0);

    if (arg0.len < 2)
        return false;

    if (!mem.eql(u8, arg0[arg0.len - 2 .. arg0.len], "id"))
        return false;

    if (arg0.len == 2)
        return true;

    return arg0[arg0.len - 3] == '/';
}

const testing = std.testing;

test "libnss getpwuid_r and getpwnam_r" {
    var tf = try File.TestDB.init(testing.allocator);
    defer tf.deinit();
    const state = State{
        .file = tf.file,
        .omit_members = false,
        .initgroups_dyn_allocator = testing.failing_allocator,
    };

    var buf: [1024]u8 = undefined;
    var errno: c_int = 0;

    var passwd: CUser = undefined;
    try testing.expectEqual(
        c.NSS_STATUS_SUCCESS,
        getpwuid_r(&state.file.db, 128, &passwd, &buf, buf.len, &errno),
    );
    try testing.expectEqual(@as(c_int, 0), errno);
    try testVidmantas(passwd);

    passwd = undefined;
    try testing.expectEqual(
        c.NSS_STATUS_SUCCESS,
        getpwnam_r(&state.file.db, "vidmantas", &passwd, &buf, buf.len, &errno),
    );
    try testing.expectEqual(@as(c_int, 0), errno);
    try testVidmantas(passwd);

    passwd = undefined;
    try testing.expectEqual(
        c.NSS_STATUS_NOTFOUND,
        getpwnam_r(&state.file.db, "does not exist", &passwd, &buf, buf.len, &errno),
    );
    try testing.expectEqual(@enumToInt(os.E.NOENT), @intCast(u16, errno));

    passwd = undefined;
    var small_buffer: [1]u8 = undefined;
    try testing.expectEqual(
        c.NSS_STATUS_TRYAGAIN,
        getpwuid_r(&state.file.db, 0, &passwd, &small_buffer, small_buffer.len, &errno),
    );
    try testing.expectEqual(@enumToInt(os.E.RANGE), @intCast(u16, errno));
}

fn testVidmantas(u: CUser) !void {
    try testing.expectEqualStrings("vidmantas", mem.sliceTo(u.pw_name, 0));
    try testing.expectEqual(@as(u32, 128), u.pw_uid);
    try testing.expectEqual(@as(u32, 128), u.pw_gid);
    try testing.expectEqualStrings("Vidmantas Kaminskas", mem.sliceTo(u.pw_gecos, 0));
    try testing.expectEqualStrings("/bin/bash", mem.sliceTo(u.pw_shell, 0));
}

test "libnss getgrgid_r and getgrnam_r" {
    var tf = try File.TestDB.init(testing.allocator);
    defer tf.deinit();

    const state = State{
        .file = tf.file,
        .omit_members = false,
        .initgroups_dyn_allocator = testing.failing_allocator,
    };

    var buf: [1024]u8 = undefined;
    var errno: c_int = 0;
    var group: CGroup = undefined;
    try testing.expectEqual(
        c.NSS_STATUS_SUCCESS,
        getgrgid_r(&state, 100000, &group, &buf, buf.len, &errno),
    );
    try testing.expectEqual(@as(c_int, 0), errno);
    try testSvcGroup(group);

    group = undefined;
    try testing.expectEqual(
        c.NSS_STATUS_SUCCESS,
        getgrnam_r(&state, "service-group", &group, &buf, buf.len, &errno),
    );
    try testing.expectEqual(@as(c_int, 0), errno);
    try testSvcGroup(group);

    group = undefined;
    try testing.expectEqual(
        c.NSS_STATUS_NOTFOUND,
        getgrnam_r(&state, "does not exist", &group, &buf, buf.len, &errno),
    );
    try testing.expectEqual(@enumToInt(os.E.NOENT), @intCast(u16, errno));
}

fn testSvcGroup(g: CGroup) !void {
    try testing.expectEqual(@as(u32, 100000), g.gr_gid);
    try testing.expectEqualStrings("service-group", mem.sliceTo(g.gr_name, 0));
    const members = g.gr_mem;
    try testing.expect(members[0] != null);
    try testing.expect(members[1] != null);
    try testing.expectEqualStrings("root", mem.sliceTo(members[0].?, 0));
    try testing.expectEqualStrings("vidmantas", mem.sliceTo(members[1].?, 0));
    try testing.expect(members[2] == null);
}

test "libnss initgroups_dyn" {
    const allocator = testing.allocator;

    var tf = try File.TestDB.init(allocator);
    defer tf.deinit();
    var state = State{
        .file = tf.file,
        .omit_members = false,
        .initgroups_dyn_allocator = allocator,
    };

    var size: c_long = 3; // size of the gids array is this
    var groups = try allocator.alloc(u32, @intCast(usize, size)); // buffer too small
    defer allocator.free(groups);
    var errno: c_int = 42; // canary
    groups[0] = 42; // canary
    var start: c_long = 1; // canary is there

    const status = initgroups_dyn(
        &state,
        "vidmantas",
        128, // gid
        &start,
        &size,
        &groups.ptr,
        3, // limit: besides canary accept 2 groups
        &errno,
    );
    try testing.expectEqual(c.NSS_STATUS_SUCCESS, status);
    try testing.expectEqual(@as(c_int, 42), errno);
    try testing.expectEqual(@as(u32, 42), groups[0]);
    try testing.expectEqual(@as(u32, 9999), groups[1]);
    try testing.expectEqual(@as(u32, 100000), groups[2]);
}

test "libnss isId" {
    const examples = [_]struct {
        cmdline: []const u8,
        want: bool,
    }{
        .{ .want = true, .cmdline = "id\x00foo" },
        .{ .want = true, .cmdline = "/usr/bin/id\x00foo" },
        .{ .want = true, .cmdline = "id" },
        .{ .want = true, .cmdline = "/usr/bin/id" },
        .{ .want = false, .cmdline = "" },
        .{ .want = false, .cmdline = "d\x00" },
        .{ .want = false, .cmdline = "/fakeid" },
    };

    for (examples) |tt| {
        const fd = try os.memfd_create("test_libnss_isId", 0);
        defer os.close(fd);
        _ = try os.write(fd, tt.cmdline);
        try os.lseek_SET(fd, 0);
        const got = isId(fd);
        try testing.expectEqual(tt.want, got);
    }
}