<?php
/**
 * Usage limits for the AI Advisor.
 *
 * @package TurtleReport
 */

namespace TurtleReport\Advisor;

use TurtleReport\Support\OptionUpdate;
use TurtleReport\Support\TransientStore;

defined( 'ABSPATH' ) || exit;

/**
 * Counts questions per day and per month for the whole site, plus the
 * tokens they used, and stops questions once a limit is reached. Each
 * person can also only ask one question at a time, and a few per minute.
 *
 * Stored in one small option (not autoloaded). Days and months follow the
 * site's time zone.
 */
final class Usage {

	/** Option. */
	const OPTION = 'turtlereport_ai_usage';

	/** Questions one person may ask per minute. */
	const PER_MINUTE = 6;

	/** Outstanding requests expire after the longest provider retry sequence. */
	const RESERVATION_TTL = 360;

	/**
	 * Reservations owned by this request, scoped by blog/user.
	 *
	 * @var array<string,string>
	 */
	private $owners = array();

	/**
	 * The counters, reset when the day or month changes.
	 *
	 * @param int|null $now Timestamp.
	 * @return array{day:string,day_n:int,month:string,month_n:int,in:int,out:int}
	 */
	public function get( $now = null ) {
		$now = null === $now ? time() : (int) $now;
		$u   = get_option( self::OPTION, array() );
		$u   = self::normalize( is_array( $u ) ? $u : array(), $now );
		unset( $u['pending'] );
		return $u;
	}

	/**
	 * Normalizes counters and discards expired reservations under the write lock.
	 *
	 * @param array $u Stored usage.
	 * @param int   $now Timestamp.
	 * @return array
	 */
	private static function normalize( array $u, $now ) {
		$day   = wp_date( 'Y-m-d', $now );
		$month = wp_date( 'Y-m', $now );
		$u     = array_merge(
			array(
				'day'     => $day,
				'day_n'   => 0,
				'month'   => $month,
				'month_n' => 0,
				'in'      => 0,
				'out'     => 0,
				'pending' => array(),
			),
			$u
		);
		if ( $u['day'] !== $day ) {
			$u['day']   = $day;
			$u['day_n'] = 0;
		}
		if ( $u['month'] !== $month ) {
			$u['month']   = $month;
			$u['month_n'] = 0;
			$u['in']      = 0;
			$u['out']     = 0;
		}
		foreach ( array( 'day_n', 'month_n', 'in', 'out' ) as $key ) {
			$u[ $key ] = is_numeric( $u[ $key ] ) ? max( 0, (int) $u[ $key ] ) : 0;
		}
		$u['pending'] = is_array( $u['pending'] ) ? array_filter(
			$u['pending'],
			static function ( $row ) use ( $now ) {
				return is_array( $row ) && isset( $row['owner'], $row['expires'], $row['day'], $row['month'] ) && is_string( $row['owner'] ) && (int) $row['expires'] > $now;
			}
		) : array();
		return $u;
	}

	/**
	 * Whether another question is allowed.
	 *
	 * @param array    $limits  day and month limits (0 = no limit).
	 * @param int      $user_id Who is asking.
	 * @param int|null $now     Timestamp.
	 * @return true|\WP_Error
	 */
	public function allow( array $limits, $user_id, $now = null ) {
		$now = null === $now ? time() : (int) $now;
		$raw = get_option( self::OPTION, array() );
		return self::check( self::normalize( is_array( $raw ) ? $raw : array(), $now ), $limits, $user_id, $now );
	}

	/**
	 * Checks answered questions plus in-flight reservations.
	 *
	 * @param array $u Usage.
	 * @param array $limits Limits.
	 * @param int   $user_id User.
	 * @param int   $now Timestamp.
	 * @return true|\WP_Error
	 */
	private static function check( array $u, array $limits, $user_id, $now ) {
		if ( isset( $u['pending'][ $user_id ] ) ) {
			return new \WP_Error( 'turtlereport_ai_busy', __( 'Your previous question is still being answered.', 'turtlereport' ), array( 'status' => 429 ) );
		}
		$day   = $u['day_n'];
		$month = $u['month_n'];
		foreach ( $u['pending'] as $reservation ) {
			$day   += $reservation['day'] === $u['day'] ? 1 : 0;
			$month += $reservation['month'] === $u['month'] ? 1 : 0;
		}
		if ( ! empty( $limits['day'] ) && $day >= (int) $limits['day'] ) {
			/* translators: %d: number of questions. */
			return new \WP_Error( 'turtlereport_ai_day_limit', sprintf( _n( 'Today\'s limit of %d AI question has been reached. It resets tomorrow, or an administrator can raise it in Settings → AI Advisor.', 'Today\'s limit of %d AI questions has been reached. It resets tomorrow, or an administrator can raise it in Settings → AI Advisor.', (int) $limits['day'], 'turtlereport' ), (int) $limits['day'] ), array( 'status' => 429 ) );
		}
		if ( ! empty( $limits['month'] ) && $month >= (int) $limits['month'] ) {
			/* translators: %d: number of questions. */
			return new \WP_Error( 'turtlereport_ai_month_limit', sprintf( _n( 'This month\'s limit of %d AI question has been reached. An administrator can raise it in Settings → AI Advisor.', 'This month\'s limit of %d AI questions has been reached. An administrator can raise it in Settings → AI Advisor.', (int) $limits['month'], 'turtlereport' ), (int) $limits['month'] ), array( 'status' => 429 ) );
		}
		$recent = get_transient( 'turtlereport_ai_rate_' . (int) $user_id );
		$recent = is_array( $recent ) ? array_values(
			array_filter(
				$recent,
				static function ( $t ) use ( $now ) {
					return is_int( $t ) && $t > $now - 60;
				}
			)
		) : array();
		if ( count( $recent ) >= self::PER_MINUTE ) {
			return new \WP_Error( 'turtlereport_ai_rate', __( 'That\'s a lot of questions in a minute. Wait a moment and ask again.', 'turtlereport' ), array( 'status' => 429 ) );
		}
		return true;
	}

	/**
	 * Atomically reserves one question within the site-wide limits.
	 *
	 * @param array $limits Limits.
	 * @param int   $user_id User.
	 * @return true|\WP_Error
	 */
	public function reserve( array $limits, $user_id ) {
		$owner   = wp_generate_uuid4();
		$allowed = true;
		OptionUpdate::change(
			self::OPTION,
			static function ( array $stored ) use ( $limits, $user_id, $owner, &$allowed ) {
				$now     = time();
				$u       = self::normalize( $stored, $now );
				$allowed = self::check( $u, $limits, $user_id, $now );
				if ( is_wp_error( $allowed ) ) {
					return $u;
				}
				$u['pending'][ $user_id ] = array(
					'owner'   => $owner,
					'expires' => $now + self::RESERVATION_TTL,
					'day'     => $u['day'],
					'month'   => $u['month'],
				);
				return $u;
			}
		);
		if ( true === $allowed ) {
			$this->owners[ get_current_blog_id() . '|' . (int) $user_id ] = $owner;
			$recent   = get_transient( 'turtlereport_ai_rate_' . (int) $user_id );
			$recent   = is_array( $recent ) ? $recent : array();
			$recent[] = time();
			TransientStore::set( 'turtlereport_ai_rate_' . (int) $user_id, array_slice( $recent, -10 ), 120 );
		}
		return $allowed;
	}

	/**
	 * Finalizes only this request's reservation, without releasing a new owner.
	 *
	 * @param int  $user_id User.
	 * @param int  $in Input tokens.
	 * @param int  $out Output tokens.
	 * @param bool $answered Whether an answer was produced.
	 * @return void
	 */
	public function finish( $user_id, $in = 0, $out = 0, $answered = false ) {
		$key   = get_current_blog_id() . '|' . (int) $user_id;
		$owner = isset( $this->owners[ $key ] ) ? $this->owners[ $key ] : '';
		if ( '' === $owner ) {
			return;
		}
		OptionUpdate::change(
			self::OPTION,
			static function ( array $stored ) use ( $owner, $user_id, $in, $out, $answered ) {
				$u = self::normalize( $stored, time() );
				if ( ! isset( $u['pending'][ $user_id ] ) || $owner !== $u['pending'][ $user_id ]['owner'] ) {
					return $u;
				}
				unset( $u['pending'][ $user_id ] );
				if ( $answered ) {
					++$u['day_n'];
					++$u['month_n'];
					$u['in']  += max( 0, (int) $in );
					$u['out'] += max( 0, (int) $out );
				}
				return $u;
			}
		);
		unset( $this->owners[ $key ] );
	}

	/**
	 * Records an answered question.
	 *
	 * @param int      $in  Input tokens.
	 * @param int      $out Output tokens.
	 * @param bool     $question Counts as a question (a connection test doesn't).
	 * @param int|null $now Timestamp.
	 * @return array The new counters.
	 */
	public function record( $in, $out, $question = true, $now = null ) {
		$u = OptionUpdate::change(
			self::OPTION,
			static function ( array $stored ) use ( $in, $out, $question, $now ) {
				$u = self::normalize( $stored, null === $now ? time() : (int) $now );
				if ( $question ) {
					++$u['day_n'];
					++$u['month_n'];
				}
				$u['in']  += max( 0, (int) $in );
				$u['out'] += max( 0, (int) $out );
				return $u;
			}
		);
		unset( $u['pending'] );
		return $u;
	}
}