Moonshark/modules/sqlite/sqlite.lua

local sqlite = {}

local Connection = {}
Connection.__index = Connection

function Connection:close()
	if self._id then
		local ok = moonshark.sql_close(self._id)
		self._id = nil
		return ok
	end
	return false
end

function Connection:ping()
	if not self._id then
		error("Connection is closed")
	end
	return moonshark.sql_ping(self._id)
end

function Connection:query(query_str, ...)
	if not self._id then
		error("Connection is closed")
	end
	return moonshark.sql_query(self._id, string.normalize_whitespace(query_str), ...)
end

function Connection:exec(query_str, ...)
	if not self._id then
		error("Connection is closed")
	end
	return moonshark.sql_exec(self._id, string.normalize_whitespace(query_str), ...)
end

function Connection:query_row(query_str, ...)
	local results = self:query(query_str, ...)
	return results and #results > 0 and results[1] or nil
end

function Connection:query_value(query_str, ...)
	local row = self:query_row(query_str, ...)
	if row then
		for _, value in pairs(row) do
			return value
		end
	end
	return nil
end

function Connection:begin()
	local result = self:exec("BEGIN")
	if result then
		return {
			conn = self,
			active = true,
			commit = function(tx)
				if tx.active then
					tx.active = false
					return tx.conn:exec("COMMIT")
				end
				return false
			end,
			rollback = function(tx)
				if tx.active then
					tx.active = false
					return tx.conn:exec("ROLLBACK")
				end
				return false
			end,
			query = function(tx, query_str, ...)
				if not tx.active then error("Transaction is not active") end
				return tx.conn:query(query_str, ...)
			end,
			exec = function(tx, query_str, ...)
				if not tx.active then error("Transaction is not active") end
				return tx.conn:exec(query_str, ...)
			end,
			query_row = function(tx, query_str, ...)
				if not tx.active then error("Transaction is not active") end
				return tx.conn:query_row(query_str, ...)
			end,
			query_value = function(tx, query_str, ...)
				if not tx.active then error("Transaction is not active") end
				return tx.conn:query_value(query_str, ...)
			end
		}
	end
	return nil
end

function Connection:insert(table_name, data)
	if string.is_blank(table_name) then
		error("Table name cannot be empty")
	end

	local keys = table.keys(data)
	local values = table.values(data)
	local placeholders = string.trim_right(string.rep("?, ", #keys), ", ")

	local query = string.parse("INSERT INTO {{table}} ({{columns}}) VALUES ({{placeholders}})", {
		table = table_name,
		columns = table.concat(keys, ", "),
		placeholders = placeholders
	})

	return self:exec(query, unpack(values))
end

function Connection:upsert(table_name, data, conflict_columns)
	if string.is_blank(table_name) then
		error("Table name cannot be empty")
	end

	local keys = table.keys(data)
	local values = table.values(data)
	local placeholders = string.trim_right(string.rep("?, ", #keys), ", ")

	local updates = table.map(keys, function(key) return key .. " = excluded." .. key end)
	local updates_str = table.concat(updates, ", ")

	local conflict_clause = ""
	if conflict_columns then
		if type(conflict_columns) == "string" then
			conflict_clause = "(" .. conflict_columns .. ")"
		else
			conflict_clause = "(" .. table.concat(conflict_columns, ", ") .. ")"
		end
	end

	local query = string.parse("INSERT INTO {{table}} ({{columns}}) VALUES ({{placeholders}}) ON CONFLICT {{conflict}} DO UPDATE SET {{updates}}", {
		table = table_name,
		columns = table.concat(keys, ", "),
		placeholders = placeholders,
		conflict = conflict_clause,
		updates = updates_str
	})

	return self:exec(query, unpack(values))
end

function Connection:update(table_name, data, where_clause, ...)
	if string.is_blank(table_name) then
		error("Table name cannot be empty")
	end
	if string.is_blank(where_clause) then
		error("WHERE clause cannot be empty for UPDATE")
	end

	local keys = table.keys(data)
	local values = table.values(data)
	local sets = table.map(keys, function(key) return key .. " = ?" end)
	local sets_str = table.concat(sets, ", ")

	local query = string.parse("UPDATE {{table}} SET {{sets}} WHERE {{where}}", {
		table = table_name,
		sets = sets_str,
		where = where_clause
	})

	table.extend(values, {...})
	return self:exec(query, unpack(values))
end

function Connection:delete(table_name, where_clause, ...)
	if string.is_blank(table_name) then
		error("Table name cannot be empty")
	end
	if string.is_blank(where_clause) then
		error("WHERE clause cannot be empty for DELETE")
	end

	local query = string.parse("DELETE FROM {{table}} WHERE {{where}}", {
		table = table_name,
		where = where_clause
	})
	return self:exec(query, ...)
end

function Connection:select(table_name, columns, where_clause, ...)
	if string.is_blank(table_name) then
		error("Table name cannot be empty")
	end

	columns = columns or "*"
	if type(columns) == "table" then
		columns = table.concat(columns, ", ")
	end

	if where_clause and not string.is_blank(where_clause) then
		local query = string.parse("SELECT {{columns}} FROM {{table}} WHERE {{where}}", {
			columns = columns,
			table = table_name,
			where = where_clause
		})
		return self:query(query, ...)
	else
		local query = string.parse("SELECT {{columns}} FROM {{table}}", {
			columns = columns,
			table = table_name
		})
		return self:query(query)
	end
end

function Connection:table_exists(table_name)
	if string.is_blank(table_name) then return false end
	return self:query_value("SELECT name FROM sqlite_master WHERE type='table' AND name=?", string.trim(table_name)) ~= nil
end

function Connection:column_exists(table_name, column_name)
	if string.is_blank(table_name) or string.is_blank(column_name) then return false end

	local result = self:query(string.parse("PRAGMA table_info({{table}})", {table = table_name}))
	if result then
		return table.any(result, function(row)
			return string.iequals(row.name, string.trim(column_name))
		end)
	end
	return false
end

function Connection:create_table(table_name, schema)
	if string.is_blank(table_name) or string.is_blank(schema) then
		error("Table name and schema cannot be empty")
	end

	local query = string.parse("CREATE TABLE IF NOT EXISTS {{table}} ({{schema}})", {
		table = table_name,
		schema = string.trim(schema)
	})
	return self:exec(query)
end

function Connection:drop_table(table_name)
	if string.is_blank(table_name) then
		error("Table name cannot be empty")
	end
	return self:exec(string.parse("DROP TABLE IF EXISTS {{table}}", {table = table_name}))
end

function Connection:add_column(table_name, column_def)
	if string.is_blank(table_name) or string.is_blank(column_def) then
		error("Table name and column definition cannot be empty")
	end

	local query = string.parse("ALTER TABLE {{table}} ADD COLUMN {{column}}", {
		table = table_name,
		column = string.trim(column_def)
	})
	return self:exec(query)
end

function Connection:create_index(index_name, table_name, columns, unique)
	if string.is_blank(index_name) or string.is_blank(table_name) then
		error("Index name and table name cannot be empty")
	end

	local unique_clause = unique and "UNIQUE " or ""
	local columns_str = type(columns) == "table" and table.concat(columns, ", ") or tostring(columns)

	local query = string.parse("CREATE {{unique}}INDEX IF NOT EXISTS {{index}} ON {{table}} ({{columns}})", {
		unique = unique_clause,
		index = index_name,
		table = table_name,
		columns = columns_str
	})
	return self:exec(query)
end

function Connection:drop_index(index_name)
	if string.is_blank(index_name) then
		error("Index name cannot be empty")
	end
	return self:exec(string.parse("DROP INDEX IF EXISTS {{index}}", {index = index_name}))
end

-- SQLite-specific functions
function Connection:vacuum()
	return self:exec("VACUUM")
end

function Connection:analyze()
	return self:exec("ANALYZE")
end

function Connection:integrity_check()
	return self:query("PRAGMA integrity_check")
end

function Connection:foreign_keys(enabled)
	local value = enabled and "ON" or "OFF"
	return self:exec(string.parse("PRAGMA foreign_keys = {{value}}", {value = value}))
end

function Connection:journal_mode(mode)
	mode = string.upper(mode or "WAL")
	local valid_modes = {"DELETE", "TRUNCATE", "PERSIST", "MEMORY", "WAL", "OFF"}

	if not table.contains(valid_modes, mode) then
		error("Invalid journal mode: " .. mode)
	end

	return self:query(string.parse("PRAGMA journal_mode = {{mode}}", {mode = mode}))
end

function Connection:synchronous(level)
	level = string.upper(level or "NORMAL")
	local valid_levels = {"OFF", "NORMAL", "FULL", "EXTRA"}

	if not table.contains(valid_levels, level) then
		error("Invalid synchronous level: " .. level)
	end

	return self:exec(string.parse("PRAGMA synchronous = {{level}}", {level = level}))
end

function Connection:cache_size(size)
	size = size or -64000
	if type(size) ~= "number" then
		error("Cache size must be a number")
	end
	return self:exec(string.parse("PRAGMA cache_size = {{size}}", {size = tostring(size)}))
end

function Connection:temp_store(mode)
	mode = string.upper(mode or "MEMORY")
	local valid_modes = {"DEFAULT", "FILE", "MEMORY"}

	if not table.contains(valid_modes, mode) then
		error("Invalid temp_store mode: " .. mode)
	end

	return self:exec(string.parse("PRAGMA temp_store = {{mode}}", {mode = mode}))
end

-- Connection management
function sqlite.open(database_path)
	database_path = database_path or ":memory:"
	if database_path ~= ":memory:" and string.is_blank(database_path) then
		database_path = ":memory:"
	end

	local conn_id = moonshark.sql_connect("sqlite", string.trim(database_path))
	if conn_id then
		return setmetatable({_id = conn_id}, Connection)
	end
	return nil
end

sqlite.connect = sqlite.open

-- Quick execution functions
function sqlite.query(database_path, query_str, ...)
	local conn = sqlite.open(database_path)
	if not conn then
		error(string.parse("Failed to open SQLite database: {{path}}", {path = database_path or ":memory:"}))
	end

	local results = conn:query(query_str, ...)
	conn:close()
	return results
end

function sqlite.exec(database_path, query_str, ...)
	local conn = sqlite.open(database_path)
	if not conn then
		error(string.parse("Failed to open SQLite database: {{path}}", {path = database_path or ":memory:"}))
	end

	local result = conn:exec(query_str, ...)
	conn:close()
	return result
end

function sqlite.query_row(database_path, query_str, ...)
	local results = sqlite.query(database_path, query_str, ...)
	return results and #results > 0 and results[1] or nil
end

function sqlite.query_value(database_path, query_str, ...)
	local row = sqlite.query_row(database_path, query_str, ...)
	if row then
		for _, value in pairs(row) do
			return value
		end
	end
	return nil
end

-- Migration helpers
function sqlite.migrate(database_path, migrations)
	local conn = sqlite.open(database_path)
	if not conn then
		error("Failed to open SQLite database for migration")
	end

	conn:create_table("_migrations", "id INTEGER PRIMARY KEY, name TEXT UNIQUE, applied_at DATETIME DEFAULT CURRENT_TIMESTAMP")

	local tx = conn:begin()
	if not tx then
		conn:close()
		error("Failed to begin migration transaction")
	end

	for _, migration in ipairs(migrations) do
		if not migration.name or string.is_blank(migration.name) then
			tx:rollback()
			conn:close()
			error("Migration must have a non-empty name")
		end

		local existing = conn:query_value("SELECT id FROM _migrations WHERE name = ?", string.trim(migration.name))
		if not existing then
			local ok, err = pcall(function()
				if type(migration.up) == "string" then
					conn:exec(migration.up)
				elseif type(migration.up) == "function" then
					migration.up(conn)
				else
					error("Migration 'up' must be string or function")
				end
			end)

			if ok then
				conn:exec("INSERT INTO _migrations (name) VALUES (?)", string.trim(migration.name))
				print(string.parse("Applied migration: {{name}}", {name = migration.name}))
			else
				tx:rollback()
				conn:close()
				error(string.parse("Migration '{{name}}' failed: {{error}}", {
					name = migration.name,
					error = err or "unknown error"
				}))
			end
		end
	end

	tx:commit()
	conn:close()
	return true
end

-- Result processing utilities
function sqlite.to_array(results, column_name)
	if not results or table.is_empty(results) then return {} end
	if string.is_blank(column_name) then error("Column name cannot be empty") end
	return table.map(results, function(row) return row[column_name] end)
end

function sqlite.to_map(results, key_column, value_column)
	if not results or table.is_empty(results) then return {} end
	if string.is_blank(key_column) then error("Key column name cannot be empty") end

	local map = {}
	for _, row in ipairs(results) do
		local key = row[key_column]
		map[key] = value_column and row[value_column] or row
	end
	return map
end

function sqlite.group_by(results, column_name)
	if not results or table.is_empty(results) then return {} end
	if string.is_blank(column_name) then error("Column name cannot be empty") end
	return table.group_by(results, function(row) return row[column_name] end)
end

function sqlite.print_results(results)
	if not results or table.is_empty(results) then
		print("No results")
		return
	end

	local columns = table.keys(results[1])
	table.sort(columns)

	-- Calculate column widths
	local widths = {}
	for _, col in ipairs(columns) do
		widths[col] = string.length(col)
		for _, row in ipairs(results) do
			local value = tostring(row[col] or "")
			widths[col] = math.max(widths[col], string.length(value))
		end
	end

	-- Print header and separator
	local header_parts = table.map(columns, function(col) return string.pad_right(col, widths[col]) end)
	local separator_parts = table.map(columns, function(col) return string.rep("-", widths[col]) end)

	print(table.concat(header_parts, " | "))
	print(table.concat(separator_parts, "-+-"))

	-- Print rows
	for _, row in ipairs(results) do
		local value_parts = table.map(columns, function(col)
			local value = tostring(row[col] or "")
			return string.pad_right(value, widths[col])
		end)
		print(table.concat(value_parts, " | "))
	end
end

return sqlite