const std = @import("std");
const encode = @import("encode.zig");
const Bloom = @import("bloom.zig").Bloom;
const cache = @import("cache.zig");
const compress = @import("compress.zig");

const Allocator = std.mem.Allocator;
const Kind = encode.Kind;
const Record = encode.Record;
const BlockCache = cache.BlockCache;
const BlockKey = cache.BlockKey;
const Codec = compress.Codec;
const readU32 = encode.readU32;
const readU64 = encode.readU64;

pub const OwnedHit = struct {
    seq: u64,
    kind: Kind,
    value: []u8,
};

fn order(a: []const u8, b: []const u8) std.math.Order {
    return std.mem.order(u8, a, b);
}

const magic: u32 = 0xC0BA1701;
const version: u32 = 1;
const footerLen = 4 + 4 + 8 + 4 + 8 + 4 + 8 + 8;
const recordHeader = 1 + 8 + 4 + 4;
const defaultBlockTarget = 4096;

pub const Builder = struct {
    allocator: Allocator,
    file: std.fs.File,
    offset: u64,
    blockBuf: std.ArrayList(u8),
    blockFirstKey: std.ArrayList(u8),
    indexBuf: std.ArrayList(u8),
    numEntries: u64,
    logicalBytes: u64,
    maxSeq: u64,
    bloom: Bloom,
    blockTarget: usize,

    pub fn init(allocator: Allocator, dir: std.fs.Dir, name: []const u8, expectedKeys: usize) !Builder {
        const file = try dir.createFile(name, .{ .truncate = true });
        return .{
            .allocator = allocator,
            .file = file,
            .offset = 0,
            .blockBuf = std.ArrayList(u8).init(allocator),
            .blockFirstKey = std.ArrayList(u8).init(allocator),
            .indexBuf = std.ArrayList(u8).init(allocator),
            .numEntries = 0,
            .logicalBytes = 0,
            .maxSeq = 0,
            .bloom = try Bloom.init(allocator, expectedKeys, 0.01),
            .blockTarget = defaultBlockTarget,
        };
    }

    pub fn deinit(self: *Builder) void {
        self.blockBuf.deinit();
        self.blockFirstKey.deinit();
        self.indexBuf.deinit();
        self.bloom.deinit();
        self.file.close();
    }

    pub fn add(self: *Builder, record: Record) !void {
        if (self.blockBuf.items.len == 0) {
            self.blockFirstKey.clearRetainingCapacity();
            try self.blockFirstKey.appendSlice(record.key);
        }
        const writer = encode.ByteWriter.init(&self.blockBuf);
        try writer.byte(@intFromEnum(record.kind));
        try writer.putU64(record.seq);
        try writer.putU32(@intCast(record.key.len));
        try writer.putU32(@intCast(record.value.len));
        try writer.slice(record.key);
        try writer.slice(record.value);

        self.bloom.add(record.key);
        self.numEntries += 1;
        self.logicalBytes += recordHeader + record.key.len + record.value.len;
        if (record.seq > self.maxSeq) self.maxSeq = record.seq;

        if (self.blockBuf.items.len >= self.blockTarget) try self.flushBlock();
    }

    fn flushBlock(self: *Builder) !void {
        if (self.blockBuf.items.len == 0) return;
        const originalLen: u32 = @intCast(self.blockBuf.items.len);

        const encoded = try compress.encode(self.allocator, self.blockBuf.items);
        defer encoded.deinit(self.allocator);

        var physical = std.ArrayList(u8).init(self.allocator);
        defer physical.deinit();
        const physWriter = encode.ByteWriter.init(&physical);
        try physWriter.byte(@intFromEnum(encoded.codec));
        try physWriter.putU32(originalLen);
        try physWriter.slice(encoded.bytes);

        const physLen: u32 = @intCast(physical.items.len);
        try self.file.writeAll(physical.items);
        var crcBuf: [4]u8 = undefined;
        std.mem.writeInt(u32, &crcBuf, encode.crc(physical.items), encode.endian);
        try self.file.writeAll(&crcBuf);

        const indexWriter = encode.ByteWriter.init(&self.indexBuf);
        try indexWriter.putU64(self.offset);
        try indexWriter.putU32(physLen);
        try indexWriter.putU32(@intCast(self.blockFirstKey.items.len));
        try indexWriter.slice(self.blockFirstKey.items);

        self.offset += physLen + 4;
        self.blockBuf.clearRetainingCapacity();
    }

    pub fn finish(self: *Builder) !void {
        try self.flushBlock();

        var bloomBuf = std.ArrayList(u8).init(self.allocator);
        defer bloomBuf.deinit();
        try self.bloom.serialize(encode.ByteWriter.init(&bloomBuf));
        const bloomOffset = self.offset;
        const bloomLen: u32 = @intCast(bloomBuf.items.len);
        try self.writeBlock(bloomBuf.items);

        const indexOffset = self.offset;
        const indexLen: u32 = @intCast(self.indexBuf.items.len);
        try self.writeBlock(self.indexBuf.items);

        var footer = std.ArrayList(u8).init(self.allocator);
        defer footer.deinit();
        const fw = encode.ByteWriter.init(&footer);
        try fw.putU32(magic);
        try fw.putU32(version);
        try fw.putU64(indexOffset);
        try fw.putU32(indexLen);
        try fw.putU64(bloomOffset);
        try fw.putU32(bloomLen);
        try fw.putU64(self.numEntries);
        try fw.putU64(self.maxSeq);
        try self.file.writeAll(footer.items);
        try self.file.sync();
    }

    fn writeBlock(self: *Builder, bytes: []const u8) !void {
        try self.file.writeAll(bytes);
        var crcBuf: [4]u8 = undefined;
        std.mem.writeInt(u32, &crcBuf, encode.crc(bytes), encode.endian);
        try self.file.writeAll(&crcBuf);
        self.offset += bytes.len + 4;
    }
};

const IndexEntry = struct {
    offset: u64,
    blockLen: u32,
    firstKey: []const u8,
};

const blockPrefix = 1 + 4;

pub const Reader = struct {
    allocator: Allocator,
    buffer: []u8,
    dataBuf: []u8,
    index: []IndexEntry,
    bloom: Bloom,
    fileId: u64,
    maxSeq: u64,
    numEntries: u64,
    physicalBytes: u64,
    level: usize,
    minKey: []const u8,
    maxKey: []const u8,

    pub fn open(allocator: Allocator, dir: std.fs.Dir, name: []const u8, fileId: u64) !Reader {
        const file = try dir.openFile(name, .{});
        defer file.close();
        const buffer = try file.readToEndAlloc(allocator, std.math.maxInt(usize));
        errdefer allocator.free(buffer);
        const physicalBytes = buffer.len;

        if (buffer.len < footerLen) return error.CorruptSstable;
        const footerStart = buffer.len - footerLen;
        if (readU32(buffer, footerStart) != magic) return error.CorruptSstable;
        const indexOffset = readU64(buffer, footerStart + 8);
        const indexLen = readU32(buffer, footerStart + 16);
        const bloomOffset = readU64(buffer, footerStart + 20);
        const bloomLen = readU32(buffer, footerStart + 28);
        const numEntries = readU64(buffer, footerStart + 32);
        const maxSeq = readU64(buffer, footerStart + 40);

        const bloomData = try verifiedBlock(buffer, bloomOffset, bloomLen);
        var bloom = try Bloom.deserialize(allocator, bloomData);
        errdefer bloom.deinit();

        const indexData = try verifiedBlock(buffer, indexOffset, indexLen);

        var decoded = std.ArrayList(u8).init(allocator);
        errdefer decoded.deinit();
        var entries = std.ArrayList(IndexEntry).init(allocator);
        errdefer entries.deinit();

        var pos: usize = 0;
        while (pos < indexData.len) {
            const physOffset = readU64(indexData, pos);
            const physLen = readU32(indexData, pos + 8);
            const firstKeyLen = readU32(indexData, pos + 12);
            const keyStart = pos + 16;
            const firstKey = indexData[keyStart .. keyStart + firstKeyLen];

            const physical = try verifiedBlock(buffer, physOffset, physLen);
            if (physical.len < blockPrefix) return error.CorruptSstable;
            const codec: Codec = @enumFromInt(physical[0]);
            const originalLen = readU32(physical, 1);
            const payload = physical[blockPrefix..];
            const plainBlock = try compress.decode(allocator, codec, payload, originalLen);
            defer allocator.free(plainBlock);

            const decOffset = decoded.items.len;
            try decoded.appendSlice(plainBlock);

            try entries.append(.{ .offset = decOffset, .blockLen = @intCast(plainBlock.len), .firstKey = firstKey });
            pos = keyStart + firstKeyLen;
        }

        const dataBuf = try decoded.toOwnedSlice();
        errdefer allocator.free(dataBuf);
        const indexSlice = try entries.toOwnedSlice();
        errdefer allocator.free(indexSlice);

        var minKey: []const u8 = buffer[0..0];
        var maxKey: []const u8 = buffer[0..0];
        if (indexSlice.len > 0) {
            minKey = indexSlice[0].firstKey;
            const lastEntry = indexSlice[indexSlice.len - 1];
            const lastBlock = dataBuf[lastEntry.offset .. lastEntry.offset + lastEntry.blockLen];
            var blockPos: usize = 0;
            while (blockPos < lastBlock.len) {
                const rec = decodeRecord(lastBlock, blockPos);
                maxKey = rec.record.key;
                blockPos = rec.next;
            }
        }

        return .{
            .allocator = allocator,
            .buffer = buffer,
            .dataBuf = dataBuf,
            .index = indexSlice,
            .bloom = bloom,
            .fileId = fileId,
            .maxSeq = maxSeq,
            .numEntries = numEntries,
            .physicalBytes = physicalBytes,
            .level = 0,
            .minKey = minKey,
            .maxKey = maxKey,
        };
    }

    pub fn deinit(self: *Reader) void {
        self.allocator.free(self.index);
        self.bloom.deinit();
        self.allocator.free(self.dataBuf);
        self.allocator.free(self.buffer);
    }

    pub fn mayContain(self: *const Reader, key: []const u8) bool {
        return self.bloom.mayContain(key);
    }

    pub fn find(self: *const Reader, key: []const u8) ?Record {
        if (!self.bloom.mayContain(key)) return null;
        const blockIdx = self.findBlock(key) orelse return null;
        const block = self.blockBytes(blockIdx);
        var pos: usize = 0;
        while (pos < block.len) {
            const decoded = decodeRecord(block, pos);
            switch (order(decoded.record.key, key)) {
                .eq => return decoded.record,
                .gt => return null,
                .lt => pos = decoded.next,
            }
        }
        return null;
    }

    pub fn findAsOf(self: *const Reader, key: []const u8, maxSeq: u64) ?Record {
        const rec = self.find(key) orelse return null;
        if (rec.seq > maxSeq) return null;
        return rec;
    }

    pub fn findCachedAsOf(self: *const Reader, blockCache: *BlockCache, allocator: Allocator, key: []const u8, maxSeq: u64) !?OwnedHit {
        if (!self.bloom.mayContain(key)) return null;
        const blockIdx = self.findBlock(key) orelse return null;
        const entry = self.index[blockIdx];
        const blockKey = BlockKey{ .fileId = self.fileId, .offset = entry.offset };

        const cached = blockCache.get(blockKey);
        const block = cached orelse self.dataBuf[entry.offset .. entry.offset + entry.blockLen];
        if (cached == null) try blockCache.put(blockKey, block);

        var pos: usize = 0;
        while (pos < block.len) {
            const decoded = decodeRecord(block, pos);
            switch (order(decoded.record.key, key)) {
                .eq => {
                    if (decoded.record.seq > maxSeq) return null;
                    return OwnedHit{
                        .seq = decoded.record.seq,
                        .kind = decoded.record.kind,
                        .value = try allocator.dupe(u8, decoded.record.value),
                    };
                },
                .gt => return null,
                .lt => pos = decoded.next,
            }
        }
        return null;
    }

    pub fn cursor(self: *const Reader, start: ?[]const u8) Cursor {
        var c = Cursor{ .reader = self, .blockIdx = 0, .pos = 0, .nextPos = 0, .cur = null };
        if (self.index.len == 0) return c;
        c.blockIdx = if (start) |s| (self.findBlock(s) orelse 0) else 0;
        c.refresh();
        if (start) |s| {
            while (c.valid() and order(c.cur.?.key, s) == .lt) c.advance();
        }
        return c;
    }

    pub fn reverseCursor(self: *const Reader, allocator: Allocator, end: ?[]const u8, endInclusive: bool) !ReverseCursor {
        var c = ReverseCursor{
            .reader = self,
            .allocator = allocator,
            .positions = std.ArrayList(usize).init(allocator),
            .blockIdx = 0,
            .idx = 0,
            .exhausted = false,
            .cur = null,
        };
        if (self.index.len == 0) {
            c.exhausted = true;
            return c;
        }

        var blockIdx: usize = if (end) |e| (self.findBlock(e) orelse {
            c.exhausted = true;
            return c;
        }) else self.index.len - 1;

        try c.loadBlock(blockIdx);
        if (end) |e| {
            var chosen: ?usize = null;
            for (c.positions.items, 0..) |p, i| {
                const rec = decodeRecord(self.blockBytes(blockIdx), p);
                const within = if (endInclusive) order(rec.record.key, e) != .gt else order(rec.record.key, e) == .lt;
                if (within) chosen = i else break;
            }
            if (chosen) |i| {
                c.idx = i;
            } else {
                if (blockIdx == 0) {
                    c.exhausted = true;
                    c.cur = null;
                    return c;
                }
                blockIdx -= 1;
                try c.loadBlock(blockIdx);
                c.idx = c.positions.items.len - 1;
            }
        } else {
            c.idx = c.positions.items.len - 1;
        }
        c.setCurrent();
        return c;
    }

    fn findBlock(self: *const Reader, key: []const u8) ?usize {
        var lo: usize = 0;
        var hi: usize = self.index.len;
        var ans: ?usize = null;
        while (lo < hi) {
            const mid = lo + (hi - lo) / 2;
            if (order(self.index[mid].firstKey, key) != .gt) {
                ans = mid;
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return ans;
    }

    fn blockBytes(self: *const Reader, idx: usize) []const u8 {
        const entry = self.index[idx];
        return self.dataBuf[entry.offset .. entry.offset + entry.blockLen];
    }

    pub fn blockCount(self: *const Reader) usize {
        return self.index.len;
    }
};

pub const Cursor = struct {
    reader: *const Reader,
    blockIdx: usize,
    pos: usize,
    nextPos: usize,
    cur: ?Record,

    pub fn valid(self: *const Cursor) bool {
        return self.cur != null;
    }

    pub fn record(self: *const Cursor) Record {
        return self.cur.?;
    }

    pub fn advance(self: *Cursor) void {
        self.pos = self.nextPos;
        self.refresh();
    }

    fn refresh(self: *Cursor) void {
        while (self.blockIdx < self.reader.index.len) {
            const block = self.reader.blockBytes(self.blockIdx);
            if (self.pos >= block.len) {
                self.blockIdx += 1;
                self.pos = 0;
                continue;
            }
            const decoded = decodeRecord(block, self.pos);
            self.cur = decoded.record;
            self.nextPos = decoded.next;
            return;
        }
        self.cur = null;
    }
};

pub const ReverseCursor = struct {
    reader: *const Reader,
    allocator: Allocator,
    positions: std.ArrayList(usize),
    blockIdx: usize,
    idx: usize,
    exhausted: bool,
    cur: ?Record,

    pub fn deinit(self: *ReverseCursor) void {
        self.positions.deinit();
    }

    pub fn valid(self: *const ReverseCursor) bool {
        return self.cur != null;
    }

    pub fn record(self: *const ReverseCursor) Record {
        return self.cur.?;
    }

    pub fn advance(self: *ReverseCursor) void {
        if (self.exhausted) return;
        if (self.idx == 0) {
            if (self.blockIdx == 0) {
                self.exhausted = true;
                self.cur = null;
                return;
            }
            self.loadBlock(self.blockIdx - 1) catch {
                self.exhausted = true;
                self.cur = null;
                return;
            };
            self.idx = self.positions.items.len - 1;
        } else {
            self.idx -= 1;
        }
        self.setCurrent();
    }

    fn loadBlock(self: *ReverseCursor, blockIdx: usize) !void {
        self.blockIdx = blockIdx;
        self.positions.clearRetainingCapacity();
        const block = self.reader.blockBytes(blockIdx);
        var pos: usize = 0;
        while (pos < block.len) {
            try self.positions.append(pos);
            pos = decodeRecord(block, pos).next;
        }
    }

    fn setCurrent(self: *ReverseCursor) void {
        const block = self.reader.blockBytes(self.blockIdx);
        self.cur = decodeRecord(block, self.positions.items[self.idx]).record;
    }
};

fn verifiedBlock(buffer: []const u8, offset: u64, len: u32) ![]const u8 {
    const start: usize = @intCast(offset);
    const end = start + len;
    if (end + 4 > buffer.len) return error.CorruptSstable;
    const block = buffer[start..end];
    if (encode.crc(block) != readU32(buffer, end)) return error.CorruptSstable;
    return block;
}

const Decoded = struct {
    record: Record,
    next: usize,
};

fn decodeRecord(block: []const u8, pos: usize) Decoded {
    const kind: Kind = @enumFromInt(block[pos]);
    const seq = readU64(block, pos + 1);
    const keyLen = readU32(block, pos + 9);
    const valLen = readU32(block, pos + 13);
    const keyStart = pos + recordHeader;
    const key = block[keyStart .. keyStart + keyLen];
    const valStart = keyStart + keyLen;
    const value = block[valStart .. valStart + valLen];
    return .{
        .record = .{ .key = key, .value = value, .seq = seq, .kind = kind },
        .next = valStart + valLen,
    };
}