<?php
/**
 * Atomic DB-backed rate limiter for chatbot requests.
 *
 * @package AIPT\Api\Chatbot
 * @since 2.5.0
 */

namespace AIPT\Api\Chatbot;

/**
 * RateLimiter class
 *
 * Implements fixed-window rate limiting using an atomic
 * INSERT ... ON DUPLICATE KEY UPDATE pattern against the
 * aipt_rate_limit_counters table. This eliminates the TOCTOU
 * race condition present in the previous transient-based
 * implementation.
 *
 * Public API is intentionally unchanged from the transient-based
 * version so callers require no modification.
 *
 * @see https://en.wikipedia.org/wiki/Time-of-check_to_time-of-use
 */
class RateLimiter {

	/**
	 * Key prefix used to namespace all bucket_key values in the DB table.
	 */
	private const KEY_PREFIX = 'aipt_chat_rate_';

	private ChatbotSettings $settings;

	private string $rate_limit_table;

	public function __construct( ?ChatbotSettings $settings = null ) {
		global $wpdb;
		$this->settings          = $settings ?? new ChatbotSettings();
		$this->rate_limit_table  = $wpdb->prefix . 'aipt_rate_limit_counters';
	}

	/**
	 * Check if request is allowed under rate limit.
	 *
	 * @param string $session_id Visitor session ID.
	 * @return array{allowed: bool, retry_after: int} Rate limit status.
	 */
	public function check( string $session_id ): array {
		if ( ! $this->isEnabled() ) {
			return [ 'allowed' => true, 'retry_after' => 0 ];
		}

		$key    = $this->buildKey( $session_id );
		$limit  = $this->getMaxRequests();
		$window = $this->getWindowSeconds();

		return $this->incrementAndCheck( $key, $limit, $window );
	}

	/**
	 * Check if rate limiting is enabled.
	 *
	 * @return bool True if enabled.
	 */
	public function isEnabled(): bool {
		return (bool) $this->settings->get( 'rate_limit.enabled', true );
	}

	/**
	 * Get maximum requests per window.
	 *
	 * @return int Maximum requests.
	 */
	public function getMaxRequests(): int {
		return (int) $this->settings->get( 'rate_limit.max_requests', 10 );
	}

	/**
	 * Get rate limit window in seconds.
	 *
	 * @return int Window seconds.
	 */
	public function getWindowSeconds(): int {
		return (int) $this->settings->get( 'rate_limit.window_seconds', 60 );
	}

	/**
	 * Check if request is allowed under IP-based rate limit.
	 *
	 * Provides a secondary rate limit check based on hashed client IP.
	 * This prevents session_id spoofing from bypassing rate limits.
	 *
	 * @param string $client_ip Client IP address.
	 * @return array{allowed: bool, retry_after: int} Rate limit status.
	 */
	public function checkByIp( string $client_ip ): array {
		if ( ! $this->isEnabled() ) {
			return [ 'allowed' => true, 'retry_after' => 0 ];
		}

		$hashed_ip = $this->hashIp( $client_ip );
		$key       = self::KEY_PREFIX . 'ip_' . $hashed_ip;
		$limit     = $this->getMaxRequests();
		$window    = $this->getWindowSeconds();

		return $this->incrementAndCheck( $key, $limit, $window );
	}

	/**
	 * Check if poll request is allowed under rate limit.
	 *
	 * Uses more lenient limits than chat requests since polling is frequent
	 * but lightweight. Defaults to 60 requests per 60 seconds.
	 *
	 * @param string $session_id Visitor session ID.
	 * @return array{allowed: bool, retry_after: int} Rate limit status.
	 */
	public function checkPoll( string $session_id ): array {
		if ( ! $this->isEnabled() ) {
			return [ 'allowed' => true, 'retry_after' => 0 ];
		}

		$key    = self::KEY_PREFIX . 'poll_' . md5( $session_id );
		$limit  = $this->getPollMaxRequests();
		$window = $this->getPollWindowSeconds();

		return $this->incrementAndCheck( $key, $limit, $window );
	}

	/**
	 * Check if a poll request is allowed under the poll IP bucket.
	 *
	 * Uses the same lenient limit/window as session-keyed checkPoll() but
	 * keyed by hashed IP. The chat /enqueue + /chat surfaces share
	 * checkByIp() (the strict send-message bucket); long-running async jobs
	 * polling /status need their own pre-DB DoS protection that does NOT
	 * exhaust the send-message budget — otherwise a valid 60-second job
	 * polled every few seconds would 429 the user's next chat send.
	 *
	 * @since 2.6.1
	 * @param string $client_ip Client IP address.
	 * @return array{allowed: bool, retry_after: int} Rate limit status.
	 */
	public function checkPollByIp( string $client_ip ): array {
		if ( ! $this->isEnabled() ) {
			return [ 'allowed' => true, 'retry_after' => 0 ];
		}

		$hashed_ip = $this->hashIp( $client_ip );
		$key       = self::KEY_PREFIX . 'poll_ip_' . $hashed_ip;
		$limit     = $this->getPollMaxRequests();
		$window    = $this->getPollWindowSeconds();

		return $this->incrementAndCheck( $key, $limit, $window );
	}

	/**
	 * Get maximum poll requests per window.
	 *
	 * @return int Maximum poll requests.
	 */
	public function getPollMaxRequests(): int {
		return (int) $this->settings->get( 'rate_limit.poll_max_requests', 60 );
	}

	/**
	 * Get poll rate limit window in seconds.
	 *
	 * @return int Window seconds.
	 */
	public function getPollWindowSeconds(): int {
		return (int) $this->settings->get( 'rate_limit.poll_window_seconds', 60 );
	}

	/**
	 * Get client IP address securely.
	 *
	 * Uses REMOTE_ADDR as the most reliable source by default.
	 * Only trusts proxy headers when the request comes from a known/trusted proxy IP.
	 *
	 * @return string Client IP address.
	 */
	public function getClientIp(): string {
		$remote_addr = '';
		if ( ! empty( $_SERVER['REMOTE_ADDR'] ) ) {
			$remote_addr = sanitize_text_field( wp_unslash( $_SERVER['REMOTE_ADDR'] ) );
			if ( ! filter_var( $remote_addr, FILTER_VALIDATE_IP ) ) {
				$remote_addr = '';
			}
		}

		$trusted_proxies  = $this->getTrustedProxyRanges();
		$is_trusted_proxy = $this->isIpInRanges( $remote_addr, $trusted_proxies );

		if ( $is_trusted_proxy ) {
			if ( ! empty( $_SERVER['HTTP_CF_CONNECTING_IP'] ) && ! empty( $_SERVER['HTTP_CF_RAY'] ) ) {
				$ip = sanitize_text_field( wp_unslash( $_SERVER['HTTP_CF_CONNECTING_IP'] ) );
				if ( filter_var( $ip, FILTER_VALIDATE_IP ) ) {
					return $ip;
				}
			}

			if ( ! empty( $_SERVER['HTTP_X_REAL_IP'] ) ) {
				$ip = sanitize_text_field( wp_unslash( $_SERVER['HTTP_X_REAL_IP'] ) );
				if ( filter_var( $ip, FILTER_VALIDATE_IP ) ) {
					return $ip;
				}
			}

			if ( ! empty( $_SERVER['HTTP_X_FORWARDED_FOR'] ) ) {
				$ip = sanitize_text_field( wp_unslash( $_SERVER['HTTP_X_FORWARDED_FOR'] ) );
				if ( strpos( $ip, ',' ) !== false ) {
					$ip = trim( explode( ',', $ip )[0] );
				}
				if ( filter_var( $ip, FILTER_VALIDATE_IP ) ) {
					return $ip;
				}
			}
		}

		return $remote_addr ?: '0.0.0.0';
	}

	/**
	 * Atomically increment a rate-limit counter and return whether the request
	 * is within the allowed window.
	 *
	 * The INSERT ... ON DUPLICATE KEY UPDATE pattern is a single round-trip to
	 * MySQL and is serialised by the engine's row-level lock on the PRIMARY KEY,
	 * eliminating the TOCTOU race that affected the old transient-based approach.
	 *
	 * Window reset logic (executed server-side to avoid PHP/MySQL clock skew):
	 *   - If the stored window is still open → increment counter, keep window_start.
	 *   - If the window has expired         → reset counter to 1, start new window.
	 *
	 * @param string $bucket_key Unique identifier for this rate-limit bucket.
	 * @param int    $limit      Maximum allowed requests per window.
	 * @param int    $window     Window length in seconds.
	 * @return array{allowed: bool, retry_after: int}
	 */
	private function incrementAndCheck( string $bucket_key, int $limit, int $window ): array {
		global $wpdb;

		$write_result = $wpdb->query(
			$wpdb->prepare(
				"INSERT INTO {$this->rate_limit_table}
					(bucket_key, counter, window_start, window_seconds)
				VALUES (%s, LAST_INSERT_ID(1), UNIX_TIMESTAMP(), %d)
				ON DUPLICATE KEY UPDATE
					counter        = LAST_INSERT_ID(IF(window_start + window_seconds > UNIX_TIMESTAMP(), counter + 1, 1)),
					window_start   = IF(window_start + window_seconds > UNIX_TIMESTAMP(), window_start, UNIX_TIMESTAMP()),
					window_seconds = VALUES(window_seconds)",
				$bucket_key,
				$window
			)
		);

		if ( false === $write_result ) {
			ChatbotLogger::error(
				'rate_limit_store_unavailable',
				[
					'wpdb_error' => isset( $wpdb->last_error ) ? (string) $wpdb->last_error : '',
				]
			);

			return [
				'allowed'     => false,

				'retry_after' => max( 1, $window ),
				'error_code'  => 'rate_limit_store_unavailable',
				'degraded'    => true,
			];
		}

		$count = (int) $wpdb->insert_id;
		if ( $count <= 0 ) {

			return [
				'allowed'     => false,
				'retry_after' => max( 1, min( $window, 30 ) ),
			];
		}

		if ( $count > $limit ) {

			$row = $wpdb->get_row(
				$wpdb->prepare(
					"SELECT window_start, window_seconds
					FROM {$this->rate_limit_table}
					WHERE bucket_key = %s",
					$bucket_key
				)
			);

			$remaining = $window;
			if ( null !== $row ) {
				$remaining = max( 0, ( (int) $row->window_start + (int) $row->window_seconds ) - time() );
				if ( 0 === $remaining ) {
					$remaining = 1; 
				}
			}

			return [
				'allowed'     => false,
				'retry_after' => $remaining,
			];
		}

		return [ 'allowed' => true, 'retry_after' => 0 ];
	}

	/**
	 * Build a rate-limit bucket key from a session ID and optional hashed IP.
	 *
	 * @param string $session_id Session identifier.
	 * @return string Bucket key (≤ 191 characters).
	 */
	private function buildKey( string $session_id ): string {
		$key_mode = $this->settings->get( 'rate_limit.key_mode', 'session' );
		$key      = $session_id;

		if ( 'session_ip_hashed' === $key_mode ) {
			$ip        = $this->getClientIp();
			$hashed_ip = $this->hashIp( $ip );
			$key      .= '_' . $hashed_ip;
		}

		return self::KEY_PREFIX . md5( $key );
	}

	/**
	 * Hash an IP address using SHA-256 with a site-specific salt.
	 *
	 * Never stores or logs the raw IP address.
	 *
	 * @param string $ip IP address to hash.
	 * @return string 64-character hex string.
	 */
	private function hashIp( string $ip ): string {
		return hash( 'sha256', $ip . wp_salt( 'auth' ) );
	}

	/**
	 * Get trusted proxy IP ranges.
	 *
	 * Returns Cloudflare and localhost IP ranges.
	 * Can be extended via filter for custom proxy configurations.
	 *
	 * @return array Array of CIDR ranges.
	 */
	private function getTrustedProxyRanges(): array {
		$default_ranges = [

			'173.245.48.0/20',
			'103.21.244.0/22',
			'103.22.200.0/22',
			'103.31.4.0/22',
			'141.101.64.0/18',
			'108.162.192.0/18',
			'190.93.240.0/20',
			'188.114.96.0/20',
			'197.234.240.0/22',
			'198.41.128.0/17',
			'162.158.0.0/15',
			'104.16.0.0/13',
			'104.24.0.0/14',
			'172.64.0.0/13',
			'131.0.72.0/22',

			'127.0.0.0/8',
			'10.0.0.0/8',
			'172.16.0.0/12',
			'192.168.0.0/16',
		];

		/**
		 * Filter trusted proxy IP ranges.
		 *
		 * @since 2.5.0
		 * @param array $ranges Array of CIDR ranges.
		 */
		return apply_filters( 'aipt_chatbot_trusted_proxy_ranges', $default_ranges );
	}

	/**
	 * Check if an IP is within any of the given CIDR ranges.
	 *
	 * @param string $ip     IP address to check.
	 * @param array  $ranges Array of CIDR ranges.
	 * @return bool True if IP is in any range.
	 */
	private function isIpInRanges( string $ip, array $ranges ): bool {
		if ( empty( $ip ) || ! filter_var( $ip, FILTER_VALIDATE_IP ) ) {
			return false;
		}

		$ip_long = ip2long( $ip );
		if ( false === $ip_long ) {
			return false;
		}

		foreach ( $ranges as $range ) {
			if ( strpos( $range, '/' ) === false ) {
				if ( $ip === $range ) {
					return true;
				}
				continue;
			}

			list( $range_ip, $mask ) = explode( '/', $range, 2 );
			$range_long = ip2long( $range_ip );
			if ( false === $range_long ) {
				continue;
			}

			$mask      = (int) $mask;
			$mask_long = -1 << ( 32 - $mask );

			if ( ( $ip_long & $mask_long ) === ( $range_long & $mask_long ) ) {
				return true;
			}
		}

		return false;
	}

	/**
	 * @deprecated Use incrementAndCheck() instead.
	 */
	private function legacyCheckTransient( string $key, int $limit, int $window ): array {
		$data = get_transient( $key );

		if ( false === $data || ! is_array( $data ) ) {
			$window_start = time();
			set_transient( $key, [ 'count' => 1, 'window_start' => $window_start ], $window );
			return [ 'allowed' => true, 'retry_after' => 0 ];
		}

		$count        = isset( $data['count'] ) ? (int) $data['count'] : 0;
		$window_start = isset( $data['window_start'] ) ? (int) $data['window_start'] : time();
		$remaining    = max( 0, ( $window_start + $window ) - time() );

		if ( $remaining <= 0 ) {
			$window_start = time();
			set_transient( $key, [ 'count' => 1, 'window_start' => $window_start ], $window );
			return [ 'allowed' => true, 'retry_after' => 0 ];
		}

		if ( $count >= $limit ) {
			return [ 'allowed' => false, 'retry_after' => $remaining ];
		}

		set_transient( $key, [ 'count' => $count + 1, 'window_start' => $window_start ], $remaining );
		return [ 'allowed' => true, 'retry_after' => 0 ];
	}
}
