//! jss sealed segment format v1 — reader
//!
//! byte-exact implementation of upstream jetstream's segment file format
//! (docs/jss-format-v1.md). the format is wire-visible contract: getSegment
//! serves these raw bytes to clients. pure format code — no goroutine-shaped
//! lifecycle, no I/O beyond the caller handing us bytes.
//!
//! sealed-file parsing over an in-memory buffer (header, block index,
//! columnar block decode, collection index). blooms are decoded by the
//! resident-manifest layer. The file-level xxh3 authenticates header/footer
//! metadata and libzstd independently verifies every block-frame checksum.

const std = @import("std");
const zstd = @import("zstd.zig");

const Allocator = std.mem.Allocator;

pub const header_size = 256;
pub const magic = "jss0";
pub const version = 1;

// reader sanity caps mirroring upstream
const max_block_events = 1 << 18;
const max_blocks = 1 << 20;
const max_collections = 1 << 20;
const max_decompressed_block = 1 << 30;

pub const Kind = enum(u8) {
    create = 1,
    update = 2,
    delete = 3,
    identity = 4,
    account = 5,
    sync = 6,
    create_resync = 7,
};

pub const Error = error{
    Truncated,
    BadMagic,
    ActiveSegment,
    UnsupportedVersion,
    BadChecksum,
    BadBlockFrame,
    BadBlockEncoding,
    BadFooter,
    LimitExceeded,
    OutOfMemory,
    DecompressFailed,
};

pub const Header = struct {
    checksum: u64,
    block_count: u32,
    event_count: u32,
    unique_did_count: u32,
    min_seq: u64,
    max_seq: u64,
    min_witnessed_at: i64,
    max_witnessed_at: i64,
    footer_offset: u64,
    did_bloom_offset: u64,
    block_did_bloom_offset: u64,
    collection_index_offset: u64,
    block_index_offset: u64,

    pub fn decode(bytes: []const u8) Error!Header {
        if (bytes.len < header_size) return error.Truncated;
        if (!std.mem.eql(u8, bytes[0..4], magic)) return error.BadMagic;
        const checksum = readInt(u64, bytes, 4);
        if (checksum == 0) return error.ActiveSegment;
        if (readInt(u16, bytes, 12) != version) return error.UnsupportedVersion;
        const h: Header = .{
            .checksum = checksum,
            .block_count = readInt(u32, bytes, 14),
            .event_count = readInt(u32, bytes, 18),
            .unique_did_count = readInt(u32, bytes, 22),
            .min_seq = readInt(u64, bytes, 26),
            .max_seq = readInt(u64, bytes, 34),
            .min_witnessed_at = @bitCast(readInt(u64, bytes, 42)),
            .max_witnessed_at = @bitCast(readInt(u64, bytes, 50)),
            .footer_offset = readInt(u64, bytes, 58),
            .did_bloom_offset = readInt(u64, bytes, 66),
            .block_did_bloom_offset = readInt(u64, bytes, 74),
            .collection_index_offset = readInt(u64, bytes, 82),
            .block_index_offset = readInt(u64, bytes, 90),
        };
        if (h.block_count > max_blocks) return error.LimitExceeded;
        if (h.block_index_offset != h.footer_offset) return error.BadFooter;
        return h;
    }
};

pub const BlockIndexEntry = struct {
    offset: u64, // file offset of the 8-byte length prefix
    compressed_size: u32,
    uncompressed_size: u32,
    event_count: u32,
    min_seq: u64,
    max_seq: u64,
    min_witnessed_at: i64,
    max_witnessed_at: i64,

    pub const encoded_size = 52;

    pub fn decodePublic(bytes: *const [encoded_size]u8) BlockIndexEntry {
        return decode(bytes);
    }

    fn decode(bytes: []const u8) BlockIndexEntry {
        return .{
            .offset = readInt(u64, bytes, 0),
            .compressed_size = readInt(u32, bytes, 8),
            .uncompressed_size = readInt(u32, bytes, 12),
            .event_count = readInt(u32, bytes, 16),
            .min_seq = readInt(u64, bytes, 20),
            .max_seq = readInt(u64, bytes, 28),
            .min_witnessed_at = @bitCast(readInt(u64, bytes, 36)),
            .max_witnessed_at = @bitCast(readInt(u64, bytes, 44)),
        };
    }
};

/// one event decoded from a block; strings/payload borrow from the block's
/// decompressed buffer.
pub const Event = struct {
    seq: u64,
    witnessed_at: i64,
    indexed_at: i64, // 0 = unset; display time falls back to witnessed_at
    kind: Kind,
    did: []const u8,
    collection: []const u8,
    rkey: []const u8,
    rev: []const u8,
    payload: []const u8, // raw DAG-CBOR; empty for deletes

    /// Timestamp exposed as `time_us` on both subscriber protocols. An
    /// operator import overrides presentation only; witnessed_at remains the
    /// immutable range/cursor timestamp and sentinel zero means no override.
    pub fn displayTime(self: Event) i64 {
        return if (self.indexed_at != 0) self.indexed_at else self.witnessed_at;
    }
};

/// events decoded from one block, backed by a single decompressed buffer.
pub const Block = struct {
    events: []Event,
    buffer: []u8,

    pub fn deinit(self: *Block, allocator: Allocator) void {
        allocator.free(self.events);
        allocator.free(self.buffer);
    }
};

pub const CollectionEntry = struct {
    nsid: []const u8, // borrows from CollectionIndex.buffer
    count: u32, // segment-wide event count (0 for $-sentinels)
};

pub const CollectionIndex = struct {
    entries: []CollectionEntry,
    /// block_count × bitmask_len; bit N of row b set ⇒ collection N in block b
    bitmasks: []const u8,
    bitmask_len: u32,
    buffer: []u8,

    pub fn deinit(self: *CollectionIndex, allocator: Allocator) void {
        allocator.free(self.entries);
        allocator.free(self.buffer);
    }

    /// null when this block has no resident bitmask row. A planner must treat
    /// null as "may contain anything" and fail open; an exact reporter must
    /// not invent membership from it.
    pub fn blockRow(self: *const CollectionIndex, block: usize) ?[]const u8 {
        if (self.bitmask_len == 0) return null;
        const at = block * self.bitmask_len;
        if (at + self.bitmask_len > self.bitmasks.len) return null;
        return self.bitmasks[at..][0..self.bitmask_len];
    }

    pub fn blockHasCollection(self: *const CollectionIndex, block: usize, collection_id: usize) bool {
        if (collection_id >= self.entries.len) return false;
        const row = self.blockRow(block) orelse return false;
        return (row[collection_id / 8] >> @intCast(collection_id % 8)) & 1 == 1;
    }
};

/// a parsed sealed segment over caller-owned bytes.
pub const Sealed = struct {
    bytes: []const u8,
    header: Header,
    block_index: []BlockIndexEntry,

    pub fn deinit(self: *Sealed, allocator: Allocator) void {
        allocator.free(self.block_index);
    }

    /// parse and validate header + block index; verifies the file xxh3
    /// (header[12..256] || footer bytes).
    pub fn parse(allocator: Allocator, bytes: []const u8) Error!Sealed {
        return parseInternal(allocator, bytes, true);
    }

    /// Parse the sealed structure without rejecting a checksum mismatch.
    /// Operator inspection uses this path so it can report an otherwise
    /// readable segment as `checksum: ... (invalid)`, matching upstream.
    /// Subscriber/archive paths must continue to use `parse`.
    pub fn parseUnchecked(allocator: Allocator, bytes: []const u8) Error!Sealed {
        return parseInternal(allocator, bytes, false);
    }

    fn parseInternal(allocator: Allocator, bytes: []const u8, verify_checksum: bool) Error!Sealed {
        const header = try Header.decode(bytes);
        if (header.footer_offset > bytes.len or header.footer_offset < header_size)
            return error.Truncated;
        if (header.block_index_offset != header.footer_offset or
            header.did_bloom_offset < header.block_index_offset or
            header.block_did_bloom_offset < header.did_bloom_offset or
            header.collection_index_offset < header.block_did_bloom_offset or
            header.collection_index_offset > bytes.len)
            return error.BadFooter;

        if (verify_checksum and !checksumValid(bytes, header)) return error.BadChecksum;

        const index_bytes_len = @as(u64, header.block_count) * BlockIndexEntry.encoded_size;
        if (header.block_index_offset + index_bytes_len > bytes.len) return error.Truncated;

        const entries = try allocator.alloc(BlockIndexEntry, header.block_count);
        errdefer allocator.free(entries);
        for (entries, 0..) |*entry, i| {
            const at = header.block_index_offset + i * BlockIndexEntry.encoded_size;
            entry.* = BlockIndexEntry.decode(bytes[at..][0..BlockIndexEntry.encoded_size]);
            if (entry.offset < header_size or
                entry.offset + 8 + entry.compressed_size > header.footer_offset)
                return error.BadBlockFrame;
        }

        return .{ .bytes = bytes, .header = header, .block_index = entries };
    }

    /// decompress and decode block `idx`.
    pub fn readBlock(self: *const Sealed, allocator: Allocator, idx: usize) Error!Block {
        const entry = self.block_index[idx];
        const frame_len = readInt(u64, self.bytes, entry.offset);
        if (frame_len != entry.compressed_size) return error.BadBlockFrame;
        const frame = self.bytes[entry.offset + 8 ..][0..entry.compressed_size];
        if (entry.uncompressed_size > max_decompressed_block) return error.LimitExceeded;

        const buffer = decompress(allocator, frame, entry.uncompressed_size) catch |err| switch (err) {
            error.OutOfMemory => return error.OutOfMemory,
            else => return error.DecompressFailed,
        };
        errdefer allocator.free(buffer);
        return decodeBlock(allocator, buffer);
    }

    /// per-block DID bloom region: block_count blobs of identical size.
    /// returns the raw gloom blobs (borrowing self.bytes).
    pub fn blockBloomBlobs(self: *const Sealed, allocator: Allocator) Error![][]const u8 {
        const start = self.header.block_did_bloom_offset;
        if (start + 8 > self.bytes.len) return error.Truncated;
        const count = readInt(u32, self.bytes, start);
        const size = readInt(u32, self.bytes, start + 4);
        if (count != self.header.block_count) return error.BadFooter;
        const region_end = start + 8 + @as(u64, count) * size;
        if (region_end > self.bytes.len) return error.Truncated;
        const blobs = try allocator.alloc([]const u8, count);
        for (blobs, 0..) |*b, i| {
            const at = start + 8 + i * @as(u64, size);
            b.* = self.bytes[at..][0..size];
        }
        return blobs;
    }

    /// parse the collection block index from the footer.
    pub fn readCollectionIndex(self: *const Sealed, allocator: Allocator) Error!CollectionIndex {
        const start = self.header.collection_index_offset;
        if (start + 16 > self.bytes.len) return error.Truncated;
        return decodeCollectionIndex(allocator, self.bytes[start..], self.header.block_count);
    }
};

pub fn checksumValid(bytes: []const u8, header: Header) bool {
    if (bytes.len < header_size or header.footer_offset > bytes.len) return false;
    var hasher = std.hash.XxHash3.init(0);
    hasher.update(bytes[12..header_size]);
    hasher.update(bytes[header.footer_offset..]);
    return hasher.final() == header.checksum;
}

/// Decode the collection footer from a standalone footer slice. The resident
/// manifest uses this at startup so planSnapshot never has to reopen segment
/// data files on its request path.
pub fn decodeCollectionIndex(allocator: Allocator, bytes: []const u8, expected_block_count: u32) Error!CollectionIndex {
    if (bytes.len < 16) return error.Truncated;
    const collection_count = readInt(u32, bytes, 0);
    const block_count = readInt(u32, bytes, 4);
    const bitmask_len = readInt(u32, bytes, 8);
    const uncompressed_size = readInt(u32, bytes, 12);
    if (collection_count > max_collections) return error.LimitExceeded;
    if (bitmask_len != (collection_count + 7) / 8) return error.BadFooter;
    if (block_count != expected_block_count) return error.BadFooter;

    const body_frame = bytes[16..];
    const buffer = decompress(allocator, body_frame, uncompressed_size) catch |err| switch (err) {
        error.OutOfMemory => return error.OutOfMemory,
        else => return error.DecompressFailed,
    };
    errdefer allocator.free(buffer);

    const entries = try allocator.alloc(CollectionEntry, collection_count);
    errdefer allocator.free(entries);
    var pos: usize = 0;
    for (entries) |*e| {
        if (pos + 5 > buffer.len) return error.BadFooter;
        const len = buffer[pos];
        const count = readInt(u32, buffer, pos + 1);
        pos += 5;
        if (pos + len > buffer.len) return error.BadFooter;
        e.* = .{ .nsid = buffer[pos..][0..len], .count = count };
        pos += len;
    }
    const bitmask_bytes = @as(usize, block_count) * bitmask_len;
    if (pos + bitmask_bytes != buffer.len) return error.BadFooter;

    return .{
        .entries = entries,
        .bitmasks = buffer[pos..],
        .bitmask_len = bitmask_len,
        .buffer = buffer,
    };
}

/// Read and decode one block from an open segment without loading the whole
/// segment. Cursor resolution uses this path because a full-network segment
/// can be much larger than the one block needed to translate a timestamp.
/// Like upstream's cursor reader, this deliberately skips the file checksum;
/// the frame length, decompressed size, and columnar encoding are still
/// validated here.
pub fn readBlockFile(
    allocator: Allocator,
    io: std.Io,
    file: *std.Io.File,
    entry: BlockIndexEntry,
) !Block {
    if (entry.uncompressed_size > max_decompressed_block) return error.LimitExceeded;

    var frame_len_buf: [8]u8 = undefined;
    _ = try file.readPositionalAll(io, &frame_len_buf, entry.offset);
    if (std.mem.readInt(u64, &frame_len_buf, .little) != entry.compressed_size)
        return error.BadBlockFrame;

    const frame = try allocator.alloc(u8, entry.compressed_size);
    defer allocator.free(frame);
    _ = try file.readPositionalAll(io, frame, entry.offset + 8);
    const buffer = decompress(allocator, frame, entry.uncompressed_size) catch |err| switch (err) {
        error.OutOfMemory => return error.OutOfMemory,
        else => return error.DecompressFailed,
    };
    errdefer allocator.free(buffer);
    return decodeBlock(allocator, buffer);
}

pub fn readHeaderFile(io: std.Io, file: *std.Io.File) !Header {
    var bytes: [header_size]u8 = undefined;
    if (try file.readPositionalAll(io, &bytes, 0) != bytes.len) return error.Truncated;
    return Header.decode(&bytes);
}

pub fn readBlockIndexFile(io: std.Io, file: *std.Io.File, header: Header, index: usize) !BlockIndexEntry {
    if (index >= header.block_count) return error.BadBlockFrame;
    var bytes: [BlockIndexEntry.encoded_size]u8 = undefined;
    const offset = header.block_index_offset + @as(u64, @intCast(index)) * BlockIndexEntry.encoded_size;
    if (try file.readPositionalAll(io, &bytes, offset) != bytes.len) return error.Truncated;
    const entry = BlockIndexEntry.decodePublic(&bytes);
    if (entry.offset < header_size or entry.offset + 8 + entry.compressed_size > header.footer_offset)
        return error.BadBlockFrame;
    return entry;
}

/// decode a decompressed columnar block. takes ownership of `buffer` (freed
/// via Block.deinit; events borrow from it).
pub fn decodeBlock(allocator: Allocator, buffer: []u8) Error!Block {
    if (buffer.len < 4) return error.BadBlockEncoding;
    const n = readInt(u32, buffer, 0);
    if (n > max_block_events) return error.LimitExceeded;
    if (n == 0) {
        // compaction can leave empty blocks: exactly the 4 count bytes
        if (buffer.len != 4) return error.BadBlockEncoding;
        return .{ .events = try allocator.alloc(Event, 0), .buffer = buffer };
    }

    const counts: usize = n;
    // fixed-width column region
    const fixed_len = 4 + counts * (8 + 8 + 8 + 1 + 1 + 2 + 1 + 1 + 4);
    if (buffer.len < fixed_len) return error.BadBlockEncoding;

    var off: usize = 4;
    const seqs = off;
    off += counts * 8;
    const witnessed = off;
    off += counts * 8;
    const indexed = off;
    off += counts * 8;
    const kinds = off;
    off += counts;
    const collection_lens = off;
    off += counts;
    const did_lens = off;
    off += counts * 2;
    const rkey_lens = off;
    off += counts;
    const rev_lens = off;
    off += counts;
    const event_lens = off;
    off += counts * 4;

    const events = try allocator.alloc(Event, counts);
    errdefer allocator.free(events);

    // blobs follow, in column order; walk with running offsets
    const col_pos: usize = off;
    var total_collections: usize = 0;
    var total_dids: usize = 0;
    var total_rkeys: usize = 0;
    var total_revs: usize = 0;
    var total_payloads: usize = 0;
    for (0..counts) |i| {
        total_collections += buffer[collection_lens + i];
        total_dids += readInt(u16, buffer, did_lens + i * 2);
        total_rkeys += buffer[rkey_lens + i];
        total_revs += buffer[rev_lens + i];
        total_payloads += readInt(u32, buffer, event_lens + i * 4);
    }
    const did_pos = col_pos + total_collections;
    const rkey_pos = did_pos + total_dids;
    const rev_pos = rkey_pos + total_rkeys;
    const payload_pos = rev_pos + total_revs;
    if (payload_pos + total_payloads != buffer.len) return error.BadBlockEncoding;

    var c: usize = col_pos;
    var d: usize = did_pos;
    var rk: usize = rkey_pos;
    var rv: usize = rev_pos;
    var p: usize = payload_pos;
    for (events, 0..) |*ev, i| {
        const kind_raw = buffer[kinds + i];
        if (kind_raw < 1 or kind_raw > 7) return error.BadBlockEncoding;
        const clen = buffer[collection_lens + i];
        const dlen = readInt(u16, buffer, did_lens + i * 2);
        const rklen = buffer[rkey_lens + i];
        const rvlen = buffer[rev_lens + i];
        const plen = readInt(u32, buffer, event_lens + i * 4);
        ev.* = .{
            .seq = readInt(u64, buffer, seqs + i * 8),
            .witnessed_at = @bitCast(readInt(u64, buffer, witnessed + i * 8)),
            .indexed_at = @bitCast(readInt(u64, buffer, indexed + i * 8)),
            .kind = @fromBackingInt(@intCast(kind_raw)),
            .collection = buffer[c..][0..clen],
            .did = buffer[d..][0..dlen],
            .rkey = buffer[rk..][0..rklen],
            .rev = buffer[rv..][0..rvlen],
            .payload = buffer[p..][0..plen],
        };
        c += clen;
        d += dlen;
        rk += rklen;
        rv += rvlen;
        p += plen;
    }

    return .{ .events = events, .buffer = buffer };
}

fn decompress(allocator: Allocator, frame: []const u8, expected_len: u32) ![]u8 {
    return zstd.decompressExact(allocator, frame, expected_len);
}

inline fn readInt(comptime T: type, bytes: []const u8, offset: usize) T {
    return std.mem.readInt(T, bytes[offset..][0..@sizeOf(T)], .little);
}

// === tests ===
// differential tests against an upstream-produced fixture live in
// segment_fixture_test.zig (skipped when the fixture is absent).

const testing = std.testing;

test "header rejects bad magic, active files, wrong version" {
    var buf: [header_size]u8 = @splat(0);
    try testing.expectError(error.BadMagic, Header.decode(&buf));
    @memcpy(buf[0..4], magic);
    try testing.expectError(error.ActiveSegment, Header.decode(&buf));
    std.mem.writeInt(u64, buf[4..12], 0xdead, .little);
    std.mem.writeInt(u16, buf[12..14], 2, .little);
    try testing.expectError(error.UnsupportedVersion, Header.decode(&buf));
    try testing.expectError(error.Truncated, Header.decode(buf[0..100]));
}

test "unchecked sealed parse exposes checksum corruption only to inspection" {
    const writer = @import("segment_writer.zig");
    var active = try writer.ActiveWriter.init(testing.allocator);
    defer active.deinit();
    try active.append(.{
        .seq = 1,
        .witnessed_at = 1,
        .indexed_at = 0,
        .kind = .create,
        .did = "did:plc:test",
        .collection = "app.bsky.feed.post",
        .rkey = "r",
        .rev = "v",
        .payload = "p",
    });
    try active.seal();
    active.buf.items[active.buf.items.len - 1] ^= 0x01;

    try testing.expectError(error.BadChecksum, Sealed.parse(testing.allocator, active.bytes()));
    var inspected = try Sealed.parseUnchecked(testing.allocator, active.bytes());
    defer inspected.deinit(testing.allocator);
    try testing.expect(!checksumValid(active.bytes(), inspected.header));
}

test "empty block decodes to zero events" {
    const buffer = try testing.allocator.dupe(u8, &[_]u8{ 0, 0, 0, 0 });
    var block = try decodeBlock(testing.allocator, buffer);
    defer block.deinit(testing.allocator);
    try testing.expectEqual(@as(usize, 0), block.events.len);
}

test "block with trailing garbage is rejected" {
    const buffer = try testing.allocator.dupe(u8, &[_]u8{ 0, 0, 0, 0, 0xff });
    defer testing.allocator.free(buffer);
    try testing.expectError(error.BadBlockEncoding, decodeBlock(testing.allocator, @constCast(buffer)));
}