path: root/base/src/main/java/bjc/utils/misc/NamedPreparedStatement.java

package bjc.utils.misc;

import java.sql.*;
import java.util.*;
import java.util.function.Function;

/**
 * NamedPreparedStatement
 *
 * Features
 * - :name  -> single parameter (bind once, all repeats covered)
 * - @name  -> list parameter (expands to ?,?,? based on values you set; repeats ok)
 * - Empty-list toggle:
 *      .emptyListAsNull()                  // IN (@ids) -> IN (NULL)                  [default]
 *      .emptyListAsLiteral("YOUR_TEXT")    // IN (@ids) -> IN (YOUR_TEXT)
 * - Skips strings ('...'/\"...\") and comments (-- ..., /* ... * /)
 * - Respects PostgreSQL type casts (::type)
 * - Lazy compilation (or call compile() to do it eagerly)
 *
 * Usage:
 *   String sql = "SELECT * FROM t WHERE a = :a AND b IN (@bs)";
 *   try (NamedPreparedStatement nps = NamedPreparedStatement.prepare(conn, sql)) {
 *       nps.setInt("a", 7);
 *       nps.setList("bs", 1, 2, 3);
 *       // optional toggles
 *       nps.emptyListAsLiteral("SELECT NULL WHERE 1=0");
 *       // eager build if you need the PS now:
 *       PreparedStatement ps = nps.compile(); // built & bound, not executed
 *       try (ResultSet rs = ps.executeQuery()) { ... }
 *   }
 */
public class NamedPreparedStatement extends BoundPreparedStatement implements AutoCloseable {
	/* ---------- Construction ---------- */

	private final Connection conn;
	private final List<Token> tokens; // parsed tokens from original SQL
	private final String originalSql;

	private final Integer autoGeneratedKeys; // null if not requested
	private final Integer resultSetType; // null unless custom rs type
	private final Integer resultSetConcurrency; // null unless custom rs type

	// current compiled statement
	private PreparedStatement currentPs;
	private String lastCompiledSql;

	// empty-list behavior
	private EmptyListMode emptyListMode = EmptyListMode.AS_NULL;
	private String emptyListLiteral = "NULL";

	private NamedPreparedStatement(Connection conn, String sql, Integer autoGeneratedKeys, Integer resultSetType,
			Integer resultSetConcurrency) {
		super();
		this.conn = Objects.requireNonNull(conn, "conn");
		this.originalSql = Objects.requireNonNull(sql, "sql");
		this.tokens = parse(sql);
		this.autoGeneratedKeys = autoGeneratedKeys;
		this.resultSetType = resultSetType;
		this.resultSetConcurrency = resultSetConcurrency;
	}

	/**
	 * Default options. 
	 * 
	 * @param conn The DB connection to use
	 * @param sql The SQL query to prepare
	 * @return The prepared statement, bound to the given DB connection
	 */
	public static NamedPreparedStatement prepare(Connection conn, String sql) {
		return new NamedPreparedStatement(conn, sql, null, null, null);
	}

	/**
	 * Request generated keys. 
	 * 
	 * @param conn The DB connection to use
	 * @param sql The SQL query to prepare
	 * @return The prepared statement, bound to the given DB connection and set to return generated keys.
	 */
	public static NamedPreparedStatement prepareWithKeys(Connection conn, String sql) {
		return new NamedPreparedStatement(conn, sql, Statement.RETURN_GENERATED_KEYS, null, null);
	}

	/**
	 * Custom result set type/concurrency. 
	 * 
	 * @param conn The DB connection to use
	 * @param sql The SQL query to prepare
	 * @param resultSetType The result-set type to use
	 * @param resultSetConcurrency The result-set concurrency to use
	 * @return The prepared statement, configured as specified
	 */
	public static NamedPreparedStatement prepare(Connection conn, String sql, int resultSetType,
			int resultSetConcurrency) {
		return new NamedPreparedStatement(conn, sql, null, resultSetType, resultSetConcurrency);
	}

	/* ---------- Empty-list toggle ---------- */

	/**
	 * Expand empty @lists as NULL (IN (@x) -> IN (NULL)).
	 *  
	 * @return The prepared statement
	 */
	public NamedPreparedStatement emptyListAsNull() {
		this.emptyListMode = EmptyListMode.AS_NULL;
		return this;
	}

	/**
	 * Expand empty @lists as the provided literal (inserted verbatim). Example
	 * (Postgres): emptyListAsLiteral("SELECT NULL WHERE 1=0") → IN (@x) becomes IN
	 * (SELECT NULL WHERE 1=0)
	 * 
	 * @param literal The literal to expand empty lists to
	 * @return The prepared statement
	 */
	public NamedPreparedStatement emptyListAsLiteral(String literal) {
		if (literal == null || literal.isBlank())
			throw new IllegalArgumentException("Empty-list literal must be non-blank.");
		this.emptyListMode = EmptyListMode.AS_CUSTOM_LITERAL;
		this.emptyListLiteral = literal;
		return this;
	}

	/* ---------- Eager compile ---------- */

	/**
	 * Build final SQL (expanding @lists), create PreparedStatement, and bind all
	 * parameters; does NOT execute. Returns the prepared statement instance.
	 * 
	 * @return The fully bound prepared statement
	 * @throws SQLException If something went wrong preparing the statement
	 */
	public PreparedStatement compile() throws SQLException {
		buildAndBind();
		return currentPs;
	}
	
	/* ---------- Execution ---------- */

	/**
	 * Execute this statement
	 * @return The result of executing the query
	 * @throws SQLException If something went wrong executing the query
	 */
	public ResultSet executeQuery() throws SQLException {
		buildAndBind();
		return currentPs.executeQuery();
	}

	/**
	 * Execute this update.
	 * 
	 * @return The number of rows updated by the update
	 * @throws SQLException If something went wrong executing the update
	 */
	public int executeUpdate() throws SQLException {
		buildAndBind();
		return currentPs.executeUpdate();
	}

	/**
	 * Execute the SQL statement
	 * 
	 * @return The result of executing the SQL statement
	 * @throws SQLException If something went wrong executing the statement
	 */
	public boolean execute() throws SQLException {
		buildAndBind();
		return currentPs.execute();
	}

	/**
	 * Execute one or more batch of parameters against this SQL statement
	 * 
	 * @return The number of rows updated by each batch
	 * @throws SQLException If something went wrong executing the statement
	 */
	public int[] executeBatch() throws SQLException {
		if (batch.isEmpty()) {
			buildAndBind();
			return currentPs.executeBatch();
		}
		List<Integer> counts = new ArrayList<>(batch.size());
		List<Binder> binders = buildPreparedStatement();

		for (ParamSnapshot snap : batch) {
			// restore snapshot
			this.singles.clear();
			this.lists.clear();
			this.singles.putAll(snap.singles);
			this.lists.putAll(snap.lists);

			bindPreparedStatement(binders);

			counts.add(currentPs.executeUpdate());
		}
		batch.clear();
		return counts.stream().mapToInt(Integer::intValue).toArray();
	}

	/**
	 * Execute a batch from a recorded snapshot
	 * @param record
	 * @return The number of rows updated by each batch
	 * @throws SQLException
	 */
	public List<Integer> executeBatchFromRecord(BoundPreparedStatement record) throws SQLException {
		if (record.batch.isEmpty()) {
			buildAndBind();
			int[] batchRes = currentPs.executeBatch();
			List<Integer> resList = new ArrayList<>(batchRes.length);
			for (int i : batchRes) resList.add(i);
			return resList;
		}
		
		List<Integer> counts = new ArrayList<>(batch.size());
		List<Binder> binders = buildPreparedStatement(record);

		for (ParamSnapshot snap : record.batch) {
			// restore snapshot
			this.singles.clear();
			this.lists.clear();
			this.singles.putAll(snap.singles);
			this.lists.putAll(snap.lists);

			bindPreparedStatement(binders);

			counts.add(currentPs.executeUpdate());
		}
		record.batch.clear();
		return counts;
	}
	
	/**
	 * Get the generated keys for this prepared statement.
	 * 
	 * See {@link PreparedStatement#getGeneratedKeys()}
	 * 
	 * @return The generated keys for this prepared statement
	 * @throws SQLException If something went wrong getting the generated keys
	 */
	public ResultSet getGeneratedKeys() throws SQLException {
		if (currentPs == null)
			throw new IllegalStateException("No statement compiled/executed yet.");
		return currentPs.getGeneratedKeys();
	}

	/**
	 * Get the underlying prepared statement, or null if one hasn't been prepared yet
	 * 
	 * TODO should this return a Optional?
	 * 
	 * @return The underlying prepared statement, or null if there isn't one
	 */
	public PreparedStatement getPreparedStatement() {
		return currentPs;
	}

	/**
	 * Get the parsed SQL for this statement.
	 * 
	 * @return The parsed SQL for this statement
	 */
	public String getParsedSql() {
		return lastCompiledSql;
	}

	/**
	 * Get the original SQL for this statement.
	 * 
	 * @return The original SQL for this statement
	 */
	public String getOriginalSql() {
		return originalSql;
	}

	/**
	 * Get the available parameter names.
	 * 
	 * @return The names of the available parameters
	 */
	public Set<String> getParameterNames() {
		LinkedHashSet<String> names = new LinkedHashSet<>();
		for (Token t : tokens)
			if (t instanceof ParamToken pt)
				names.add(pt.name);
		return Collections.unmodifiableSet(names);
	}

	@Override
	public void close() throws SQLException {
		if (currentPs != null)
			currentPs.close();
	}

	@Override
	public String toString() {
		return "NamedPreparedStatement{sql='" + (lastCompiledSql != null ? lastCompiledSql : originalSql) + "'}";
	}

	/* ---------- Binding (single) ---------- */

	

	/* ---------- Binding (lists) ---------- */

	

	/* ---------- Build & bind ---------- */

	private void buildAndBind() throws SQLException {
		List<Binder> bindersInOrder = buildPreparedStatement();

		bindPreparedStatement(bindersInOrder);
	}

	private void bindPreparedStatement(List<Binder> bindersInOrder) throws SQLException {
		int idx = 1;
		for (Binder b : bindersInOrder)
			b.bind(currentPs, idx++);
	}

	private List<Binder> buildPreparedStatement() throws SQLException {
		return buildPreparedStatement(this);
	}
	
	private List<Binder> buildPreparedStatement(BoundPreparedStatement record) throws SQLException {
		StringBuilder sb = new StringBuilder(originalSql.length() + 32);
		List<Binder> bindersInOrder = new ArrayList<>(32);

		for (Token t : tokens) {
			if (t instanceof TextToken tt) {
				sb.append(tt.text);
			} else {
				ParamToken pt = (ParamToken) t;
				if (!pt.list) { // :name
					Binder b = singles.get(pt.name);
					if (b == null)
						throw new IllegalArgumentException("Missing value for parameter :" + pt.name);
					sb.append('?');
					bindersInOrder.add(b);
				} else { // @name
					List<Binder> lst = lists.get(pt.name);
					if (lst == null)
						throw new IllegalArgumentException("Missing list for parameter @" + pt.name);
					if (lst.isEmpty()) {
						if (emptyListMode == EmptyListMode.AS_NULL) {
							sb.append("NULL");
						} else {
							sb.append(emptyListLiteral);
						}
					} else {
						for (int i = 0; i < lst.size(); i++) {
							if (i > 0)
								sb.append(',');
							sb.append('?');
							bindersInOrder.add(lst.get(i));
						}
					}
				}
			}
		}

		this.lastCompiledSql = sb.toString();

		if (currentPs != null)
			try {
				currentPs.close();
			} catch (SQLException ignore) {
			}

		if (autoGeneratedKeys != null) {
			currentPs = conn.prepareStatement(lastCompiledSql, autoGeneratedKeys);
		} else if (resultSetType != null && resultSetConcurrency != null) {
			currentPs = conn.prepareStatement(lastCompiledSql, resultSetType, resultSetConcurrency);
		} else {
			currentPs = conn.prepareStatement(lastCompiledSql);
		}
		return bindersInOrder;
	}
	
	/* ---------- Parser ---------- */

	private interface Token {
	}

	private static final class TextToken implements Token {
		final String text;

		TextToken(String s) {
			this.text = s;
		}
	}

	private static final class ParamToken implements Token {
		final String name;
		final boolean list;

		ParamToken(String n, boolean l) {
			this.name = n;
			this.list = l;
		}
	}

	private static List<Token> parse(String sql) {
		ArrayList<Token> out = new ArrayList<>();
		StringBuilder buf = new StringBuilder();

		boolean inSingle = false, inDouble = false, inLine = false, inBlock = false;
		int n = sql.length();

		for (int i = 0; i < n; i++) {
			char c = sql.charAt(i);

			if (inLine) {
				buf.append(c);
				if (c == '\n' || c == '\r')
					inLine = false;
				continue;
			}
			if (inBlock) {
				buf.append(c);
				if (c == '*' && i + 1 < n && sql.charAt(i + 1) == '/') {
					buf.append('/');
					i++;
					inBlock = false;
				}
				continue;
			}
			if (inSingle) {
				buf.append(c);
				if (c == '\'') {
					if (i + 1 < n && sql.charAt(i + 1) == '\'') {
						buf.append('\'');
						i++;
					} else
						inSingle = false;
				}
				continue;
			}
			if (inDouble) {
				buf.append(c);
				if (c == '"') {
					if (i + 1 < n && sql.charAt(i + 1) == '"') {
						buf.append('"');
						i++;
					} else
						inDouble = false;
				}
				continue;
			}

			if (c == '-' && i + 1 < n && sql.charAt(i + 1) == '-') {
				buf.append("--");
				i++;
				inLine = true;
				continue;
			}
			if (c == '/' && i + 1 < n && sql.charAt(i + 1) == '*') {
				buf.append("/*");
				i++;
				inBlock = true;
				continue;
			}
			if (c == '\'') {
				buf.append(c);
				inSingle = true;
				continue;
			}
			if (c == '"') {
				buf.append(c);
				inDouble = true;
				continue;
			}

			if (c == ':' || c == '@') {
				char prev = (i > 0) ? sql.charAt(i - 1) : '\0';
				if ((c == ':' && prev == ':') || (c == '@' && prev == '@')) {
					buf.append(c);
					continue;
				} // ::type or @@sys
				if (i + 1 < n && isIdentStart(sql.charAt(i + 1))) {
					if (buf.length() > 0) {
						out.add(new TextToken(buf.toString()));
						buf.setLength(0);
					}
					int j = i + 2;
					while (j < n && isIdentPart(sql.charAt(j)))
						j++;
					String name = sql.substring(i + 1, j);
					out.add(new ParamToken(name, c == '@'));
					i = j - 1;
					continue;
				}
			}

			buf.append(c);
		}
		if (buf.length() > 0)
			out.add(new TextToken(buf.toString()));
		return out;
	}

	private static boolean isIdentStart(char c) {
		return c == '_' || Character.isLetter(c);
	}

	private static boolean isIdentPart(char c) {
		return c == '_' || Character.isLetterOrDigit(c);
	}

	/* ---------- Helpers ---------- */

	private enum EmptyListMode {
		AS_NULL, AS_CUSTOM_LITERAL
	}
}