<?php
/**
 * Security rate limiting. There is no usage quota; these limits exist only to
 * stop runaway agents, loops, brute force and denial of service.
 *
 * @package HandyMCP
 */

namespace HandyMCP\Security;

use HandyMCP\Settings;

defined( 'ABSPATH' ) || exit;

/**
 * Fixed-window counters stored in transients (object-cache aware).
 */
final class RateLimiter {

	/** @var Settings */
	private $settings;

	/**
	 * Constructor.
	 *
	 * @param Settings $settings Settings.
	 */
	public function __construct( Settings $settings ) {
		$this->settings = $settings;
	}

	/**
	 * Window length for a bucket.
	 *
	 * @param string $bucket Bucket.
	 */
	private function window( string $bucket ): int {
		if ( 'dcr' === $bucket ) {
			return HOUR_IN_SECONDS;
		}
		return 'auth_fail' === $bucket ? 10 * MINUTE_IN_SECONDS : MINUTE_IN_SECONDS;
	}

	/**
	 * Limit for a bucket.
	 *
	 * @param string $bucket read | write | destructive | scan | auth_fail | dcr.
	 */
	public function limit( string $bucket ): int {
		if ( 'dcr' === $bucket ) {
			// Client registrations per IP address and hour.
			return 10;
		}
		return max( 1, (int) $this->settings->get( 'rate_' . $bucket, 60 ) );
	}

	/**
	 * Transient key for the current window.
	 *
	 * @param string $key    Identity key.
	 * @param string $bucket Bucket.
	 */
	private function slot( string $key, string $bucket ): string {
		$window = $this->window( $bucket );
		return 'handymcp_rl_' . substr( md5( $key . '|' . $bucket . '|' . (int) floor( time() / $window ) ), 0, 24 );
	}

	/**
	 * Current count.
	 *
	 * @param string $key    Identity key.
	 * @param string $bucket Bucket.
	 */
	public function count( string $key, string $bucket ): int {
		$slot = $this->slot( $key, $bucket );
		if ( wp_using_ext_object_cache() ) {
			$v = wp_cache_get( $slot, 'handymcp_rl' );
			if ( false !== $v ) {
				return (int) $v;
			}
		}
		// Read the row directly: hit() writes it with SQL, so a cached transient could be stale.
		global $wpdb;
		$timeout = (int) $wpdb->get_var( $wpdb->prepare( "SELECT option_value FROM {$wpdb->options} WHERE option_name = %s", '_transient_timeout_' . $slot ) ); // phpcs:ignore WordPress.DB.DirectDatabaseQuery.DirectQuery, WordPress.DB.DirectDatabaseQuery.NoCaching -- counter must be current.
		if ( $timeout && $timeout < time() ) {
			return 0;
		}
		return (int) $wpdb->get_var( $wpdb->prepare( "SELECT option_value FROM {$wpdb->options} WHERE option_name = %s", '_transient_' . $slot ) ); // phpcs:ignore WordPress.DB.DirectDatabaseQuery.DirectQuery, WordPress.DB.DirectDatabaseQuery.NoCaching -- counter must be current.
	}

	/**
	 * Increment a counter.
	 *
	 * @param string $key    Identity key.
	 * @param string $bucket Bucket.
	 */
	public function hit( string $key, string $bucket ): int {
		$slot = $this->slot( $key, $bucket );
		// Atomic increment: object cache when present, otherwise one UPDATE in the database.
		if ( wp_using_ext_object_cache() ) {
			wp_cache_add( $slot, 0, 'handymcp_rl', $this->window( $bucket ) + 5 );
			$n = wp_cache_incr( $slot, 1, 'handymcp_rl' );
			if ( false !== $n ) {
				return (int) $n;
			}
		}
		global $wpdb;
		$name    = '_transient_' . $slot;
		$timeout = '_transient_timeout_' . $slot;
		// phpcs:disable WordPress.DB.DirectDatabaseQuery
		$wpdb->query( $wpdb->prepare( "INSERT INTO {$wpdb->options} (option_name, option_value, autoload) VALUES (%s, %s, 'off') ON DUPLICATE KEY UPDATE option_value = option_value", $timeout, (string) ( time() + $this->window( $bucket ) + 5 ) ) );
		$wpdb->query( $wpdb->prepare( "INSERT INTO {$wpdb->options} (option_name, option_value, autoload) VALUES (%s, '1', 'off') ON DUPLICATE KEY UPDATE option_value = option_value + 1", $name ) );
		$n = (int) $wpdb->get_var( $wpdb->prepare( "SELECT option_value FROM {$wpdb->options} WHERE option_name = %s", $name ) );
		// phpcs:enable
		wp_cache_delete( $slot, 'transient' );
		return $n;
	}

	/**
	 * Check (and by default consume) one unit.
	 *
	 * @param string $key       Identity key.
	 * @param string $bucket    Bucket.
	 * @param bool   $increment Consume a unit.
	 */
	public function allow( string $key, string $bucket, bool $increment = true ): bool {
		if ( $increment ) {
			// Count first, then compare the new total: parallel requests cannot all slip under the limit.
			return $this->hit( $key, $bucket ) <= $this->limit( $bucket );
		}
		return $this->count( $key, $bucket ) < $this->limit( $bucket );
	}

	/**
	 * Seconds until the current window resets.
	 *
	 * @param string $bucket Bucket.
	 */
	public function retry_after( string $bucket ): int {
		$w = $this->window( $bucket );
		return $w - ( time() % $w );
	}
}