<?php
/**
 * Aoedia per-user token limits.
 *
 * Each role gets a daily token budget. A user's usage is the sum of the
 * tokens logged for their chat requests since midnight in the site's time
 * zone, so the audit log is the single source of truth.
 *
 * @package Aoedia
 */

defined( 'ABSPATH' ) || exit;

/** Daily token limit for roles without a saved limit. 0 means no limit. */
const AOEDIA_DEFAULT_TOKEN_LIMIT = 100000;

/**
 * Returns the roles that can use Aoedia.
 *
 * @return array<string, string> Translated role names, keyed by role.
 */
function aoedia_get_assistant_roles(): array {
	$roles = array();
	foreach ( wp_roles()->role_objects as $name => $role ) {
		if ( $role->has_cap( 'edit_posts' ) ) {
			$roles[ $name ] = translate_user_role( wp_roles()->role_names[ $name ] );
		}
	}

	return $roles;
}

/**
 * Returns the daily token limit of each role that can use Aoedia.
 *
 * Administrators have no limit until one is saved; other roles get
 * AOEDIA_DEFAULT_TOKEN_LIMIT.
 *
 * @return array<string, int> Limits keyed by role. 0 means no limit.
 */
function aoedia_get_token_limits(): array {
	$saved  = (array) get_option( 'aoedia_token_limits', array() );
	$limits = array();

	foreach ( array_keys( aoedia_get_assistant_roles() ) as $role ) {
		if ( isset( $saved[ $role ] ) ) {
			$limits[ $role ] = max( 0, (int) $saved[ $role ] );
		} else {
			$limits[ $role ] = 'administrator' === $role ? 0 : AOEDIA_DEFAULT_TOKEN_LIMIT;
		}
	}

	return $limits;
}

/**
 * Sanitizes the saved role limits.
 *
 * @param mixed $value The submitted limits.
 * @return array<string, int> Limits keyed by role.
 */
function aoedia_sanitize_token_limits( $value ): array {
	$limits = array();
	foreach ( array_keys( aoedia_get_assistant_roles() ) as $role ) {
		$limits[ $role ] = is_array( $value ) && isset( $value[ $role ] ) ? max( 0, (int) $value[ $role ] ) : 0;
	}

	return $limits;
}

/**
 * Returns a user's daily token limit.
 *
 * A user with several roles gets the most generous limit among them.
 *
 * @param WP_User $user The user.
 * @return int The limit. 0 means no limit.
 */
function aoedia_get_user_token_limit( WP_User $user ): int {
	$limits = aoedia_get_token_limits();
	$limit  = null;

	if ( is_multisite() && is_super_admin( $user->ID ) ) {
		$limit = 0;
	} else {
		foreach ( $user->roles as $role ) {
			$role_limit = $limits[ $role ] ?? AOEDIA_DEFAULT_TOKEN_LIMIT;
			if ( 0 === $role_limit ) {
				$limit = 0;
				break;
			}
			$limit = max( (int) $limit, $role_limit );
		}
	}

	/**
	 * Filters a user's daily token limit.
	 *
	 * @param int     $limit The limit. 0 means no limit.
	 * @param WP_User $user  The user.
	 */
	return max( 0, (int) apply_filters( 'aoedia_token_limit', $limit ?? AOEDIA_DEFAULT_TOKEN_LIMIT, $user ) );
}

/**
 * Returns how many tokens a user has used today.
 *
 * @param int $user_id The user ID.
 * @return int The number of tokens.
 */
function aoedia_get_user_tokens_used( int $user_id ): int {
	global $wpdb;

	$since = ( new DateTimeImmutable( 'today', wp_timezone() ) )->setTimezone( new DateTimeZone( 'UTC' ) );

	// phpcs:ignore WordPress.DB.DirectDatabaseQuery -- Plugin table; must be current, so not cached.
	return (int) $wpdb->get_var(
		$wpdb->prepare(
			'SELECT COALESCE( SUM( tokens ), 0 ) FROM %i WHERE user_id = %d AND created_at >= %s',
			aoedia_log_table(),
			$user_id,
			$since->format( 'Y-m-d H:i:s' )
		)
	);
}

/**
 * Returns a user's token usage for today.
 *
 * @param WP_User|null $user The user. Default the current user.
 * @return array{used: int, limit: int, resets: string} Tokens used, the limit (0 for none) and when it resets, as an ISO 8601 date.
 */
function aoedia_get_token_usage( ?WP_User $user = null ): array {
	$user = $user ?? wp_get_current_user();

	return array(
		'used'   => aoedia_get_user_tokens_used( $user->ID ),
		'limit'  => aoedia_get_user_token_limit( $user ),
		'resets' => ( new DateTimeImmutable( 'tomorrow', wp_timezone() ) )->format( DATE_ATOM ),
	);
}

/**
 * Returns the error for a user who has reached their limit.
 *
 * @param array{used: int, limit: int, resets: string} $usage The user's usage.
 * @return WP_Error The error.
 */
function aoedia_token_limit_error( array $usage ): WP_Error {
	return new WP_Error(
		'aoedia_token_limit',
		sprintf(
			/* translators: 1: number of tokens, 2: time. */
			__( 'You have used your daily limit of %1$s tokens. It resets at %2$s.', 'aoedia' ),
			number_format_i18n( $usage['limit'] ),
			wp_date( get_option( 'time_format' ), strtotime( $usage['resets'] ) )
		),
		array( 'status' => 429 )
	);
}