<?php
/**
 * Rate Limiter.
 *
 * Tracks token usage per minute and enforces rate limits. Uses WordPress
 * transients for persistent storage across requests. Rate limits are scoped
 * per provider + model combination since different providers and models have
 * completely separate limits.
 *
 * @package   Agenistrator
 * @since     1.0.0
 * @author    Alex Ruban <alex@maximum.guru>
 * @author    Maximum Software Inc <contact@maximum.software>
 * @license   GPL-3.0-only
 * @link      https://agenistrator.com/
 *
 * SPDX-FileCopyrightText: 2025-2026 Alex Ruban <alex@maximum.guru>
 * SPDX-FileCopyrightText: 2025-2026 Maximum Software Inc <contact@maximum.software>
 * SPDX-License-Identifier: GPL-3.0-only
 */

namespace Agenistrator;

// Direct access prevention
if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

/**
 * Rate_Limiter class.
 */
class Rate_Limiter {
	use Singleton;

	/**
	 * Transient key prefix for storing usage data
	 */
	private const TRANSIENT_PREFIX = 'agenistrator_rl_';

	/**
	 * Window size in seconds
	 */
	private const WINDOW_SIZE = 60;

	/**
	 * Record token usage for a specific model
	 *
	 * @param LLM $model  The model instance to scope usage to
	 * @param int $tokens Number of tokens used
	 * @since     1.0.0
	 */
	public function record_usage( LLM $model, int $tokens ) {
		$key   = $this->get_scope_key( $model );
		$usage = $this->get_usage_data( $key );
		$now   = time();

		// Add new entry
		$usage[] = array(
			'time'   => $now,
			'tokens' => $tokens,
		);

		// Clean up old entries
		$usage = $this->prune_old_entries( $usage, $now );

		// Store with expiry slightly longer than window
		set_transient( $key, wp_json_encode( $usage ), self::WINDOW_SIZE + 10 );
	}

	/**
	 * Get tokens used in the current window for a specific model
	 *
	 * @param LLM $model The model instance to check usage for
	 * @return int Total tokens used in the last minute
	 * @since     1.0.0
	 */
	public function get_tokens_used( LLM $model ) {
		$key   = $this->get_scope_key( $model );
		$usage = $this->get_usage_data( $key );
		$usage = $this->prune_old_entries( $usage, time() );

		$total = 0;
		foreach ( $usage as $entry ) {
			$total += $entry['tokens'];
		}
		return $total;
	}

	/**
	 * Get available tokens before hitting rate limit for a specific model
	 *
	 * @param LLM $model The model instance to check availability for
	 * @return int Available tokens (PHP_INT_MAX if no limit configured)
	 * @since     1.0.0
	 */
	public function get_available_tokens( LLM $model ) {
		$max = $model->get_provider()->get_max_tokens_per_minute();

		// No rate limit configured for this provider
		if ( null === $max ) {
			return PHP_INT_MAX;
		}

		$used = $this->get_tokens_used( $model );
		return max( 0, $max - $used );
	}

	/**
	 * Check if we need to wait and calculate wait time for a specific model
	 *
	 * @param LLM $model         The model instance to check wait time for
	 * @param int $tokens_needed Tokens needed for the next request
	 * @return float Seconds to wait (0 if no wait needed)
	 * @since     1.0.0
	 */
	public function get_wait_time( LLM $model, int $tokens_needed ) {
		$available = $this->get_available_tokens( $model );

		if ( $tokens_needed <= $available ) {
			return 0;
		}

		// Calculate when oldest entries will expire
		$key   = $this->get_scope_key( $model );
		$usage = $this->get_usage_data( $key );
		$now   = time();
		$usage = $this->prune_old_entries( $usage, $now );

		if ( empty( $usage ) ) {
			return 0;
		}

		// Sort by time
		usort(
			$usage,
			function ( $a, $b ) {
				return $a['time'] - $b['time'];
			}
		);

		// Find how long until enough tokens are freed
		$tokens_to_free = $tokens_needed - $available;
		$freed          = 0;
		$wait_until     = $now;

		foreach ( $usage as $entry ) {
			$freed     += $entry['tokens'];
			$wait_until = $entry['time'] + self::WINDOW_SIZE;
			if ( $freed >= $tokens_to_free ) {
				break;
			}
		}

		return max( 0, $wait_until - $now + 1 );
	}

	/**
	 * Generate a scoped transient key for a model
	 *
	 * Uses provider ID and model ID to create a unique scope, ensuring
	 * that usage tracking is isolated per provider + model combination.
	 *
	 * @param LLM $model The model instance
	 * @return string Transient key (max 43 chars: 11 prefix + 32 hash)
	 */
	private function get_scope_key( LLM $model ) {
		$provider_id = $model->get_provider()->get_id();
		$model_id    = $model->get_model_id();
		return self::TRANSIENT_PREFIX . md5( $provider_id . '|' . $model_id );
	}

	/**
	 * Get usage data from transient
	 *
	 * @param string $key Transient key
	 * @return array Usage entries
	 */
	private function get_usage_data( $key ) {
		$usage_json = get_transient( $key );
		if ( ! is_string( $usage_json ) ) {
			return array();
		}
		$usage = json_decode( $usage_json, true );
		return is_array( $usage ) ? $usage : array();
	}

	/**
	 * Remove entries older than the window
	 *
	 * @param array $usage Usage entries
	 * @param int   $now   Current timestamp
	 * @return array Pruned entries
	 */
	private function prune_old_entries( array $usage, int $now ) {
		$cutoff = $now - self::WINDOW_SIZE;
		return array_values(
			array_filter(
				$usage,
				function ( $entry ) use ( $cutoff ) {
					return $entry['time'] > $cutoff;
				}
			)
		);
	}
}