<?php
/**
 * Checks and runs a model-written, read-only SQL query.
 *
 * @package Askshelf
 */

declare( strict_types = 1 );

namespace Askshelf\Data;

use InvalidArgumentException;

defined( 'ABSPATH' ) || exit;

// Exception messages here go back to the model as JSON tool results, never into HTML.
// phpcs:disable WordPress.Security.EscapeOutput.ExceptionNotEscaped

/**
 * The custom-query tool's safety layer. A query must pass every check here, then runs
 * in a READ ONLY transaction with a time limit, and its results are scrubbed:
 *
 * - One statement, starting with SELECT or WITH. No comments, no user variables (@),
 *   no statement separators except one trailing semicolon. mysqli runs one statement
 *   per call, so nothing can be stacked after it. Strings use single quotes with no
 *   backslashes, so they read the same whatever the server's SQL mode.
 * - No keyword that writes, locks, sleeps, reads files or changes the session.
 * - Tables: only the allowlist below (sales, products, stock, coupons, categories),
 *   written without the WordPress table prefix; this class adds it. Any other real table
 *   named anywhere in the query is refused, so a comma join can't reach users or
 *   postmeta. No schema-qualified names (information_schema.x, mysql.x).
 * - Columns that identify people (emails, names, addresses, phones, IPs, notes, order
 *   keys) are refused wherever they appear, and SELECT * is refused, so they can't be
 *   read from an allowed table either.
 * - At most 200 rows come back; longer strings are cut, and anything that looks like an
 *   email address is replaced.
 */
final class SqlGuard {

	public const MAX_ROWS = 200;

	private const MAX_SQL_CHARS = 4000;

	private const MAX_LIMIT = 1000;

	private const TIME_LIMIT_MS = 5000;

	/** Tables the model may read, without the WordPress prefix. */
	public const TABLES = array(
		'wc_order_stats',
		'wc_order_product_lookup',
		'wc_order_coupon_lookup',
		'wc_order_tax_lookup',
		'wc_product_meta_lookup',
		'wc_category_lookup',
		'wc_customer_lookup',
		'wc_orders',
		'woocommerce_order_items',
		'posts',
		'terms',
		'term_taxonomy',
		'term_relationships',
	);

	/** Column names that identify a person, refused anywhere in a query. */
	private const BLOCKED_COLUMNS = array(
		'email',
		'user_email',
		'billing_email',
		'first_name',
		'last_name',
		'username',
		'user_login',
		'user_pass',
		'user_nicename',
		'display_name',
		'phone',
		'billing_phone',
		'address_1',
		'address_2',
		'postcode',
		'city',
		'company',
		'customer_note',
		'post_excerpt',
		'post_password',
		'post_author',
		'order_key',
		'ip_address',
		'customer_ip_address',
		'user_agent',
		'customer_user_agent',
		'transaction_id',
		'meta_value',
		'comment_content',
		'guid',
	);

	/** Keywords and functions with no place in a read-only report query. */
	private const BLOCKED_WORDS = array(
		'update',
		'delete',
		'create',
		'alter',
		'drop',
		'truncate',
		'rename',
		'grant',
		'revoke',
		'lock',
		'unlock',
		'call',
		'do',
		'handler',
		'load',
		'set',
		'into',
		'outfile',
		'dumpfile',
		'execute',
		'prepare',
		'deallocate',
		'sleep',
		'benchmark',
		'get_lock',
		'release_lock',
		'is_free_lock',
		'is_used_lock',
		'load_file',
		'shutdown',
		'kill',
		'flush',
		'reset',
		'purge',
		'install',
		'uninstall',
		'analyze',
		'optimize',
		'repair',
		'checksum',
		'show',
		'describe',
		'explain',
		'use',
		'procedure',
		'sys_exec',
		'sys_eval',
		'database',
		'schema',
		'user',
		'current_user',
		'session_user',
		'system_user',
		'version',
		'connection_id',
		'master_pos_wait',
		'source_pos_wait',
		'found_rows',
		'last_insert_id',
		'row_count',
		'information_schema',
		'performance_schema',
		'mysql',
		'sys',
		'share',
	);

	/**
	 * Checks a query and rewrites its table names. Throws with a message the model can
	 * act on when the query isn't allowed.
	 *
	 * @param string $sql Query as the model wrote it.
	 * @return string The query to run, with prefixed table names and a LIMIT.
	 * @throws InvalidArgumentException When the query isn't allowed.
	 */
	public static function prepare( string $sql ): string {
		global $wpdb;
		$sql = trim( $sql );
		if ( '' === $sql || strlen( $sql ) > self::MAX_SQL_CHARS ) {
			throw new InvalidArgumentException( 'The query is empty or longer than 4000 characters.' );
		}

		$tokens = self::tokenize( $sql );
		if ( ! empty( $tokens ) && ';' === end( $tokens )['v'] ) {
			array_pop( $tokens );
		}
		if ( empty( $tokens ) ) {
			throw new InvalidArgumentException( 'The query is empty.' );
		}

		$first = strtolower( $tokens[0]['v'] );
		if ( 'word' !== $tokens[0]['t'] || ! in_array( $first, array( 'select', 'with' ), true ) ) {
			throw new InvalidArgumentException( 'Only SELECT queries (optionally starting with WITH) are allowed.' );
		}

		$real_tables = self::real_tables();
		$prefix      = $wpdb->prefix;
		$ctes        = self::cte_names( $tokens );
		$count       = count( $tokens );
		$depth       = 0;
		$has_limit   = false;
		// Per parenthesis depth: whether we are in a FROM list, where a comma starts a table.
		$in_from = array( 0 => false );
		$ends    = array( 'where', 'group', 'having', 'order', 'limit', 'union', 'window', 'into', 'for', 'lock', 'select' );

		for ( $i = 0; $i < $count; $i++ ) {
			$tok  = $tokens[ $i ];
			$word = strtolower( $tok['v'] );

			if ( 'sym' === $tok['t'] ) {
				if ( ';' === $tok['v'] ) {
					throw new InvalidArgumentException( 'Only one statement is allowed.' );
				}
				if ( '@' === $tok['v'] ) {
					throw new InvalidArgumentException( 'Variables (@) are not allowed.' );
				}
				if ( '(' === $tok['v'] ) {
					++$depth;
					$in_from[ $depth ] = false;
				} elseif ( ')' === $tok['v'] ) {
					unset( $in_from[ $depth ] );
					--$depth;
				} elseif ( '*' === $tok['v'] && self::is_select_star( $tokens, $i ) ) {
					throw new InvalidArgumentException( 'SELECT * is not allowed: name the columns you need.' );
				}
				continue;
			}
			if ( 'word' !== $tok['t'] && 'ident' !== $tok['t'] ) {
				continue;
			}

			if ( 'word' === $tok['t'] ) {
				if ( 'from' === $word ) {
					$in_from[ $depth ] = true;
				} elseif ( in_array( $word, $ends, true ) ) {
					$in_from[ $depth ] = false;
				}
			}

			if ( 'word' === $tok['t'] && in_array( $word, self::BLOCKED_WORDS, true ) ) {
				throw new InvalidArgumentException( sprintf( '"%s" is not allowed in a read-only report query.', strtoupper( $word ) ) );
			}
			if ( in_array( $word, self::BLOCKED_COLUMNS, true ) ) {
				throw new InvalidArgumentException( sprintf( 'The column "%s" can identify a person, so Askshelf never reads it.', $word ) );
			}
			if ( 'for' === $word && isset( $tokens[ $i + 1 ] ) && 'update' === strtolower( $tokens[ $i + 1 ]['v'] ) ) {
				throw new InvalidArgumentException( 'Locking reads are not allowed.' );
			}
			if ( 0 === $depth && 'limit' === $word ) {
				$has_limit = true;
				self::check_limit( $tokens, $i );
			}

			// A real table named anywhere must be on the allowlist.
			$unprefixed = self::unprefix( $word, $prefix );
			if ( isset( $real_tables[ $prefix . $unprefixed ] ) && ! in_array( $unprefixed, self::TABLES, true ) && ! isset( $ctes[ $word ] ) ) {
				throw new InvalidArgumentException( sprintf( 'The table "%s" is not one Askshelf may read. Allowed: %s.', $unprefixed, implode( ', ', self::TABLES ) ) );
			}

			// Table positions: after FROM, JOIN, or a comma in a FROM list.
			$prev     = $i > 0 ? strtolower( $tokens[ $i - 1 ]['v'] ) : '';
			$is_table = in_array( $prev, array( 'from', 'join', 'straight_join' ), true )
				|| ( ',' === $prev && ! empty( $in_from[ $depth ] ) );
			if ( $is_table ) {
				if ( isset( $tokens[ $i + 1 ] ) && '.' === $tokens[ $i + 1 ]['v'] ) {
					throw new InvalidArgumentException( 'Schema-qualified table names are not allowed. Write tables without a database name or prefix, e.g. wc_order_stats.' );
				}
				if ( isset( $ctes[ $word ] ) ) {
					continue;
				}
				if ( ! in_array( $unprefixed, self::TABLES, true ) ) {
					throw new InvalidArgumentException( sprintf( 'The table "%s" is not one Askshelf may read. Allowed: %s.', $tok['v'], implode( ', ', self::TABLES ) ) );
				}
				$tokens[ $i ] = array(
					't'   => 'ident',
					'v'   => '`' . $prefix . $unprefixed . '`',
					'raw' => true,
				);
			}
		}//end for

		if ( 0 !== $depth ) {
			throw new InvalidArgumentException( 'The query has unbalanced parentheses.' );
		}

		$out = self::join_tokens( $tokens );
		if ( ! $has_limit ) {
			$out .= ' LIMIT ' . ( self::MAX_ROWS + 1 );
		}
		return $out;
	}

	/**
	 * Runs a prepared query read-only, with a time limit.
	 *
	 * @param string $sql Output of prepare().
	 * @return array{columns: list<string>, rows: list<array<string, mixed>>, truncated: bool}
	 * @throws InvalidArgumentException When the database rejects the query.
	 */
	public static function run( string $sql ): array {
		global $wpdb;
		$mariadb    = false !== stripos( (string) $wpdb->db_server_info(), 'mariadb' );
		$suppressed = $wpdb->suppress_errors( true );

		// phpcs:disable WordPress.DB.DirectDatabaseQuery.DirectQuery, WordPress.DB.DirectDatabaseQuery.NoCaching, WordPress.DB.PreparedSQL.NotPrepared -- Checked by prepare(); runs read-only.
		$wpdb->query( $mariadb ? 'SET SESSION max_statement_time = ' . ( self::TIME_LIMIT_MS / 1000 ) : 'SET SESSION MAX_EXECUTION_TIME = ' . self::TIME_LIMIT_MS );
		$wpdb->query( 'START TRANSACTION READ ONLY' );
		$rows  = $wpdb->get_results( $sql, ARRAY_A );
		$error = $wpdb->last_error;
		$wpdb->query( 'ROLLBACK' );
		$wpdb->query( $mariadb ? 'SET SESSION max_statement_time = 0' : 'SET SESSION MAX_EXECUTION_TIME = 0' );
		// phpcs:enable WordPress.DB.DirectDatabaseQuery.DirectQuery, WordPress.DB.DirectDatabaseQuery.NoCaching, WordPress.DB.PreparedSQL.NotPrepared
		$wpdb->suppress_errors( $suppressed );

		if ( '' !== $error || ! is_array( $rows ) ) {
			// The database's own message helps the model fix its query; table names keep their prefix out.
			throw new InvalidArgumentException( 'The database rejected the query: ' . str_replace( $wpdb->prefix, '', $error ) );
		}

		$truncated = count( $rows ) > self::MAX_ROWS;
		$rows      = array_slice( $rows, 0, self::MAX_ROWS );
		$rows      = array_values( array_map( array( self::class, 'scrub_row' ), $rows ) );
		return array(
			'columns'   => $rows ? array_keys( $rows[0] ) : array(),
			'rows'      => $rows,
			'truncated' => $truncated,
		);
	}

	/**
	 * Cuts long strings and hides anything that looks like an email address.
	 *
	 * @param array<string, mixed> $row Row.
	 * @return array<string, mixed>
	 */
	private static function scrub_row( array $row ): array {
		foreach ( $row as $key => $value ) {
			if ( ! is_string( $value ) ) {
				continue;
			}
			$value       = (string) preg_replace( '/[^\s@<>"\']+@[^\s@<>"\']+\.[a-z]{2,}/i', '[hidden]', $value );
			$row[ $key ] = mb_strlen( $value ) > 200 ? mb_substr( $value, 0, 200 ) . '…' : $value;
			if ( is_numeric( $row[ $key ] ) && ! preg_match( '/^0\d/', $row[ $key ] ) ) {
				$row[ $key ] = $row[ $key ] + 0;
			}
		}
		return $row;
	}

	/**
	 * Splits SQL into tokens: word, ident (backticked), str, num, sym. Refuses comments.
	 *
	 * @param string $sql SQL.
	 * @return list<array{t: string, v: string}>
	 * @throws InvalidArgumentException On comments or unterminated strings.
	 */
	public static function tokenize( string $sql ): array {
		$tokens = array();
		$len    = strlen( $sql );
		$i      = 0;
		while ( $i < $len ) {
			$c = $sql[ $i ];
			if ( ctype_space( $c ) ) {
				++$i;
				continue;
			}
			$two = substr( $sql, $i, 2 );
			if ( '--' === $two || '/*' === $two || '#' === $c ) {
				throw new InvalidArgumentException( 'Comments are not allowed in the query.' );
			}
			if ( '"' === $c ) {
				throw new InvalidArgumentException( 'Use single quotes for strings and backticks for names.' );
			}
			if ( "'" === $c ) {
				$j = $i + 1;
				while ( $j < $len ) {
					if ( '\\' === $sql[ $j ] ) {
						throw new InvalidArgumentException( 'Backslashes are not allowed in strings.' );
					}
					if ( $sql[ $j ] === $c ) {
						if ( $j + 1 < $len && $sql[ $j + 1 ] === $c ) {
							$j += 2;
							continue;
						}
						break;
					}
					++$j;
				}
				if ( $j >= $len ) {
					throw new InvalidArgumentException( 'The query has an unterminated string.' );
				}
				$tokens[] = array(
					't' => 'str',
					'v' => substr( $sql, $i, $j - $i + 1 ),
				);
				$i        = $j + 1;
				continue;
			}//end if
			if ( '`' === $c ) {
				$j = strpos( $sql, '`', $i + 1 );
				if ( false === $j ) {
					throw new InvalidArgumentException( 'The query has an unterminated `identifier`.' );
				}
				$tokens[] = array(
					't' => 'ident',
					'v' => substr( $sql, $i + 1, $j - $i - 1 ),
				);
				$i        = $j + 1;
				continue;
			}
			if ( preg_match( '/\G[A-Za-z_][A-Za-z0-9_$]*/', $sql, $m, 0, $i ) ) {
				$tokens[] = array(
					't' => 'word',
					'v' => $m[0],
				);
				$i       += strlen( $m[0] );
				continue;
			}
			if ( preg_match( '/\G\d+(\.\d+)?/', $sql, $m, 0, $i ) ) {
				$tokens[] = array(
					't' => 'num',
					'v' => $m[0],
				);
				$i       += strlen( $m[0] );
				continue;
			}
			if ( preg_match( '/\G(<=>|<>|!=|<=|>=|\|\||&&|:=)/', $sql, $m, 0, $i ) ) {
				if ( ':=' === $m[0] ) {
					throw new InvalidArgumentException( 'Assignments are not allowed.' );
				}
				$tokens[] = array(
					't' => 'sym',
					'v' => $m[0],
				);
				$i       += strlen( $m[0] );
				continue;
			}
			$tokens[] = array(
				't' => 'sym',
				'v' => $c,
			);
			++$i;
		}//end while
		return $tokens;
	}

	/**
	 * Names defined by WITH name AS (…), so they may be used as tables.
	 *
	 * @param list<array{t: string, v: string}> $tokens Tokens.
	 * @return array<string, true>
	 */
	private static function cte_names( array $tokens ): array {
		$names = array();
		$count = count( $tokens );
		for ( $i = 0; $i + 2 < $count; $i++ ) {
			$is_name = in_array( $tokens[ $i ]['t'], array( 'word', 'ident' ), true );
			if ( $is_name && 'as' === strtolower( $tokens[ $i + 1 ]['v'] ) && '(' === $tokens[ $i + 2 ]['v'] ) {
				$prev = $i > 0 ? strtolower( $tokens[ $i - 1 ]['v'] ) : '';
				if ( in_array( $prev, array( 'with', ',', 'recursive' ), true ) ) {
					$names[ strtolower( $tokens[ $i ]['v'] ) ] = true;
				}
			}
		}
		return $names;
	}

	/**
	 * Whether a * is a column wildcard (SELECT *, t.*) rather than multiplication or COUNT(*).
	 *
	 * @param list<array{t: string, v: string}> $tokens Tokens.
	 * @param int                               $i      Index of the *.
	 */
	private static function is_select_star( array $tokens, int $i ): bool {
		$prev = $i > 0 ? strtolower( $tokens[ $i - 1 ]['v'] ) : '';
		if ( in_array( $prev, array( 'select', 'distinct', ',' ), true ) ) {
			return true;
		}
		// t.* at the start of a select item.
		if ( '.' === $prev && $i >= 3 ) {
			$before = strtolower( $tokens[ $i - 3 ]['v'] );
			return in_array( $before, array( 'select', 'distinct', ',' ), true );
		}
		return false;
	}

	/**
	 * LIMIT n, LIMIT m, n or LIMIT n OFFSET m, with n at most MAX_LIMIT.
	 *
	 * @param list<array{t: string, v: string}> $tokens Tokens.
	 * @param int                               $i      Index of LIMIT.
	 * @throws InvalidArgumentException When the limit is too large or not a number.
	 */
	private static function check_limit( array $tokens, int $i ): void {
		$first = $tokens[ $i + 1 ] ?? null;
		if ( null === $first || 'num' !== $first['t'] ) {
			throw new InvalidArgumentException( 'LIMIT must be a number.' );
		}
		$count = (int) $first['v'];
		if ( isset( $tokens[ $i + 2 ] ) && ',' === $tokens[ $i + 2 ]['v'] ) {
			$count = (int) ( $tokens[ $i + 3 ]['v'] ?? 0 );
		}
		if ( $count > self::MAX_LIMIT ) {
			throw new InvalidArgumentException( sprintf( 'LIMIT can be at most %d; only %d rows are returned anyway. Aggregate instead.', self::MAX_LIMIT, self::MAX_ROWS ) );
		}
	}

	/**
	 * A name without the WordPress prefix, if it has it.
	 *
	 * @param string $name   Lowercased name.
	 * @param string $prefix Prefix.
	 */
	private static function unprefix( string $name, string $prefix ): string {
		$prefix = strtolower( $prefix );
		return '' !== $prefix && str_starts_with( $name, $prefix ) ? substr( $name, strlen( $prefix ) ) : $name;
	}

	/**
	 * Tables in this database with the WordPress prefix, lowercased.
	 *
	 * @return array<string, true>
	 */
	private static function real_tables(): array {
		global $wpdb;
		// phpcs:ignore WordPress.DB.DirectDatabaseQuery.DirectQuery, WordPress.DB.DirectDatabaseQuery.NoCaching -- Table list for the allowlist check.
		$names = $wpdb->get_col( $wpdb->prepare( 'SHOW TABLES LIKE %s', $wpdb->esc_like( $wpdb->prefix ) . '%' ) );
		$out   = array();
		foreach ( (array) $names as $name ) {
			$out[ strtolower( (string) $name ) ] = true;
		}
		return $out;
	}

	/**
	 * Joins tokens back into SQL.
	 *
	 * @param list<array{t: string, v: string, raw?: bool}> $tokens Tokens.
	 */
	private static function join_tokens( array $tokens ): string {
		$out = '';
		foreach ( $tokens as $tok ) {
			$text = 'ident' === $tok['t'] && empty( $tok['raw'] ) ? '`' . str_replace( '`', '', $tok['v'] ) . '`' : $tok['v'];
			// No space before "(": MySQL refuses "COUNT (*)" unless IGNORE_SPACE is set.
			$glue = '' === $out || in_array( $tok['v'], array( '.', ',', ')', '(' ), true ) || str_ends_with( $out, '.' ) || str_ends_with( $out, '(' ) ? '' : ' ';
			$out .= $glue . $text;
		}
		return $out;
	}
}