zcatsql/src/rtree.zig

//! R-Tree Spatial Index Helpers
//!
//! Provides convenient wrappers for SQLite's R-Tree extension for spatial indexing.
//! R-Tree is compiled into SQLite by default with -DSQLITE_ENABLE_RTREE.
//!
//! R-Tree is useful for:
//! - Geographic data (latitude/longitude bounding boxes)
//! - 2D/3D spatial queries
//! - Range queries on multiple dimensions
//! - Game collision detection
//! - Any multi-dimensional range search

const std = @import("std");
const Database = @import("database.zig").Database;
const Statement = @import("statement.zig").Statement;
const Error = @import("errors.zig").Error;

/// A 2D bounding box (rectangle).
pub const BoundingBox2D = struct {
    min_x: f64,
    max_x: f64,
    min_y: f64,
    max_y: f64,

    /// Creates a bounding box from center point and size.
    pub fn fromCenter(center_x: f64, center_y: f64, width: f64, height: f64) BoundingBox2D {
        const half_w = width / 2.0;
        const half_h = height / 2.0;
        return .{
            .min_x = center_x - half_w,
            .max_x = center_x + half_w,
            .min_y = center_y - half_h,
            .max_y = center_y + half_h,
        };
    }

    /// Creates a bounding box from a point (zero-size box).
    pub fn fromPoint(x: f64, y: f64) BoundingBox2D {
        return .{
            .min_x = x,
            .max_x = x,
            .min_y = y,
            .max_y = y,
        };
    }

    /// Checks if this box intersects with another box.
    pub fn intersects(self: BoundingBox2D, other: BoundingBox2D) bool {
        return self.min_x <= other.max_x and
            self.max_x >= other.min_x and
            self.min_y <= other.max_y and
            self.max_y >= other.min_y;
    }

    /// Checks if this box contains a point.
    pub fn containsPoint(self: BoundingBox2D, x: f64, y: f64) bool {
        return x >= self.min_x and x <= self.max_x and
            y >= self.min_y and y <= self.max_y;
    }

    /// Checks if this box fully contains another box.
    pub fn contains(self: BoundingBox2D, other: BoundingBox2D) bool {
        return other.min_x >= self.min_x and
            other.max_x <= self.max_x and
            other.min_y >= self.min_y and
            other.max_y <= self.max_y;
    }

    /// Returns the area of the bounding box.
    pub fn area(self: BoundingBox2D) f64 {
        return (self.max_x - self.min_x) * (self.max_y - self.min_y);
    }

    /// Returns the center point of the bounding box.
    pub fn center(self: BoundingBox2D) struct { x: f64, y: f64 } {
        return .{
            .x = (self.min_x + self.max_x) / 2.0,
            .y = (self.min_y + self.max_y) / 2.0,
        };
    }

    /// Expands this box to include another box.
    pub fn expand(self: *BoundingBox2D, other: BoundingBox2D) void {
        self.min_x = @min(self.min_x, other.min_x);
        self.max_x = @max(self.max_x, other.max_x);
        self.min_y = @min(self.min_y, other.min_y);
        self.max_y = @max(self.max_y, other.max_y);
    }
};

/// A 3D bounding box.
pub const BoundingBox3D = struct {
    min_x: f64,
    max_x: f64,
    min_y: f64,
    max_y: f64,
    min_z: f64,
    max_z: f64,

    /// Creates a 3D bounding box from center point and size.
    pub fn fromCenter(center_x: f64, center_y: f64, center_z: f64, width: f64, height: f64, depth: f64) BoundingBox3D {
        const half_w = width / 2.0;
        const half_h = height / 2.0;
        const half_d = depth / 2.0;
        return .{
            .min_x = center_x - half_w,
            .max_x = center_x + half_w,
            .min_y = center_y - half_h,
            .max_y = center_y + half_h,
            .min_z = center_z - half_d,
            .max_z = center_z + half_d,
        };
    }

    /// Creates a bounding box from a point (zero-size box).
    pub fn fromPoint(x: f64, y: f64, z: f64) BoundingBox3D {
        return .{
            .min_x = x,
            .max_x = x,
            .min_y = y,
            .max_y = y,
            .min_z = z,
            .max_z = z,
        };
    }

    /// Checks if this box intersects with another box.
    pub fn intersects(self: BoundingBox3D, other: BoundingBox3D) bool {
        return self.min_x <= other.max_x and
            self.max_x >= other.min_x and
            self.min_y <= other.max_y and
            self.max_y >= other.min_y and
            self.min_z <= other.max_z and
            self.max_z >= other.min_z;
    }

    /// Returns the volume of the bounding box.
    pub fn volume(self: BoundingBox3D) f64 {
        return (self.max_x - self.min_x) * (self.max_y - self.min_y) * (self.max_z - self.min_z);
    }
};

/// Geographic coordinates (latitude/longitude).
pub const GeoCoord = struct {
    latitude: f64,
    longitude: f64,

    /// Creates a bounding box around this point with a given radius in degrees.
    pub fn boundingBox(self: GeoCoord, radius_degrees: f64) BoundingBox2D {
        return .{
            .min_x = self.longitude - radius_degrees,
            .max_x = self.longitude + radius_degrees,
            .min_y = self.latitude - radius_degrees,
            .max_y = self.latitude + radius_degrees,
        };
    }

    /// Approximate distance in kilometers to another coordinate (Haversine formula).
    pub fn distanceKm(self: GeoCoord, other: GeoCoord) f64 {
        const earth_radius_km = 6371.0;
        const lat1 = self.latitude * std.math.pi / 180.0;
        const lat2 = other.latitude * std.math.pi / 180.0;
        const delta_lat = (other.latitude - self.latitude) * std.math.pi / 180.0;
        const delta_lon = (other.longitude - self.longitude) * std.math.pi / 180.0;

        const a = std.math.sin(delta_lat / 2.0) * std.math.sin(delta_lat / 2.0) +
            std.math.cos(lat1) * std.math.cos(lat2) *
            std.math.sin(delta_lon / 2.0) * std.math.sin(delta_lon / 2.0);
        const c = 2.0 * std.math.atan2(std.math.sqrt(a), std.math.sqrt(1.0 - a));

        return earth_radius_km * c;
    }
};

/// R-Tree helper functions for a database connection.
pub const RTree = struct {
    db: *Database,
    allocator: std.mem.Allocator,

    const Self = @This();

    /// Creates a new R-Tree helper bound to a database connection.
    pub fn init(db: *Database, allocator: std.mem.Allocator) Self {
        return .{ .db = db, .allocator = allocator };
    }

    // ========================================================================
    // Table Management
    // ========================================================================

    /// Creates a 2D R-Tree virtual table.
    ///
    /// Example:
    /// ```zig
    /// var rtree = RTree.init(&db, allocator);
    /// try rtree.createTable2D("locations", "min_x", "max_x", "min_y", "max_y");
    /// ```
    pub fn createTable2D(
        self: *Self,
        table_name: []const u8,
        min_x_name: []const u8,
        max_x_name: []const u8,
        min_y_name: []const u8,
        max_y_name: []const u8,
    ) !void {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "CREATE VIRTUAL TABLE {s} USING rtree(id, {s}, {s}, {s}, {s})\x00",
            .{ table_name, min_x_name, max_x_name, min_y_name, max_y_name },
        );
        defer self.allocator.free(sql);
        try self.db.exec(sql[0 .. sql.len - 1 :0]);
    }

    /// Creates a 3D R-Tree virtual table.
    pub fn createTable3D(
        self: *Self,
        table_name: []const u8,
        min_x_name: []const u8,
        max_x_name: []const u8,
        min_y_name: []const u8,
        max_y_name: []const u8,
        min_z_name: []const u8,
        max_z_name: []const u8,
    ) !void {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "CREATE VIRTUAL TABLE {s} USING rtree(id, {s}, {s}, {s}, {s}, {s}, {s})\x00",
            .{ table_name, min_x_name, max_x_name, min_y_name, max_y_name, min_z_name, max_z_name },
        );
        defer self.allocator.free(sql);
        try self.db.exec(sql[0 .. sql.len - 1 :0]);
    }

    /// Creates a simple 2D R-Tree with default column names.
    pub fn createSimpleTable2D(self: *Self, table_name: []const u8) !void {
        try self.createTable2D(table_name, "min_x", "max_x", "min_y", "max_y");
    }

    /// Creates a simple 3D R-Tree with default column names.
    pub fn createSimpleTable3D(self: *Self, table_name: []const u8) !void {
        try self.createTable3D(table_name, "min_x", "max_x", "min_y", "max_y", "min_z", "max_z");
    }

    /// Creates a geographic R-Tree (using latitude/longitude).
    pub fn createGeoTable(self: *Self, table_name: []const u8) !void {
        try self.createTable2D(table_name, "min_lon", "max_lon", "min_lat", "max_lat");
    }

    /// Drops an R-Tree table.
    pub fn dropTable(self: *Self, table_name: []const u8) !void {
        const sql = try std.fmt.allocPrint(self.allocator, "DROP TABLE IF EXISTS {s}\x00", .{table_name});
        defer self.allocator.free(sql);
        try self.db.exec(sql[0 .. sql.len - 1 :0]);
    }

    // ========================================================================
    // Insert Operations
    // ========================================================================

    /// Inserts a 2D bounding box into an R-Tree.
    pub fn insert2D(self: *Self, table_name: []const u8, id: i64, box: BoundingBox2D) !void {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "INSERT INTO {s} VALUES (?, ?, ?, ?, ?)\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        defer stmt.finalize();

        try stmt.bindInt(1, id);
        try stmt.bindFloat(2, box.min_x);
        try stmt.bindFloat(3, box.max_x);
        try stmt.bindFloat(4, box.min_y);
        try stmt.bindFloat(5, box.max_y);

        _ = try stmt.step();
    }

    /// Inserts a 3D bounding box into an R-Tree.
    pub fn insert3D(self: *Self, table_name: []const u8, id: i64, box: BoundingBox3D) !void {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "INSERT INTO {s} VALUES (?, ?, ?, ?, ?, ?, ?)\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        defer stmt.finalize();

        try stmt.bindInt(1, id);
        try stmt.bindFloat(2, box.min_x);
        try stmt.bindFloat(3, box.max_x);
        try stmt.bindFloat(4, box.min_y);
        try stmt.bindFloat(5, box.max_y);
        try stmt.bindFloat(6, box.min_z);
        try stmt.bindFloat(7, box.max_z);

        _ = try stmt.step();
    }

    /// Inserts a point (as a zero-size bounding box) into a 2D R-Tree.
    pub fn insertPoint2D(self: *Self, table_name: []const u8, id: i64, x: f64, y: f64) !void {
        try self.insert2D(table_name, id, BoundingBox2D.fromPoint(x, y));
    }

    /// Inserts a geographic coordinate into a geo R-Tree.
    pub fn insertGeo(self: *Self, table_name: []const u8, id: i64, coord: GeoCoord) !void {
        try self.insert2D(table_name, id, .{
            .min_x = coord.longitude,
            .max_x = coord.longitude,
            .min_y = coord.latitude,
            .max_y = coord.latitude,
        });
    }

    // ========================================================================
    // Update Operations
    // ========================================================================

    /// Updates a 2D bounding box in an R-Tree.
    pub fn update2D(self: *Self, table_name: []const u8, id: i64, box: BoundingBox2D) !void {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "UPDATE {s} SET min_x=?, max_x=?, min_y=?, max_y=? WHERE id=?\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        defer stmt.finalize();

        try stmt.bindFloat(1, box.min_x);
        try stmt.bindFloat(2, box.max_x);
        try stmt.bindFloat(3, box.min_y);
        try stmt.bindFloat(4, box.max_y);
        try stmt.bindInt(5, id);

        _ = try stmt.step();
    }

    /// Deletes an entry from an R-Tree.
    pub fn delete(self: *Self, table_name: []const u8, id: i64) !void {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "DELETE FROM {s} WHERE id=?\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        defer stmt.finalize();

        try stmt.bindInt(1, id);
        _ = try stmt.step();
    }

    // ========================================================================
    // Query Operations
    // ========================================================================

    /// Queries for all entries that intersect with a bounding box.
    /// Returns a prepared statement with results: id, min_x, max_x, min_y, max_y
    pub fn queryIntersects2D(self: *Self, table_name: []const u8, box: BoundingBox2D) !Statement {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "SELECT * FROM {s} WHERE min_x <= ? AND max_x >= ? AND min_y <= ? AND max_y >= ?\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        errdefer stmt.finalize();

        try stmt.bindFloat(1, box.max_x);
        try stmt.bindFloat(2, box.min_x);
        try stmt.bindFloat(3, box.max_y);
        try stmt.bindFloat(4, box.min_y);

        return stmt;
    }

    /// Queries for all entries contained within a bounding box.
    pub fn queryContained2D(self: *Self, table_name: []const u8, box: BoundingBox2D) !Statement {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "SELECT * FROM {s} WHERE min_x >= ? AND max_x <= ? AND min_y >= ? AND max_y <= ?\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        errdefer stmt.finalize();

        try stmt.bindFloat(1, box.min_x);
        try stmt.bindFloat(2, box.max_x);
        try stmt.bindFloat(3, box.min_y);
        try stmt.bindFloat(4, box.max_y);

        return stmt;
    }

    /// Queries for all entries near a point (within a given radius).
    pub fn queryNearPoint2D(self: *Self, table_name: []const u8, x: f64, y: f64, radius: f64) !Statement {
        const box = BoundingBox2D{
            .min_x = x - radius,
            .max_x = x + radius,
            .min_y = y - radius,
            .max_y = y + radius,
        };
        return self.queryIntersects2D(table_name, box);
    }

    /// Queries for all entries near a geographic coordinate.
    /// radius is in degrees (approximately: 1 degree ≈ 111km at equator)
    pub fn queryNearGeo(self: *Self, table_name: []const u8, coord: GeoCoord, radius_degrees: f64) !Statement {
        return self.queryNearPoint2D(table_name, coord.longitude, coord.latitude, radius_degrees);
    }

    /// Gets all entry IDs that intersect with a bounding box.
    pub fn getIntersectingIds2D(self: *Self, table_name: []const u8, box: BoundingBox2D) ![]i64 {
        var stmt = try self.queryIntersects2D(table_name, box);
        defer stmt.finalize();

        var ids: std.ArrayListUnmanaged(i64) = .empty;
        errdefer ids.deinit(self.allocator);

        while (try stmt.step()) {
            try ids.append(self.allocator, stmt.columnInt(0));
        }

        return ids.toOwnedSlice(self.allocator);
    }

    /// Frees an array of IDs returned by getIntersectingIds2D.
    pub fn freeIds(self: *Self, ids: []i64) void {
        self.allocator.free(ids);
    }

    /// Counts entries that intersect with a bounding box.
    pub fn countIntersects2D(self: *Self, table_name: []const u8, box: BoundingBox2D) !i64 {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "SELECT COUNT(*) FROM {s} WHERE min_x <= ? AND max_x >= ? AND min_y <= ? AND max_y >= ?\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        defer stmt.finalize();

        try stmt.bindFloat(1, box.max_x);
        try stmt.bindFloat(2, box.min_x);
        try stmt.bindFloat(3, box.max_y);
        try stmt.bindFloat(4, box.min_y);

        if (try stmt.step()) {
            return stmt.columnInt(0);
        }
        return 0;
    }

    /// Gets the total count of entries in an R-Tree.
    pub fn count(self: *Self, table_name: []const u8) !i64 {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "SELECT COUNT(*) FROM {s}\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        defer stmt.finalize();

        if (try stmt.step()) {
            return stmt.columnInt(0);
        }
        return 0;
    }

    /// Gets the bounding box that encompasses all entries in the R-Tree.
    pub fn getBounds2D(self: *Self, table_name: []const u8) !?BoundingBox2D {
        const sql = try std.fmt.allocPrint(
            self.allocator,
            "SELECT MIN(min_x), MAX(max_x), MIN(min_y), MAX(max_y) FROM {s}\x00",
            .{table_name},
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        defer stmt.finalize();

        if (try stmt.step()) {
            if (stmt.columnIsNull(0)) return null;
            return BoundingBox2D{
                .min_x = stmt.columnFloat(0),
                .max_x = stmt.columnFloat(1),
                .min_y = stmt.columnFloat(2),
                .max_y = stmt.columnFloat(3),
            };
        }
        return null;
    }

    // ========================================================================
    // Join Operations (R-Tree with regular tables)
    // ========================================================================

    /// Performs a spatial join between an R-Tree and a regular table.
    /// Returns entries from both tables where the R-Tree bbox intersects the query box.
    ///
    /// Example:
    /// ```zig
    /// var stmt = try rtree.spatialJoin2D(
    ///     "locations",           // R-Tree table
    ///     "places",              // Regular table
    ///     "id",                  // Join column in places
    ///     &.{ "name", "type" },  // Columns to select from places
    ///     query_box,
    /// );
    /// ```
    pub fn spatialJoin2D(
        self: *Self,
        rtree_table: []const u8,
        data_table: []const u8,
        join_column: []const u8,
        select_columns: []const []const u8,
        box: BoundingBox2D,
    ) !Statement {
        // Build select columns
        var col_buf: [2048]u8 = undefined;
        var col_len: usize = 0;

        @memcpy(col_buf[col_len..][0..4], "r.id");
        col_len += 4;

        for (select_columns) |col| {
            @memcpy(col_buf[col_len..][0..4], ", d.");
            col_len += 4;
            @memcpy(col_buf[col_len..][0..col.len], col);
            col_len += col.len;
        }

        const sql = try std.fmt.allocPrint(
            self.allocator,
            "SELECT {s} FROM {s} r JOIN {s} d ON r.id = d.{s} WHERE r.min_x <= ? AND r.max_x >= ? AND r.min_y <= ? AND r.max_y >= ?\x00",
            .{ col_buf[0..col_len], rtree_table, data_table, join_column },
        );
        defer self.allocator.free(sql);

        var stmt = try self.db.prepare(sql[0 .. sql.len - 1 :0]);
        errdefer stmt.finalize();

        try stmt.bindFloat(1, box.max_x);
        try stmt.bindFloat(2, box.min_x);
        try stmt.bindFloat(3, box.max_y);
        try stmt.bindFloat(4, box.min_y);

        return stmt;
    }
};