<?php
if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

/**
 * How much THIS site consumed.
 *
 * The service already reported every generation's usage and the client threw it away. That was
 * not a detail: without it, whoever uses the plugin has no way of knowing they are burning their
 * budget on a consistency pass over a twenty-section page, and whoever sells it has no way of
 * answering "how much have I used this month?".
 *
 * Two options, not a table: a running total for the current month (which resets itself when the
 * month changes) and the last twenty generations for support. It is deliberately little — the
 * real accounting lives on the service, the one place a customer cannot rewrite it.
 */
class AIPB_Usage {

	const OPT_TOTALS = 'aipb_usage_totals';
	const OPT_RECENT = 'aipb_usage_recent';

	/** How many recent generations are kept. Enough to find the one somebody is asking about. */
	const KEEP = 20;

	/**
	 * Records one generation. Called where the service's answer is parsed, so it covers every
	 * road — synchronous or streamed.
	 *
	 * @param array  $answer Answer from AIPB_Remote::generate(): usage, model, request_id.
	 * @param string $mode   section | patch | global | plan | tokens | model.
	 */
	public static function record( $answer, $mode ) {
		if ( ! is_array( $answer ) ) {
			return;
		}

		$usage = isset( $answer['usage'] ) && is_array( $answer['usage'] ) ? $answer['usage'] : array();

		self::add_to_totals( $usage );
		self::remember( $answer, $mode, $usage );
	}

	/**
	 * The current month's totals.
	 *
	 * @return array{month:string,count:int,input:int,output:int,cache_read:int}
	 */
	public static function totals() {
		$totals = get_option( self::OPT_TOTALS, array() );
		$totals = is_array( $totals ) ? $totals : array();

		return array(
			'month'      => isset( $totals['month'] ) ? (string) $totals['month'] : gmdate( 'Y-m' ),
			'count'      => isset( $totals['count'] ) ? (int) $totals['count'] : 0,
			'input'      => isset( $totals['input'] ) ? (int) $totals['input'] : 0,
			'output'     => isset( $totals['output'] ) ? (int) $totals['output'] : 0,
			'cache_read' => isset( $totals['cache_read'] ) ? (int) $totals['cache_read'] : 0,
		);
	}

	/** The latest generations, newest first. Support needs them: they carry the request_id. */
	public static function recent() {
		$recent = get_option( self::OPT_RECENT, array() );

		return is_array( $recent ) ? $recent : array();
	}

	/** @param array $usage token counters as the service reported them */
	protected static function add_to_totals( $usage ) {
		$month  = gmdate( 'Y-m' );
		$totals = get_option( self::OPT_TOTALS, array() );
		$totals = is_array( $totals ) ? $totals : array();

		// New month, new counters: nobody wants to read a total that starts from the day they
		// installed the plugin, and keeping the history here would mean keeping it forever.
		if ( ! isset( $totals['month'] ) || $totals['month'] !== $month ) {
			$totals = array( 'month' => $month, 'count' => 0, 'input' => 0, 'output' => 0, 'cache_read' => 0 );
		}

		$totals['count']      = (int) $totals['count'] + 1;
		$totals['input']      = (int) $totals['input'] + self::tokens( $usage, 'input_tokens' );
		$totals['output']     = (int) $totals['output'] + self::tokens( $usage, 'output_tokens' );
		$totals['cache_read'] = (int) $totals['cache_read'] + self::tokens( $usage, 'cache_read_input_tokens' );

		// Not autoloaded: this is read on one admin screen, not on every page of the site.
		update_option( self::OPT_TOTALS, $totals, false );
	}

	protected static function remember( $answer, $mode, $usage ) {
		$recent = get_option( self::OPT_RECENT, array() );
		$recent = is_array( $recent ) ? $recent : array();

		array_unshift( $recent, array(
			'ts'    => time(),
			'mode'  => (string) $mode,
			'model' => isset( $answer['model'] ) ? (string) $answer['model'] : '',
			'req'   => isset( $answer['request_id'] ) ? (string) $answer['request_id'] : '',
			'in'    => self::tokens( $usage, 'input_tokens' ),
			'out'   => self::tokens( $usage, 'output_tokens' ),
			'user'  => get_current_user_id(),
		) );

		update_option( self::OPT_RECENT, array_slice( $recent, 0, self::KEEP ), false );
	}

	protected static function tokens( $usage, $key ) {
		return isset( $usage[ $key ] ) ? (int) $usage[ $key ] : 0;
	}
}
