kb_exfiltrator/src/main.zig

const std = @import("std");
const snapshot = @import("snapshot.zig");
const c = @cImport({
    @cInclude("sqlite3.h");
});

// Helper function to clean highlight text by removing page/margin reference numbers
// Removes patterns like "10After" -> "After" and " 5good" -> " good"
fn cleanHighlightText(allocator: std.mem.Allocator, text: []const u8) ![]const u8 {
    var result = std.ArrayList(u8).initCapacity(allocator, 0) catch unreachable;
    defer result.deinit(allocator);
    const writer = result.writer(allocator);

    var i: usize = 0;
    while (i < text.len) {
        // Check if we're at start of text or after whitespace/newline
        const at_boundary = (i == 0) or (text[i - 1] == ' ') or (text[i - 1] == '\n') or (text[i - 1] == '\t');

        if (at_boundary and std.ascii.isDigit(text[i])) {
            // Count consecutive digits
            var digit_end = i;
            while (digit_end < text.len and std.ascii.isDigit(text[digit_end])) {
                digit_end += 1;
            }

            // If digits are followed by a letter (like "10After"), skip the digits
            if (digit_end < text.len and std.ascii.isAlphabetic(text[digit_end])) {
                i = digit_end; // Skip the digits
                continue;
            }
        }

        try writer.writeByte(text[i]);
        i += 1;
    }

    return try allocator.dupe(u8, result.items);
}

// Helper function to escape JSON strings
fn escapeJson(allocator: std.mem.Allocator, s: []const u8) ![]const u8 {
    var result = std.ArrayList(u8).initCapacity(allocator, 0) catch unreachable;
    defer result.deinit(allocator);
    const writer = result.writer(allocator);

    try writer.writeAll("\"");
    for (s) |ch| {
        switch (ch) {
            '"' => try writer.writeAll("\\\""),
            '\\' => try writer.writeAll("\\\\"),
            '\n' => try writer.writeAll("\\n"),
            '\r' => try writer.writeAll("\\r"),
            '\t' => try writer.writeAll("\\t"),
            else => try writer.writeByte(ch),
        }
    }
    try writer.writeAll("\"");

    return try allocator.dupe(u8, result.items);
}

const Highlight = struct {
    text: []const u8, // Owned copy
    date: []const u8, // Owned copy
    book: []const u8, // Owned copy
    text_hash: [64]u8, // Stack allocated - no cleanup needed
    book_hash: [64]u8, // Stack allocated - no cleanup needed
    allocator: std.mem.Allocator,

    // JSON-serializable version (no allocator field)
    const JsonHighlight = struct {
        text: []const u8,
        date: []const u8,
        book: []const u8,
        text_hash: []const u8, // Slice view of the hash array
        book_hash: []const u8, // Slice view of the hash array
    };

    fn toJson(self: *const Highlight) JsonHighlight {
        return JsonHighlight{
            .text = self.text,
            .date = self.date,
            .book = self.book,
            .text_hash = &self.text_hash,
            .book_hash = &self.book_hash,
        };
    }

    fn init(allocator: std.mem.Allocator, text: []const u8, date: []const u8, book: []const u8, content_id_raw: []const u8) !Highlight {
        // Clean the text to remove page reference numbers
        const cleaned_text = try cleanHighlightText(allocator, text);
        defer allocator.free(cleaned_text);

        var highlight = Highlight{
            .text = try allocator.dupe(u8, cleaned_text), // Copy cleaned text
            .date = try allocator.dupe(u8, date), // Copy from SQLite buffer
            .book = try allocator.dupe(u8, book), // Copy from SQLite buffer
            .text_hash = undefined,
            .book_hash = undefined,
            .allocator = allocator,
        };

        // Now compute hashes into the struct fields (use cleaned text for hash)
        computeHash(cleaned_text, &highlight.text_hash);
        computeHash(content_id_raw, &highlight.book_hash);

        return highlight;
    }

    fn deinit(self: *Highlight) void {
        self.allocator.free(self.text);
        self.allocator.free(self.date);
        self.allocator.free(self.book);
    }

    fn computeHash(data: []const u8, out_hex: *[64]u8) void {
        var hash: [32]u8 = undefined; // SHA256 produces 32 bytes
        std.crypto.hash.sha2.Sha256.hash(data, &hash, .{});

        // Convert 32 bytes to 64 hex characters
        const hex_chars = "0123456789abcdef";
        for (hash, 0..) |byte, i| {
            out_hex[i * 2] = hex_chars[byte >> 4];
            out_hex[i * 2 + 1] = hex_chars[byte & 0x0F];
        }
    }
};

// Zig only collects `test` blocks from the test unit's root file, so pull in the
// snapshot module's tests explicitly.
test {
    _ = @import("snapshot.zig");
}

const usage =
    \\Usage:
    \\  kb_exfiltrator snapshot   <db> [endpoint] [timeout-secs]
    \\      Upload a consistent whole-database snapshot (VACUUM INTO). This is the
    \\      pipeline's primary mode: the server parses it, so adding new stats never
    \\      requires touching the device. Omit endpoint to leave the snapshot on disk.
    \\
    \\  kb_exfiltrator highlights <db> [endpoint]
    \\      Legacy mode: extract highlights and POST them as JSON.
    \\
    \\  kb_exfiltrator <db> [endpoint]
    \\      Backwards-compatible alias for `highlights`.
    \\
;

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

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

    if (args.len < 2) {
        std.debug.print("{s}", .{usage});
        return error.MissingArgument;
    }

    if (std.mem.eql(u8, args[1], "snapshot")) {
        if (args.len < 3) {
            std.debug.print("{s}", .{usage});
            return error.MissingArgument;
        }
        const timeout = if (args.len >= 5)
            std.fmt.parseInt(u64, args[4], 10) catch snapshot.default_timeout_secs
        else
            snapshot.default_timeout_secs;
        try snapshot.startWatchdog(timeout);
        return snapshot.run(
            allocator,
            args[2],
            if (args.len >= 4) args[3] else null,
            snapshot.default_snapshot_path,
        );
    }

    // Legacy mode uploads too, so it needs the same protection: the endpoint it posts to
    // is what hung a run for 16 days.
    try snapshot.startWatchdog(snapshot.default_timeout_secs);

    // `highlights <db> [endpoint]`, or the legacy `<db> [endpoint]` form.
    const legacy = std.mem.eql(u8, args[1], "highlights");
    const rest = if (legacy) args[2..] else args[1..];
    if (rest.len < 1) {
        std.debug.print("{s}", .{usage});
        return error.MissingArgument;
    }
    return runHighlights(allocator, rest[0], if (rest.len >= 2) rest[1] else null);
}

fn runHighlights(
    allocator: std.mem.Allocator,
    db_path: []const u8,
    endpoint_url: ?[]const u8,
) !void {
    // Open database. db_path arrives as a plain slice, so re-terminate it for the C API.
    const db_path_z = try allocator.dupeZ(u8, db_path);
    defer allocator.free(db_path_z);

    var db: ?*c.sqlite3 = null;
    const rc = c.sqlite3_open(db_path_z.ptr, &db);
    if (rc != c.SQLITE_OK) {
        std.debug.print("Failed to open database: {s}\n", .{c.sqlite3_errmsg(db)});
        return error.DatabaseOpenFailed;
    }
    defer _ = c.sqlite3_close(db);

    // Prepare query to get highlights
    const query =
        \\SELECT
        \\  Text,
        \\  DateCreated,
        \\  SUBSTR(
        \\    SUBSTR(ContentID, 40, INSTR(ContentID, '.epub') - 40 + 5),
        \\    INSTR(SUBSTR(ContentID, 40, INSTR(ContentID, '.epub') - 40 + 5), '/') + 1
        \\  ) AS Book,
        \\  ContentID
        \\FROM Bookmark
        \\WHERE Text IS NOT NULL
        \\ORDER BY DateCreated DESC;
    ;

    var stmt: ?*c.sqlite3_stmt = null;
    if (c.sqlite3_prepare_v2(db, query.ptr, @intCast(query.len), &stmt, null) != c.SQLITE_OK) {
        std.debug.print("Failed to prepare statement: {s}\n", .{c.sqlite3_errmsg(db)});
        return error.QueryPrepareFailed;
    }
    defer _ = c.sqlite3_finalize(stmt);

    // Create ArrayList to collect highlights
    var highlights = std.ArrayList(Highlight).initCapacity(allocator, 0) catch unreachable;
    defer {
        for (highlights.items) |*highlight| {
            highlight.deinit();
        }
        highlights.deinit(allocator);
    }

    // Execute query and collect results
    while (c.sqlite3_step(stmt) == c.SQLITE_ROW) {

        // Get the text (highlight) - convert C pointer to Zig slice
        const text_ptr = c.sqlite3_column_text(stmt, 0);
        const text = if (text_ptr != null)
            std.mem.span(@as([*:0]const u8, @ptrCast(text_ptr)))
        else
            "";

        // Get date - convert C pointer to Zig slice
        const date_ptr = c.sqlite3_column_text(stmt, 1);
        const date = if (date_ptr != null)
            std.mem.span(@as([*:0]const u8, @ptrCast(date_ptr)))
        else
            "unknown";

        // Get book title - convert C pointer to Zig slice
        const book_ptr = c.sqlite3_column_text(stmt, 2);
        const book = if (book_ptr != null)
            std.mem.span(@as([*:0]const u8, @ptrCast(book_ptr)))
        else
            "unknown";

        // Get raw ContentID for hash - convert C pointer to Zig slice
        const content_id_ptr = c.sqlite3_column_text(stmt, 3);
        const content_id = if (content_id_ptr != null)
            std.mem.span(@as([*:0]const u8, @ptrCast(content_id_ptr)))
        else
            "";

        // Create Highlight struct (makes owned copies and computes hashes)
        const highlight = try Highlight.init(allocator, text, date, book, content_id);

        // Add to ArrayList
        try highlights.append(allocator, highlight);
    }

    std.debug.print("Extracted {d} highlights from database\n", .{highlights.items.len});

    // Build JSON manually
    var json_buffer = std.ArrayList(u8).initCapacity(allocator, 0) catch unreachable;
    defer json_buffer.deinit(allocator);
    const writer = json_buffer.writer(allocator);

    try writer.writeAll("[");
    for (highlights.items, 0..) |*highlight, i| {
        if (i > 0) try writer.writeAll(",");
        try writer.writeAll("{");

        const escaped_text = try escapeJson(allocator, highlight.text);
        defer allocator.free(escaped_text);
        const escaped_book = try escapeJson(allocator, highlight.book);
        defer allocator.free(escaped_book);

        try writer.print("\"text\":{s},", .{escaped_text});
        try writer.print("\"date\":\"{s}\",", .{highlight.date});
        try writer.print("\"book\":{s},", .{escaped_book});
        try writer.print("\"text_hash\":\"{s}\",", .{highlight.text_hash});
        try writer.print("\"book_hash\":\"{s}\"", .{highlight.book_hash});
        try writer.writeAll("}");
    }
    try writer.writeAll("]");

    // If endpoint URL provided, POST the JSON; otherwise just print it
    if (endpoint_url) |url| {
        std.debug.print("Sending {d} highlights ({d} bytes) to {s}...\n", .{highlights.items.len, json_buffer.items.len, url});

        var client: std.http.Client = .{ .allocator = allocator };
        defer client.deinit();

        const uri = try std.Uri.parse(url);

        const result = try client.fetch(.{
            .location = .{ .uri = uri },
            .method = .POST,
            .payload = json_buffer.items,
            .headers = .{
                .content_type = .{ .override = "application/json" },
            },
        });

        const status_code = @intFromEnum(result.status);
        if (status_code >= 200 and status_code < 300) {
            std.debug.print("✓ Successfully sent highlights to server (HTTP {})\n", .{status_code});
        } else {
            std.debug.print("✗ Server returned error: HTTP {}\n", .{status_code});
            return error.HttpRequestFailed;
        }
    } else {
        std.debug.print("No endpoint provided. JSON payload ({d} bytes):\n{s}\n", .{json_buffer.items.len, json_buffer.items});
    }
}