const SQLite = @This();

io: Io,
allocator: Allocator,
/// Pool of read-only connections. Reads run concurrently across them (WAL).
readers: []Connection,
avail: []*Connection,
avail_count: usize,
pool_mutex: Mutex = .init,
pool_cond: Condition = .init,
/// Single writer connection. WAL permits only one writer, so all writes
/// serialize here behind `writer_mutex` — this avoids the SQLITE_BUSY retry
/// churn that many competing writer connections would incur.
writer: Connection,
writer_mutex: Mutex = .init,
/// Interface
store: Store = .{
    .vtable = &.{
        .get = get,
        .put = put,
        .putExpire = putExpire,
        .fetchRemove = fetchRemove,
        .remove = remove,
        .append = append,
        .prepend = prepend,
        .pop = pop,
        .popFirst = popFirst,
    },
},

pub const Config = struct {
    path: [:0]const u8,
    truncate: bool = false,
    /// Number of read connections in the pool (a separate writer connection
    /// is always created). Defaults to the CPU count when null.
    pool_size: ?usize = null,
    busy_timeout_ms: c_int = 5000,
};

/// A connection and its per-connection prepared statements. Prepared
/// statements belong to a single connection, so each connection owns its own
/// set. The same layout is used for both reader and writer connections.
const Connection = struct {
    conn: zqlite.Conn,
    get_stmt: zqlite.Stmt,
    put_stmt: zqlite.Stmt,
    fetch_stmt: zqlite.Stmt,
    remove_stmt: zqlite.Stmt,
    append_stmt: zqlite.Stmt,
    prepend_stmt: zqlite.Stmt,
    pop_last_stmt: zqlite.Stmt,
    pop_first_stmt: zqlite.Stmt,
    array_delete_stmt: zqlite.Stmt,

    fn open(config: Config) !Connection {
        const flags = zqlite.c.SQLITE_OPEN_CREATE | zqlite.c.SQLITE_OPEN_READWRITE;
        const conn = try zqlite.open(config.path, flags);
        errdefer conn.close();

        try conn.busyTimeout(config.busy_timeout_ms);
        try conn.execNoArgs("PRAGMA journal_mode=WAL");
        try conn.execNoArgs("PRAGMA synchronous=NORMAL");
        try conn.execNoArgs("PRAGMA temp_store=MEMORY");
        try conn.execNoArgs("PRAGMA mmap_size=268435456");
        try conn.execNoArgs(
            "CREATE TABLE IF NOT EXISTS kv_strings" ++
                "(key TEXT PRIMARY KEY, value TEXT NOT NULL, expiry INTEGER)",
        );
        try conn.execNoArgs(
            "CREATE TABLE IF NOT EXISTS kv_arrays" ++
                "(id INTEGER PRIMARY KEY AUTOINCREMENT, key TEXT NOT NULL, value TEXT NOT NULL, position INTEGER NOT NULL)",
        );
        try conn.execNoArgs("CREATE INDEX IF NOT EXISTS kv_arrays_key_pos ON kv_arrays (key, position)");

        return .{
            .conn = conn,
            .get_stmt = try conn.prepare("SELECT value, expiry FROM kv_strings WHERE key = ?1"),
            .put_stmt = try conn.prepare("INSERT OR REPLACE INTO kv_strings (key, value, expiry) VALUES (?1, ?2, ?3)"),
            .fetch_stmt = try conn.prepare("SELECT value FROM kv_strings WHERE key = ?1"),
            .remove_stmt = try conn.prepare("DELETE FROM kv_strings WHERE key = ?1"),
            .append_stmt = try conn.prepare("INSERT INTO kv_arrays (key, value, position) VALUES (?1, ?2, COALESCE((SELECT MAX(position) FROM kv_arrays WHERE key = ?1), -1) + 1)"),
            .prepend_stmt = try conn.prepare("INSERT INTO kv_arrays (key, value, position) VALUES (?1, ?2, COALESCE((SELECT MIN(position) FROM kv_arrays WHERE key = ?1), 1) - 1)"),
            .pop_last_stmt = try conn.prepare("SELECT id, value FROM kv_arrays WHERE key = ?1 ORDER BY position DESC LIMIT 1"),
            .pop_first_stmt = try conn.prepare("SELECT id, value FROM kv_arrays WHERE key = ?1 ORDER BY position ASC LIMIT 1"),
            .array_delete_stmt = try conn.prepare("DELETE FROM kv_arrays WHERE id = ?1"),
        };
    }

    fn deinit(self: *Connection) void {
        self.get_stmt.deinit();
        self.put_stmt.deinit();
        self.fetch_stmt.deinit();
        self.remove_stmt.deinit();
        self.append_stmt.deinit();
        self.prepend_stmt.deinit();
        self.pop_last_stmt.deinit();
        self.pop_first_stmt.deinit();
        self.array_delete_stmt.deinit();
        self.conn.close();
    }
};

pub fn init(io: Io, allocator: Allocator, comptime config: Config) !SQLite {
    if (config.truncate) {
        Dir.deleteFileAbsolute(io, config.path) catch |err| switch (err) {
            error.FileNotFound => {},
            else => return err,
        };
    }

    const size = config.pool_size orelse (std.Thread.getCpuCount() catch 4);
    const readers = try allocator.alloc(Connection, size);
    errdefer allocator.free(readers);
    const avail = try allocator.alloc(*Connection, size);
    errdefer allocator.free(avail);

    var opened: usize = 0;
    errdefer for (0..opened) |i| readers[i].deinit();
    while (opened < size) : (opened += 1) {
        readers[opened] = try Connection.open(config);
    }

    for (readers, 0..) |*c, i| avail[i] = c;

    const writer = try Connection.open(config);

    return .{
        .io = io,
        .allocator = allocator,
        .readers = readers,
        .avail = avail,
        .avail_count = size,
        .writer = writer,
    };
}

pub fn deinit(self: *SQLite) void {
    for (self.readers) |*c| c.deinit();
    self.writer.deinit();
    self.allocator.free(self.readers);
    self.allocator.free(self.avail);
    self.* = undefined;
}

/// Check out a read connection, blocking until one is free.
fn acquireReader(self: *SQLite) *Connection {
    self.pool_mutex.lockUncancelable(self.io);
    while (self.avail_count == 0) self.pool_cond.waitUncancelable(self.io, &self.pool_mutex);
    self.avail_count -= 1;
    const c = self.avail[self.avail_count];
    self.pool_mutex.unlock(self.io);
    return c;
}

fn releaseReader(self: *SQLite, c: *Connection) void {
    self.pool_mutex.lockUncancelable(self.io);
    self.avail[self.avail_count] = c;
    self.avail_count += 1;
    self.pool_mutex.unlock(self.io);
    self.pool_cond.signal(self.io);
}

fn get(s: *Store, allocator: Allocator, key: []const u8) !?[]const u8 {
    const self: *SQLite = @fieldParentPtr("store", s);
    const c = self.acquireReader();
    defer self.releaseReader(c);

    const stmt = c.get_stmt;
    defer stmt.reset() catch {};
    try stmt.bind(.{key});
    if (try stmt.step() == false) return null;

    if (stmt.nullableInt(1)) |expiry| {
        const now = Io.Timestamp.now(self.io, .real).toMilliseconds();
        if (expiry < now) return null;
    }

    return try allocator.dupe(u8, stmt.text(0));
}

fn put(s: *Store, key: []const u8, value: []const u8) !void {
    const self: *SQLite = @fieldParentPtr("store", s);
    try self.putMaybeExpire(key, value, null);
}

fn putExpire(s: *Store, key: []const u8, value: []const u8, expiration: i32) !void {
    const self: *SQLite = @fieldParentPtr("store", s);
    const now = Io.Timestamp.now(self.io, .real).toMilliseconds();
    try self.putMaybeExpire(key, value, now + (@as(i64, expiration) * std.time.ms_per_s));
}

fn putMaybeExpire(self: *SQLite, key: []const u8, value: []const u8, expiry: ?i64) !void {
    self.writer_mutex.lockUncancelable(self.io);
    defer self.writer_mutex.unlock(self.io);

    const stmt = self.writer.put_stmt;
    defer stmt.reset() catch {};
    try stmt.bind(.{ key, value, expiry });
    try stmt.stepToCompletion();
}

fn fetchRemove(s: *Store, allocator: Allocator, key: []const u8) !?[]const u8 {
    const self: *SQLite = @fieldParentPtr("store", s);
    self.writer_mutex.lockUncancelable(self.io);
    defer self.writer_mutex.unlock(self.io);
    const c = &self.writer;

    const sel = c.fetch_stmt;
    defer sel.reset() catch {};
    try sel.bind(.{key});
    if (try sel.step() == false) return null;

    const value = try allocator.dupe(u8, sel.text(0));
    errdefer allocator.free(value);
    sel.reset() catch {};

    const del = c.remove_stmt;
    defer del.reset() catch {};
    try del.bind(.{key});
    try del.stepToCompletion();

    return value;
}

fn remove(s: *Store, key: []const u8) !void {
    const self: *SQLite = @fieldParentPtr("store", s);
    self.writer_mutex.lockUncancelable(self.io);
    defer self.writer_mutex.unlock(self.io);

    const stmt = self.writer.remove_stmt;
    defer stmt.reset() catch {};
    try stmt.bind(.{key});
    try stmt.stepToCompletion();
}

fn append(s: *Store, key: []const u8, value: []const u8) !void {
    const self: *SQLite = @fieldParentPtr("store", s);
    self.writer_mutex.lockUncancelable(self.io);
    defer self.writer_mutex.unlock(self.io);

    const stmt = self.writer.append_stmt;
    defer stmt.reset() catch {};
    try stmt.bind(.{ key, value });
    try stmt.stepToCompletion();
}

fn prepend(s: *Store, key: []const u8, value: []const u8) !void {
    const self: *SQLite = @fieldParentPtr("store", s);
    self.writer_mutex.lockUncancelable(self.io);
    defer self.writer_mutex.unlock(self.io);

    const stmt = self.writer.prepend_stmt;
    defer stmt.reset() catch {};
    try stmt.bind(.{ key, value });
    try stmt.stepToCompletion();
}

fn pop(s: *Store, allocator: Allocator, key: []const u8) !?[]const u8 {
    const self: *SQLite = @fieldParentPtr("store", s);
    self.writer_mutex.lockUncancelable(self.io);
    defer self.writer_mutex.unlock(self.io);
    return popStmt(&self.writer, allocator, self.writer.pop_last_stmt, key);
}

fn popFirst(s: *Store, allocator: Allocator, key: []const u8) !?[]const u8 {
    const self: *SQLite = @fieldParentPtr("store", s);
    self.writer_mutex.lockUncancelable(self.io);
    defer self.writer_mutex.unlock(self.io);
    return popStmt(&self.writer, allocator, self.writer.pop_first_stmt, key);
}

fn popStmt(c: *Connection, allocator: Allocator, sel: zqlite.Stmt, key: []const u8) !?[]const u8 {
    defer sel.reset() catch {};
    try sel.bind(.{key});
    if (try sel.step() == false) return null;

    const row_id = sel.int(0);
    const value = try allocator.dupe(u8, sel.text(1));
    errdefer allocator.free(value);
    sel.reset() catch {};

    const del = c.array_delete_stmt;
    defer del.reset() catch {};
    try del.bind(.{row_id});
    try del.stepToCompletion();

    return value;
}

const t = std.testing;
test "sqlite overwrite" {
    var kv: SQLite = try .init(t.io, t.allocator, .{
        .path = "/tmp/jetkv_sqlite.db",
        .truncate = true,
    });
    defer kv.deinit();
    const backend = &kv.store;

    try backend.put("foo", "bar");
    try backend.put("foo", "baz");

    if (try backend.get(t.allocator, "foo")) |v| {
        defer t.allocator.free(v);
        try t.expectEqualStrings("baz", v);
    } else try t.expect(false);
}

const std = @import("std");
const Io = std.Io;
const Dir = Io.Dir;
const Mutex = Io.Mutex;
const Condition = Io.Condition;
const Allocator = std.mem.Allocator;
const Store = @import("Store.zig");
const zqlite = @import("zqlite");