const std = @import("std");
const info = std.log.info;
const debug = std.log.debug;
const print = std.debug.print;
const eql = std.mem.eql;
const Allocator = std.mem.Allocator;
const testing = std.testing;
const expect = testing.expect;
const ArrayList = std.ArrayList;
const Base32 = @import("base32.zig").Base32;
const Base32Error = @import("base32.zig").Base32Error;

// otpauth://totp/AWS+Dev?secret=47STA47VFCMMLLWOLHWO3KY7MYNC36MLCDTHOLIYKJCTTSSAMKVM7YA3VWT2AJEP&digits=6&icon=Amazon
// otpauth://totp/Gitea%20%28git.lucares.de%29:andi?algorithm=SHA1&digits=6&issuer=Gitea&period=30&secret=MSP53Q672UJMSCLQVRCJKMMZKK7MWSMYFL77OQ24JBM65RZWY7F2Y45FMTNLYM36

pub fn main() !void {
    mainInternal() catch |err| {
        switch (err) {
            Base32Error.InvalidCharacter, Base32Error.InvalidPadding => {
                try std.io.getStdErr().writer().print("The secret is invalid.\n", .{});
            },
            ArgumentError.UnknownParameter => {
                // do nothing, error message is already written (because the message contains the name of the unknown parameter)
            },
            ArgumentError.InvalidOtpAuthUrl => {},
            else => |leftover_err| {
                try std.io.getStdErr().writer().print("{?}\n", .{leftover_err});
            },
        }
        std.process.exit(1);
    };
}

fn mainInternal() !void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const args = try std.process.argsAlloc(allocator);

    const arg: Args = try parseArgs(allocator, args);
    //print("parsed Args: {?any}\n", arg);
    if (arg.list) {
        const names = try executeGetList(allocator, arg.config_location);

        for (0..names.items.len) |i| {
            try std.io.getStdOut().writer().print("{s}\n", .{names.items[i]});
        }
    } else if (arg.show != null) {
        const authenticator = try getAuthenticator(allocator, arg.show.?, arg.config_location);
        const code = try authenticator.code(allocator, std.time.timestamp);
        try std.io.getStdOut().writer().print("{s}\n", .{code});
    } else {
        printHelp();
    }
}

fn printHelp() void {
    const msg =
        \\Usage: zig-totp command [options]
        \\
        \\Commands:
        \\  list                   List the configured authentiators
        \\  show NAME              Show the code for the authenticator with name NAME
        \\
        \\Options:
        \\  --config path          The path to a config file (default is $XDG_CONFIG_HOME/zig-totp or $HOME/.zig-totp).
        \\ 
    ;
    std.debug.print("{s}", .{msg});
}

const Args = struct {
    list: bool,
    show: ?[]const u8,
    config_location: []const u8,
};

const OtpAuthUrl = struct { name: []const u8, secretEncoded: []const u8, url: []const u8, period: u32, digits: u4 };

const Authenticator = struct {
    url: OtpAuthUrl,

    /// algorithm based on https://www.rfc-editor.org/rfc/rfc6238
    pub fn code(self: Authenticator, allocator: Allocator, timeFunc: *const fn () i64) ![]const u8 {
        if (self.url.name.len > 0) {}

        const secret = try Base32.decodeU8(allocator, self.url.secretEncoded);
        defer allocator.free(secret);
        if (false) {
            debug("secret: {s}\n", .{secret});
        }

        const intervalNumber = @divTrunc(timeFunc(), self.url.period);

        var intervalAsU8Array: [8]u8 = undefined;
        std.mem.writePackedInt(i64, intervalAsU8Array[0..], 0, intervalNumber, .big);
        //debug("interval packed: {X}\n", .{intervalAsU8Array});

        //debug("intervalNumber: {d}\n", .{intervalNumber});

        var out: [std.crypto.auth.hmac.HmacSha1.mac_length]u8 = undefined;
        std.crypto.auth.hmac.HmacSha1.create(out[0..], &intervalAsU8Array, secret);
        //debug("hmac: {X}\n", .{out});

        // take the 4 least significant bits of the hash and use them as byte offset
        const leastSignificantByte = out[std.crypto.auth.hmac.HmacSha1.mac_length - 1];
        const byteIndex = leastSignificantByte & 0b1111;
        //debug("index: {d}\n", .{byteIndex});

        const x: [4]u8 = [4]u8{ out[byteIndex], out[byteIndex + 1], out[byteIndex + 2], out[byteIndex + 3] };

        const tokenBase = std.mem.readInt(i32, &x, .big) & 0x7fffffff;
        //debug("tokenBase: {d}\n", .{tokenBase});
        const token = @mod(tokenBase, (std.math.pow(i64, 10, self.url.digits)));

        const result = try zeroPad(allocator, self.url.digits, token);
        //debug("code as 0-padded string: {s} ({d})\n", .{ result, token });
        return result;
    }
};

fn zeroPad(allocator: Allocator, digits: u4, x: anytype) ![]u8 {
    const result = try allocator.alloc(u8, digits);

    const s = try std.fmt.allocPrint(allocator, "{d}", .{x});
    defer allocator.free(s);

    var i: usize = 0;
    var j: usize = 0;

    while (i < digits) {
        if (j < s.len) {
            result[@as(usize, digits) - i - 1] = s[s.len - j - 1];
        } else {
            result[@as(usize, digits) - i - 1] = '0';
        }

        j += 1;
        i += 1;
    }

    return result;
}

const ArgumentError = error{ InvalidOtpAuthUrl, UnknownParameter, MissingConfigLocation, AuthenticatorNotFound };

fn parseArgs(allocator: Allocator, args: []const []const u8) !Args {
    var result = Args{
        .list = false,
        .show = null,
        .config_location = try configLocation(allocator),
    };
    var i: u17 = 1;
    while (i < args.len) : (i += 1) {
        const arg = args[i];

        if (eql(u8, "--config", arg)) {
            i += 1;
            if (i >= args.len) {
                return error.MissingConfigLocation;
            }
            result.config_location = args[i];
        } else if (eql(u8, "list", arg)) {
            result.list = true;
        } else if (eql(u8, "show", arg)) {
            i += 1;
            if (i >= args.len) {
                return error.MissingConfigLocation;
            }
            result.show = args[i];
        } else {
            std.debug.print("unknown argument: {s}\n", .{arg});
            return ArgumentError.UnknownParameter;
        }
    }
    return result;
}

fn parseOtpAuthUrl(allocator: Allocator, url: []const u8) !OtpAuthUrl {
    const uri = try std.Uri.parse(url);

    if (!eql(u8, uri.scheme, "otpauth")) {
        return error.InvalidOtpAuthUrl;
    }

    if (uri.query == null) {
        return error.InvalidOtpAuthUrl;
    }

    var name = try uri.path.toRawMaybeAlloc(allocator);
    name = std.mem.trimLeft(u8, name, "/");

    const params = try parseQueryString(allocator, try uri.query.?.toRawMaybeAlloc(allocator));

    return OtpAuthUrl{
        .name = name,
        .secretEncoded = params.get("secret").?,
        .url = url,
        .period = if (params.get("period") != null) try std.fmt.parseInt(u32, params.get("period").?, 10) else 30,
        .digits = if (params.get("digits") != null) try std.fmt.parseInt(u4, params.get("digits").?, 10) else 6,
    };
}

fn parseQueryString(allocator: Allocator, query: []const u8) !std.StringHashMap([]const u8) {
    var map = std.StringHashMap([]const u8).init(allocator);

    var it = std.mem.splitSequence(u8, query, "&");

    while (it.peek() != null) {
        const next = it.next().?;

        var paramIt = std.mem.splitSequence(u8, next, "=");
        const name = paramIt.next();
        const value = paramIt.next();

        if (name) |n| {
            if (value) |v| {
                try map.put(n, v);
            } else {
                try map.put(n, "");
            }
        }
    }

    return map;
}

fn configLocation(allocator: Allocator) ![]const u8 {
    const xdg_config_home: ?[]const u8 = std.process.getEnvVarOwned(allocator, "XDG_CONFIG_HOME") catch null;

    if (xdg_config_home) |base| {
        const config_location = try std.mem.concat(allocator, u8, &[_][]const u8{ base, "/zig-totp" });
        //debug("config_location: {s}", .{config_location});
        return config_location;
    }

    const home = std.process.getEnvVarOwned(allocator, "HOME") catch null;
    const base = home orelse unreachable;
    const config_location = try std.mem.concat(allocator, u8, &[_][]const u8{ base, "/.zig-totp" });

    //debug("config_location: {s}", .{config_location});
    return config_location;
}

/// Reads the config file into a list of OtpAuthUrl.
/// The caller is responsible to free the memory.
fn read_config(allocator: Allocator, config_location: []const u8) !ArrayList(OtpAuthUrl) {
    const file = try std.fs.cwd().openFile(config_location, .{});
    defer file.close();

    var buf_reader = std.io.bufferedReader(file.reader());
    var in_stream = buf_reader.reader();

    var authenticators = std.ArrayList(OtpAuthUrl).init(allocator);

    var line_no: usize = 1;
    while (try in_stream.readUntilDelimiterOrEofAlloc(allocator, '\n', 1024 * 1024)) |line| {
        //defer allocator.free(line);
        if (line.len > 0 and std.mem.trim(u8, line, " \r").len > 0) {
            const authenticatorResult = parseOtpAuthUrl(allocator, line);
            if (authenticatorResult) |authenticator| {
                try authenticators.append(authenticator);
            } else |err| {
                switch (err) {
                    error.InvalidOtpAuthUrl => {
                        try std.io.getStdErr().writer().print("invalid otpauth url in line {d}\n", .{line_no});
                    },
                    else => {
                        return err;
                    },
                }
            }
        }
        line_no += 1;
    }
    return authenticators;
}

fn executeGetList(allocator: Allocator, config_location: []const u8) !std.ArrayList([]const u8) {
    var names = std.ArrayList([]const u8).init(allocator);

    const authenticators = try read_config(allocator, config_location);
    defer authenticators.deinit();

    for (0..authenticators.items.len) |i| {
        try names.append(authenticators.items[i].name);
    }

    return names;
}

fn getAuthenticator(allocator: Allocator, name: []const u8, config_location: []const u8) !Authenticator {
    const otpAuthUrls = try read_config(allocator, config_location);

    for (0..otpAuthUrls.items.len) |i| {
        const url: OtpAuthUrl = otpAuthUrls.items[i];
        if (eql(u8, url.name, name)) {
            return Authenticator{ .url = url };
        }
    }
    return error.AuthenticatorNotFound;
}

test "parse command line parameter: 'list'" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const arg = try parseArgs(allocator, &[_][]const u8{ "path/of/executable", "list" });
    try expect(arg.list);
}

test "read list of entries" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    std.fs.cwd().deleteTree("test-tmp") catch unreachable;
    try std.fs.cwd().makeDir("test-tmp");
    const file = try std.fs.cwd().createFile("test-tmp/zig-totp", .{ .read = true });
    defer {
        file.close();
        std.fs.cwd().deleteTree("test-tmp") catch unreachable;
    }

    _ = try file.write("otpauth://totp/token1?secret=c2VjcmV0Cg==\notpauth://totp/token2?secret=c2VjcmV0Cg==");

    const list: std.ArrayList([]const u8) = try executeGetList(allocator, "test-tmp/zig-totp");
    try std.testing.expectEqualStrings("token1", list.items[0]);
    try std.testing.expectEqualStrings("token2", list.items[1]);
    try std.testing.expectEqual(2, list.items.len);
}

test "read list of entries, ignoring invalid otpauth urls" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    std.fs.cwd().deleteTree("test-tmp") catch unreachable;
    try std.fs.cwd().makeDir("test-tmp");
    const file = try std.fs.cwd().createFile("test-tmp/zig-totp", .{ .read = true });
    defer {
        file.close();
        std.fs.cwd().deleteTree("test-tmp") catch unreachable;
    }

    _ = try file.write("otpauth://totp/token1?secret=c2VjcmV0Cg==\notpauth://invalid\n");

    const list: std.ArrayList([]const u8) = try executeGetList(allocator, "test-tmp/zig-totp");

    try std.testing.expectEqualStrings("token1", list.items[0]);
    try std.testing.expectEqual(1, list.items.len);
}

test "parse oth pauth url" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const actual: OtpAuthUrl = try parseOtpAuthUrl(allocator, "otpauth://totp/foo?secret=MFQQ&period=31&digits=7");

    try std.testing.expectEqualStrings("MFQQ", actual.secretEncoded);
    try std.testing.expectEqualStrings("foo", actual.name);
    try std.testing.expectEqual(31, actual.period);
    try std.testing.expectEqual(7, actual.digits);
}

test "zero padding" {
    const actual = try zeroPad(std.testing.allocator, 6, 123);
    defer std.testing.allocator.free(actual);

    try std.testing.expectEqualStrings("000123", actual);
}

test "authenticator generate code with 6 digits" {
    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "TestAuthenticator", .secretEncoded = "HJ5PX4IFQKD37HFNXLJYBAVD5G6NMMUOKUCDXJ4XTLUUBWSXRXEN5SR4MLAZA5M2", .url = "not needed", .period = 30, .digits = 6 } };

    const code = try authenticator.code(std.testing.allocator, _closure_return_1725695340);
    defer std.testing.allocator.free(code);
    try std.testing.expectEqualStrings("218139", code);
}
test "authenticator generate code with 10 digits and zero padding" {
    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "TestAuthenticator", .secretEncoded = "HJ5PX4IFQKD37HFNXLJYBAVD5G6NMMUOKUCDXJ4XTLUUBWSXRXEN5SR4MLAZA5M2", .url = "not needed", .period = 60, .digits = 10 } };

    const code = try authenticator.code(std.testing.allocator, _closure_return_1725695340);
    defer std.testing.allocator.free(code);
    try std.testing.expectEqualStrings("0844221464", code);
}

test "testcases from https://www.rfc-editor.org/rfc/rfc6238#appendix-A" {
    const secretForSha1 = "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"; // plain: "12345678901234567890"  hex: "3132333435363738393031323334353637383930"

    try testTOTP(OtpAuthUrl{ .name = "", .secretEncoded = secretForSha1, .url = "", .period = 30, .digits = 8 }, _closure_return_59, "94287082");

    try testTOTP(OtpAuthUrl{ .name = "", .secretEncoded = secretForSha1, .url = "", .period = 30, .digits = 8 }, _closure_return_1111111109, "07081804");

    try testTOTP(OtpAuthUrl{ .name = "", .secretEncoded = secretForSha1, .url = "", .period = 30, .digits = 8 }, _closure_return_1111111111, "14050471");

    try testTOTP(OtpAuthUrl{ .name = "", .secretEncoded = secretForSha1, .url = "", .period = 30, .digits = 8 }, _closure_return_1234567890, "89005924");

    try testTOTP(OtpAuthUrl{ .name = "", .secretEncoded = secretForSha1, .url = "", .period = 30, .digits = 8 }, _closure_return_2000000000, "69279037");

    try testTOTP(OtpAuthUrl{ .name = "", .secretEncoded = secretForSha1, .url = "", .period = 30, .digits = 8 }, _closure_return_20000000000, "65353130");
}

fn testTOTP(otpAuthUrl: OtpAuthUrl, timeFunc: *const fn () i64, expected: []const u8) !void {
    const authenticator = Authenticator{ .url = otpAuthUrl };

    const code = try authenticator.code(std.testing.allocator, timeFunc);
    defer std.testing.allocator.free(code);
    try std.testing.expectEqualStrings(expected, code);
}

fn _closure_return_1725695340() i64 {
    return 1725695340;
}

fn _closure_return_59() i64 {
    return 59;
}

fn _closure_return_1111111109() i64 {
    return 1111111109;
}

fn _closure_return_1111111111() i64 {
    return 1111111111;
}

fn _closure_return_1234567890() i64 {
    return 1234567890;
}

fn _closure_return_2000000000() i64 {
    return 2000000000;
}

fn _closure_return_20000000000() i64 {
    return 20000000000;
}

test "secret is not Base32 - invalid padding" {
    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "", .secretEncoded = "H", .url = "", .period = 60, .digits = 6 } };

    const actualError = authenticator.code(std.testing.allocator, _closure_return_1725695340);
    try std.testing.expectError(Base32Error.InvalidPadding, actualError);
}
test "secret is not Base32 - invalid padding, input has too many bytes" {
    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "", .secretEncoded = "HAX", .url = "", .period = 60, .digits = 6 } };

    const actualError = authenticator.code(std.testing.allocator, _closure_return_1725695340);
    try std.testing.expectError(Base32Error.InvalidPadding, actualError);
}
test "secret is not Base32 - invalid padding, trailing bits are not 0" {
    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "", .secretEncoded = "HX", .url = "", .period = 60, .digits = 6 } };

    const actualError = authenticator.code(std.testing.allocator, _closure_return_1725695340);
    try std.testing.expectError(Base32Error.InvalidPadding, actualError);
}

test "secret is not Base32 - invalid character" {
    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "", .secretEncoded = "0", .url = "", .period = 60, .digits = 6 } };

    const actualError = authenticator.code(std.testing.allocator, _closure_return_1725695340);
    try std.testing.expectError(Base32Error.InvalidCharacter, actualError);
}