const std = @import("std");
const clock = @import("../core/clock.zig");
const config = @import("../core/config.zig");
const zat = @import("zat");

pub const Account = struct {
    handle: []const u8,
    did: []const u8,
    email: []const u8,
    password: []const u8,
};

pub const TokenClaims = struct {
    did: []const u8,
    scope: []const u8,
    jti: []const u8,
    cnf_jkt: ?[]const u8 = null,
    exp: i64,
};

var jwt_counter: usize = 0;

pub fn passwordMatches(account: Account, password: []const u8) bool {
    if (passwordHashMatches(account.password, password)) return true;
    return std.mem.eql(u8, password, account.password);
}

pub fn hashPassword(allocator: std.mem.Allocator, password: []const u8, salt: [16]u8) ![]const u8 {
    var derived: [32]u8 = undefined;
    try std.crypto.pwhash.pbkdf2(&derived, password, &salt, password_hash_rounds, std.crypto.auth.hmac.sha2.HmacSha256);
    const salt_hex = std.fmt.bytesToHex(salt, .lower);
    const derived_hex = std.fmt.bytesToHex(derived, .lower);
    return std.fmt.allocPrint(allocator, "pbkdf2-sha256:{d}:{s}:{s}", .{ password_hash_rounds, &salt_hex, &derived_hex });
}

pub fn passwordHashMatches(stored: []const u8, password: []const u8) bool {
    var parts = std.mem.splitScalar(u8, stored, ':');
    const scheme = parts.next() orelse return false;
    if (!std.mem.eql(u8, scheme, "pbkdf2-sha256")) return false;
    const rounds_text = parts.next() orelse return false;
    const salt_hex = parts.next() orelse return false;
    const derived_hex = parts.next() orelse return false;
    if (parts.next() != null or salt_hex.len != 32 or derived_hex.len != 64) return false;
    const rounds = std.fmt.parseInt(u32, rounds_text, 10) catch return false;
    var salt: [16]u8 = undefined;
    _ = std.fmt.hexToBytes(&salt, salt_hex) catch return false;
    var expected: [32]u8 = undefined;
    _ = std.fmt.hexToBytes(&expected, derived_hex) catch return false;
    var actual: [32]u8 = undefined;
    std.crypto.pwhash.pbkdf2(&actual, password, &salt, rounds, std.crypto.auth.hmac.sha2.HmacSha256) catch return false;
    return std.crypto.timing_safe.eql([32]u8, actual, expected);
}

const password_hash_rounds: u32 = 100_000;

pub const SessionToken = struct {
    token: []const u8,
    jti: []const u8,
    iat: i64,
    exp: i64,
};

pub const SessionPair = struct {
    access: SessionToken,
    refresh: SessionToken,
};

pub fn createSessionJwt(allocator: std.mem.Allocator, kind: []const u8, account: Account) ![]const u8 {
    const token = try createSessionToken(allocator, kind, account);
    allocator.free(token.jti);
    return token.token;
}

pub fn createDpopSessionJwt(allocator: std.mem.Allocator, kind: []const u8, account: Account, dpop_jkt: []const u8) ![]const u8 {
    const token = try createSessionTokenWithScopeAndCnf(allocator, account, kind, tryScope(kind), dpop_jkt);
    allocator.free(token.jti);
    return token.token;
}

pub fn createSessionPair(allocator: std.mem.Allocator, account: Account) !SessionPair {
    return createSessionPairWithAccessScope(allocator, account, "com.atproto.access");
}

pub fn createSessionPairWithAccessScope(allocator: std.mem.Allocator, account: Account, access_scope: []const u8) !SessionPair {
    const access = try createSessionTokenWithScope(allocator, account, "access", access_scope);
    errdefer {
        allocator.free(access.token);
        allocator.free(access.jti);
    }
    const refresh = try createSessionTokenWithScope(allocator, account, "refresh", "com.atproto.refresh");
    return .{ .access = access, .refresh = refresh };
}

pub fn createSessionToken(allocator: std.mem.Allocator, kind: []const u8, account: Account) !SessionToken {
    const scope = try std.fmt.allocPrint(allocator, "com.atproto.{s}", .{kind});
    defer allocator.free(scope);
    return createSessionTokenWithScope(allocator, account, kind, scope);
}

fn createSessionTokenWithScope(allocator: std.mem.Allocator, account: Account, kind: []const u8, scope: []const u8) !SessionToken {
    return createSessionTokenWithScopeAndCnf(allocator, account, kind, scope, null);
}

fn createSessionTokenWithScopeAndCnf(
    allocator: std.mem.Allocator,
    account: Account,
    kind: []const u8,
    scope: []const u8,
    dpop_jkt: ?[]const u8,
) !SessionToken {
    return createSessionTokenAt(allocator, account, kind, scope, dpop_jkt, clock.nowNanos());
}

fn createSessionTokenAt(
    allocator: std.mem.Allocator,
    account: Account,
    kind: []const u8,
    scope: []const u8,
    dpop_jkt: ?[]const u8,
    issued_at_ns: i96,
) !SessionToken {
    const jti = @atomicRmw(usize, &jwt_counter, .Add, 1, .monotonic);
    const iat: i64 = @intCast(@divTrunc(issued_at_ns, std.time.ns_per_s));
    const nsec: i64 = @intCast(@mod(issued_at_ns, std.time.ns_per_s));
    const lifetime: i64 = if (std.mem.eql(u8, kind, "refresh")) 14 * 24 * 60 * 60 else 2 * 60 * 60;
    const jti_text = try std.fmt.allocPrint(allocator, "zds-{d}-{d}-{d}", .{ iat, nsec, jti });
    return signSessionToken(allocator, account, scope, dpop_jkt, jti_text, iat, iat + lifetime);
}

/// mints a session pair that names an existing row's jtis and expiries, so a
/// refresh replayed inside the rotation grace window hands back the same
/// successor the winning request got
pub fn reissueSessionPair(
    allocator: std.mem.Allocator,
    account: Account,
    access_scope: []const u8,
    access_jti: []const u8,
    access_exp: i64,
    refresh_jti: []const u8,
    refresh_exp: i64,
) !SessionPair {
    const iat: i64 = @intCast(@divTrunc(clock.nowNanos(), std.time.ns_per_s));
    const access = try signSessionToken(allocator, account, access_scope, null, try allocator.dupe(u8, access_jti), iat, access_exp);
    errdefer {
        allocator.free(access.token);
        allocator.free(access.jti);
    }
    const refresh = try signSessionToken(allocator, account, "com.atproto.refresh", null, try allocator.dupe(u8, refresh_jti), iat, refresh_exp);
    return .{ .access = access, .refresh = refresh };
}

/// takes ownership of `jti_text`
fn signSessionToken(
    allocator: std.mem.Allocator,
    account: Account,
    scope: []const u8,
    dpop_jkt: ?[]const u8,
    jti_text: []const u8,
    iat: i64,
    exp: i64,
) !SessionToken {
    errdefer allocator.free(jti_text);
    const header = "{\"alg\":\"HS256\",\"typ\":\"JWT\"}";
    const payload = if (dpop_jkt) |jkt|
        try std.fmt.allocPrint(
            allocator,
            "{{\"sub\":{f},\"scope\":{f},\"iat\":{d},\"exp\":{d},\"jti\":{f},\"cnf\":{{\"jkt\":{f}}}}}",
            .{ std.json.fmt(account.did, .{}), std.json.fmt(scope, .{}), iat, exp, std.json.fmt(jti_text, .{}), std.json.fmt(jkt, .{}) },
        )
    else
        try std.fmt.allocPrint(
            allocator,
            "{{\"sub\":{f},\"scope\":{f},\"iat\":{d},\"exp\":{d},\"jti\":{f}}}",
            .{ std.json.fmt(account.did, .{}), std.json.fmt(scope, .{}), iat, exp, std.json.fmt(jti_text, .{}) },
        );
    defer allocator.free(payload);

    const header_b64 = try zat.jwt.base64UrlEncode(allocator, header);
    defer allocator.free(header_b64);
    const payload_b64 = try zat.jwt.base64UrlEncode(allocator, payload);
    defer allocator.free(payload_b64);

    const signing_input = try std.fmt.allocPrint(allocator, "{s}.{s}", .{ header_b64, payload_b64 });
    defer allocator.free(signing_input);
    var signature: [std.crypto.auth.hmac.sha2.HmacSha256.mac_length]u8 = undefined;
    std.crypto.auth.hmac.sha2.HmacSha256.create(&signature, signing_input, config.jwtSecret());
    const signature_b64 = try zat.jwt.base64UrlEncode(allocator, &signature);
    defer allocator.free(signature_b64);
    return .{
        .token = try std.fmt.allocPrint(allocator, "{s}.{s}", .{ signing_input, signature_b64 }),
        .jti = jti_text,
        .iat = iat,
        .exp = exp,
    };
}

fn tryScope(kind: []const u8) []const u8 {
    if (std.mem.eql(u8, kind, "access")) return "com.atproto.access";
    if (std.mem.eql(u8, kind, "refresh")) return "com.atproto.refresh";
    return "com.atproto.access";
}

pub fn createServiceJwt(
    allocator: std.mem.Allocator,
    account: Account,
    audience: []const u8,
    lexicon_method: ?[]const u8,
) ![]const u8 {
    const keypair = try serviceKeypair();
    return createServiceJwtWithKeypair(allocator, account, audience, lexicon_method, null, &keypair);
}

pub fn createServiceJwtWithKeypair(
    allocator: std.mem.Allocator,
    account: Account,
    audience: []const u8,
    lexicon_method: ?[]const u8,
    requested_exp: ?i64,
    keypair: *const zat.Keypair,
) ![]const u8 {
    const header = try std.fmt.allocPrint(
        allocator,
        "{{\"typ\":\"JWT\",\"alg\":\"{s}\"}}",
        .{@tagName(keypair.algorithm())},
    );
    defer allocator.free(header);
    const iat = unixNow();
    const exp = requested_exp orelse iat + 60;
    const jti = @atomicRmw(usize, &jwt_counter, .Add, 1, .monotonic);

    const payload = if (lexicon_method) |lxm|
        try std.fmt.allocPrint(
            allocator,
            "{{\"iss\":\"{s}\",\"sub\":\"{s}\",\"aud\":\"{s}\",\"exp\":{d},\"iat\":{d},\"jti\":\"zds-{d}-{d}\",\"lxm\":\"{s}\"}}",
            .{ account.did, account.did, audience, exp, iat, iat, jti, lxm },
        )
    else
        try std.fmt.allocPrint(
            allocator,
            "{{\"iss\":\"{s}\",\"sub\":\"{s}\",\"aud\":\"{s}\",\"exp\":{d},\"iat\":{d},\"jti\":\"zds-{d}-{d}\"}}",
            .{ account.did, account.did, audience, exp, iat, iat, jti },
        );
    defer allocator.free(payload);

    return zat.oauth.createJwt(allocator, header, payload, keypair);
}

fn unixNow() i64 {
    return clock.nowSeconds();
}

pub fn serviceKeypair() !zat.Keypair {
    return zat.Keypair.fromSecretKey(.p256, .{
        0x01, 0x7a, 0x64, 0x73, 0x2d, 0x64, 0x65, 0x76,
        0x2d, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65,
        0x2d, 0x61, 0x75, 0x74, 0x68, 0x2d, 0x6b, 0x65,
        0x79, 0x2d, 0x30, 0x30, 0x30, 0x30, 0x30, 0x31,
    });
}

pub const SessionJwtError = error{ Invalid, Expired };

pub fn claimsFromSessionJwt(allocator: std.mem.Allocator, token: []const u8) ?TokenClaims {
    return parseSessionJwt(allocator, token) catch null;
}

/// verifies the signature before looking at `exp`, so `Expired` is only ever
/// reported for a token this server minted
pub fn parseSessionJwt(allocator: std.mem.Allocator, token: []const u8) SessionJwtError!TokenClaims {
    var parts = std.mem.splitScalar(u8, token, '.');
    const header_part = parts.next() orelse return error.Invalid;
    const payload_part = parts.next() orelse return error.Invalid;
    const signature_part = parts.next() orelse return error.Invalid;
    if (payload_part.len == 0 or signature_part.len == 0 or parts.next() != null) return error.Invalid;

    const signing_input_len = header_part.len + 1 + payload_part.len;
    if (token.len < signing_input_len or token[header_part.len] != '.') return error.Invalid;
    const signing_input = token[0..signing_input_len];

    const header_json = zat.jwt.base64UrlDecode(allocator, header_part) catch return error.Invalid;
    defer allocator.free(header_json);
    const header = std.json.parseFromSlice(std.json.Value, allocator, header_json, .{}) catch return error.Invalid;
    defer header.deinit();
    const alg = zat.json.getString(header.value, "alg") orelse return error.Invalid;
    if (!std.mem.eql(u8, alg, "HS256")) return error.Invalid;

    const signature = zat.jwt.base64UrlDecode(allocator, signature_part) catch return error.Invalid;
    defer allocator.free(signature);
    if (signature.len != std.crypto.auth.hmac.sha2.HmacSha256.mac_length) return error.Invalid;
    var expected: [std.crypto.auth.hmac.sha2.HmacSha256.mac_length]u8 = undefined;
    std.crypto.auth.hmac.sha2.HmacSha256.create(&expected, signing_input, config.jwtSecret());
    var actual: [std.crypto.auth.hmac.sha2.HmacSha256.mac_length]u8 = undefined;
    @memcpy(&actual, signature);
    if (!std.crypto.timing_safe.eql([std.crypto.auth.hmac.sha2.HmacSha256.mac_length]u8, actual, expected)) return error.Invalid;

    const payload = zat.jwt.base64UrlDecode(allocator, payload_part) catch return error.Invalid;
    defer allocator.free(payload);

    const parsed = std.json.parseFromSlice(std.json.Value, allocator, payload, .{}) catch return error.Invalid;
    defer parsed.deinit();
    const subject = zat.json.getString(parsed.value, "sub") orelse return error.Invalid;
    if (zat.Did.parse(subject) == null) return error.Invalid;
    const scope = zat.json.getString(parsed.value, "scope") orelse return error.Invalid;
    if (!std.mem.startsWith(u8, scope, "com.atproto.")) return error.Invalid;
    const jti = zat.json.getString(parsed.value, "jti") orelse return error.Invalid;
    const exp = zat.json.getInt(parsed.value, "exp") orelse return error.Invalid;
    if (exp <= unixNow()) return error.Expired;
    const cnf_jkt = if (zat.json.getPath(parsed.value, "cnf.jkt")) |jkt_value| switch (jkt_value) {
        .string => |jkt| allocator.dupe(u8, jkt) catch return error.Invalid,
        else => return error.Invalid,
    } else null;
    return .{
        .did = allocator.dupe(u8, subject) catch return error.Invalid,
        .scope = allocator.dupe(u8, scope) catch return error.Invalid,
        .jti = allocator.dupe(u8, jti) catch return error.Invalid,
        .cnf_jkt = cnf_jkt,
        .exp = exp,
    };
}

test "writes jwt-shaped session token" {
    const alice = testAccount();
    const token = try createSessionJwt(std.testing.allocator, "access", alice);
    defer std.testing.allocator.free(token);
    try std.testing.expect(std.mem.count(u8, token, ".") == 2);
    const claims = claimsFromSessionJwt(std.testing.allocator, token).?;
    defer std.testing.allocator.free(claims.did);
    defer std.testing.allocator.free(claims.scope);
    defer std.testing.allocator.free(claims.jti);
    defer if (claims.cnf_jkt) |jkt| std.testing.allocator.free(jkt);
    try std.testing.expectEqualStrings(alice.did, claims.did);
}

test "expired session jwt parses as Expired, not Invalid" {
    const alice = testAccount();
    const three_hours_ago: i96 = clock.nowNanos() - 3 * 60 * 60 * std.time.ns_per_s;
    const stale = try createSessionTokenAt(std.testing.allocator, alice, "access", "com.atproto.access", null, three_hours_ago);
    defer std.testing.allocator.free(stale.token);
    defer std.testing.allocator.free(stale.jti);
    try std.testing.expectError(error.Expired, parseSessionJwt(std.testing.allocator, stale.token));
    try std.testing.expect(claimsFromSessionJwt(std.testing.allocator, stale.token) == null);

    const forged = try std.fmt.allocPrint(std.testing.allocator, "{s}x", .{stale.token});
    defer std.testing.allocator.free(forged);
    try std.testing.expectError(error.Invalid, parseSessionJwt(std.testing.allocator, forged));
}

test "reissued session pair keeps the row's jtis and expiries" {
    const alice = testAccount();
    const pair = try reissueSessionPair(std.testing.allocator, alice, "com.atproto.appPass", "access-jti", 4102444800, "refresh-jti", 4102444801);
    defer std.testing.allocator.free(pair.access.token);
    defer std.testing.allocator.free(pair.access.jti);
    defer std.testing.allocator.free(pair.refresh.token);
    defer std.testing.allocator.free(pair.refresh.jti);
    const access = try parseSessionJwt(std.testing.allocator, pair.access.token);
    defer std.testing.allocator.free(access.did);
    defer std.testing.allocator.free(access.scope);
    defer std.testing.allocator.free(access.jti);
    try std.testing.expectEqualStrings("access-jti", access.jti);
    try std.testing.expectEqualStrings("com.atproto.appPass", access.scope);
    try std.testing.expectEqual(@as(i64, 4102444800), access.exp);
    const refresh = try parseSessionJwt(std.testing.allocator, pair.refresh.token);
    defer std.testing.allocator.free(refresh.did);
    defer std.testing.allocator.free(refresh.scope);
    defer std.testing.allocator.free(refresh.jti);
    try std.testing.expectEqualStrings("refresh-jti", refresh.jti);
    try std.testing.expectEqual(@as(i64, 4102444801), refresh.exp);
}

test "app-password session pair uses appPass access scope" {
    const alice = testAccount();
    const pair = try createSessionPairWithAccessScope(std.testing.allocator, alice, "com.atproto.appPass");
    defer std.testing.allocator.free(pair.access.token);
    defer std.testing.allocator.free(pair.access.jti);
    defer std.testing.allocator.free(pair.refresh.token);
    defer std.testing.allocator.free(pair.refresh.jti);
    const access_claims = claimsFromSessionJwt(std.testing.allocator, pair.access.token).?;
    defer std.testing.allocator.free(access_claims.did);
    defer std.testing.allocator.free(access_claims.scope);
    defer std.testing.allocator.free(access_claims.jti);
    defer if (access_claims.cnf_jkt) |jkt| std.testing.allocator.free(jkt);
    try std.testing.expectEqualStrings("com.atproto.appPass", access_claims.scope);
    const refresh_claims = claimsFromSessionJwt(std.testing.allocator, pair.refresh.token).?;
    defer std.testing.allocator.free(refresh_claims.did);
    defer std.testing.allocator.free(refresh_claims.scope);
    defer std.testing.allocator.free(refresh_claims.jti);
    defer if (refresh_claims.cnf_jkt) |jkt| std.testing.allocator.free(jkt);
    try std.testing.expectEqualStrings("com.atproto.refresh", refresh_claims.scope);
}

test "dpop session token carries cnf jkt" {
    const alice = testAccount();
    const token = try createDpopSessionJwt(std.testing.allocator, "access", alice, "test-jkt");
    defer std.testing.allocator.free(token);
    const claims = claimsFromSessionJwt(std.testing.allocator, token).?;
    defer std.testing.allocator.free(claims.did);
    defer std.testing.allocator.free(claims.scope);
    defer std.testing.allocator.free(claims.jti);
    defer if (claims.cnf_jkt) |jkt| std.testing.allocator.free(jkt);
    try std.testing.expectEqualStrings("test-jkt", claims.cnf_jkt.?);
}

test "creates signed service auth token" {
    const alice = testAccount();
    const token = try createServiceJwt(std.testing.allocator, alice, "did:plc:service", "xyz.fake.inbox.send");
    defer std.testing.allocator.free(token);
    try std.testing.expect(std.mem.count(u8, token, ".") == 2);
}

fn testAccount() Account {
    return .{
        .handle = "alice.test",
        .did = "did:plc:cmadossymmii3izkabdbp5en",
        .email = "alice@test.com",
        .password = "password",
    };
}