//! minimal zstd bindings
//!
//! backed by the vendored facebook/zstd source (pinned by url+hash in
//! build.zig.zon, compiled in build.zig — no system library). Both directions
//! use libzstd: the writer needs a compressor (absent from Zig std), and the
//! reader needs content-checksum verification matching upstream's decoder.

const std = @import("std");

const c = struct {
    extern fn ZSTD_createCCtx() ?*anyopaque;
    extern fn ZSTD_freeCCtx(cctx: ?*anyopaque) usize;
    extern fn ZSTD_CCtx_setParameter(cctx: ?*anyopaque, param: c_int, value: c_int) usize;
    extern fn ZSTD_compress2(cctx: ?*anyopaque, dst: [*]u8, dst_cap: usize, src: [*]const u8, src_len: usize) usize;
    extern fn ZSTD_compressBound(src_len: usize) usize;
    extern fn ZSTD_decompress(dst: [*]u8, dst_cap: usize, src: [*]const u8, src_len: usize) usize;
    extern fn ZSTD_isError(code: usize) c_uint;

    // ZSTD_cParameter values (zstd.h, stable API)
    const ZSTD_c_compressionLevel: c_int = 100;
    const ZSTD_c_checksumFlag: c_int = 201;
};

pub const level_default: c_int = 3;

/// compress `src` into one zstd frame with a content checksum.
/// returned slice is allocated from `allocator`.
pub fn compress(allocator: std.mem.Allocator, src: []const u8) ![]u8 {
    const cctx = c.ZSTD_createCCtx() orelse return error.OutOfMemory;
    defer _ = c.ZSTD_freeCCtx(cctx);
    if (c.ZSTD_isError(c.ZSTD_CCtx_setParameter(cctx, c.ZSTD_c_compressionLevel, level_default)) != 0)
        return error.CompressFailed;
    if (c.ZSTD_isError(c.ZSTD_CCtx_setParameter(cctx, c.ZSTD_c_checksumFlag, 1)) != 0)
        return error.CompressFailed;

    const bound = c.ZSTD_compressBound(src.len);
    const dst = try allocator.alloc(u8, bound);
    errdefer allocator.free(dst);
    const written = c.ZSTD_compress2(cctx, dst.ptr, dst.len, src.ptr, src.len);
    if (c.ZSTD_isError(written) != 0) return error.CompressFailed;
    return allocator.realloc(dst, written);
}

/// Decompress one frame into its declared exact-size output. libzstd verifies
/// the frame content checksum when present; all JSS writers enable it.
pub fn decompressExact(allocator: std.mem.Allocator, frame: []const u8, expected_len: usize) ![]u8 {
    const dst = try allocator.alloc(u8, expected_len);
    errdefer allocator.free(dst);
    const written = c.ZSTD_decompress(dst.ptr, dst.len, frame.ptr, frame.len);
    if (c.ZSTD_isError(written) != 0 or written != expected_len)
        return error.DecompressFailed;
    return dst;
}

// === tests ===

fn repeated(comptime unit: []const u8, comptime count: usize) *const [unit.len * count]u8 {
    return comptime blk: {
        var out: [unit.len * count]u8 = undefined;
        for (0..count) |i| @memcpy(out[i * unit.len ..][0..unit.len], unit);
        const final = out;
        break :blk &final;
    };
}

test "round-trip through std decompress" {
    const src = repeated("the quick brown fox jumps over the lazy dog", 100);
    const frame = try compress(std.testing.allocator, src);
    defer std.testing.allocator.free(frame);
    try std.testing.expect(frame.len < src.len);

    var in: std.Io.Reader = .fixed(frame);
    var stream: std.compress.zstd.Decompress = .init(&in, &.{}, .{});
    var out: std.Io.Writer.Allocating = .init(std.testing.allocator);
    defer out.deinit();
    _ = try stream.reader.streamRemaining(&out.writer);
    try std.testing.expectEqualStrings(src, out.written());
}

test "libzstd decompress verifies the frame content checksum" {
    const src = repeated("checksum-protected block", 100);
    const frame = try compress(std.testing.allocator, src);
    defer std.testing.allocator.free(frame);

    const decoded = try decompressExact(std.testing.allocator, frame, src.len);
    defer std.testing.allocator.free(decoded);
    try std.testing.expectEqualStrings(src, decoded);

    frame[frame.len - 1] ^= 0x01;
    try std.testing.expectError(
        error.DecompressFailed,
        decompressExact(std.testing.allocator, frame, src.len),
    );
}

// === dictionary compression (the /subscribe wire's v1 scheme) ===

const cdict = struct {
    extern fn ZSTD_createCCtxParams() ?*anyopaque;
    extern fn ZSTD_freeCCtxParams(p: ?*anyopaque) usize;
    extern fn ZSTD_CCtx_setPledgedSrcSize(cctx: ?*anyopaque, size: c_ulonglong) usize;
    extern fn ZSTD_CCtx_loadDictionary(cctx: ?*anyopaque, dict: [*]const u8, size: usize) usize;
    extern fn ZSTD_CCtx_reset(cctx: ?*anyopaque, reset: c_int) usize;
};

const ZSTD_c_windowLog: c_int = 101;
const ZSTD_reset_session_only: c_int = 1;

/// per-message dictionary compressor matching upstream's v1 encoder:
/// dict 1612007021, 128 KiB window, content checksum, level 3. one per
/// subscriber (compression contexts are not thread-safe).
pub const DictCompressor = struct {
    cctx: *anyopaque,

    pub fn init(dictionary: []const u8) !DictCompressor {
        return initLevel(dictionary, level_default);
    }

    pub fn initLevel(dictionary: []const u8, level: c_int) !DictCompressor {
        const cctx = c.ZSTD_createCCtx() orelse return error.OutOfMemory;
        errdefer _ = c.ZSTD_freeCCtx(cctx);
        if (c.ZSTD_isError(c.ZSTD_CCtx_setParameter(cctx, c.ZSTD_c_compressionLevel, level)) != 0)
            return error.CompressFailed;
        if (c.ZSTD_isError(c.ZSTD_CCtx_setParameter(cctx, c.ZSTD_c_checksumFlag, 1)) != 0)
            return error.CompressFailed;
        if (c.ZSTD_isError(c.ZSTD_CCtx_setParameter(cctx, ZSTD_c_windowLog, 17)) != 0)
            return error.CompressFailed;
        if (c.ZSTD_isError(cdict.ZSTD_CCtx_loadDictionary(cctx, dictionary.ptr, dictionary.len)) != 0)
            return error.CompressFailed;
        return .{ .cctx = cctx };
    }

    pub fn deinit(self: *DictCompressor) void {
        _ = c.ZSTD_freeCCtx(self.cctx);
    }

    /// compress one message into a standalone zstd frame referencing the
    /// dictionary. the loaded dictionary persists across session resets.
    pub fn compress(self: *DictCompressor, allocator: std.mem.Allocator, src: []const u8) ![]u8 {
        if (c.ZSTD_isError(cdict.ZSTD_CCtx_reset(self.cctx, ZSTD_reset_session_only)) != 0)
            return error.CompressFailed;
        const bound = c.ZSTD_compressBound(src.len);
        const dst = try allocator.alloc(u8, bound);
        errdefer allocator.free(dst);
        const written = c.ZSTD_compress2(self.cctx, dst.ptr, dst.len, src.ptr, src.len);
        if (c.ZSTD_isError(written) != 0) return error.CompressFailed;
        return allocator.realloc(dst, written);
    }
};

pub const v1_dictionary = @embedFile("../../assets/zstd_dictionary");

const v2_dictionary_b64 = @embedFile("../../assets/zstd_dictionary_v2.b64");
pub const v2_dictionary = blk: {
    @setEvalBranchQuota(200_000);
    const decoder = std.base64.standard.decoderWithIgnore("\r\n");
    var decoded: [65_536]u8 = undefined;
    const written = decoder.decode(&decoded, v2_dictionary_b64) catch
        @compileError("invalid embedded v2 zstd dictionary base64");
    if (written != decoded.len) @compileError("embedded v2 zstd dictionary has the wrong length");
    break :blk decoded;
};

pub fn dictionaryId(dictionary: []const u8) !u32 {
    if (dictionary.len < 8) return error.InvalidDictionary;
    if (std.mem.readInt(u32, dictionary[0..4], .little) != 0xEC30A437)
        return error.InvalidDictionary;
    const id = std.mem.readInt(u32, dictionary[4..8], .little);
    if (id == 0) return error.InvalidDictionary;
    return id;
}

pub const v2_dictionary_id = dictionaryId(&v2_dictionary) catch
    @compileError("invalid embedded v2 zstd dictionary header");

test "dictionary round-trip carries the v1 dict id" {
    var comp = try DictCompressor.init(v1_dictionary);
    defer comp.deinit();
    const msg = "{\"did\":\"did:plc:abc\",\"time_us\":1,\"kind\":\"commit\"}";
    const frame = try comp.compress(std.testing.allocator, msg);
    defer std.testing.allocator.free(frame);
    // frame header carries the dict id: 1612007021 = 0x6015466d LE
    try std.testing.expect(std.mem.indexOf(u8, frame[0..16], "\x6d\x46\x15\x60") != null);

    // second message reuses the dictionary after session reset
    const frame2 = try comp.compress(std.testing.allocator, msg);
    defer std.testing.allocator.free(frame2);
    try std.testing.expect(frame2.len > 0);
}

test "v2 dictionary is the pinned upstream artifact" {
    try std.testing.expectEqual(@as(usize, 65_536), v2_dictionary.len);
    try std.testing.expectEqual(@as(u32, 20_260_811), v2_dictionary_id);
    try std.testing.expectEqual(@as(u32, 0xEC30A437), std.mem.readInt(u32, v2_dictionary[0..4], .little));
    var digest: [32]u8 = undefined;
    std.crypto.hash.sha2.Sha256.hash(&v2_dictionary, &digest, .{});
    var expected: [32]u8 = undefined;
    _ = try std.fmt.hexToBytes(&expected, "f631a689a829b84b042ba45bcd7a7c2f792467192c5832fcfc1fc48afbec96aa");
    try std.testing.expectEqualSlices(u8, &expected, &digest);
}