wp-agentic-writer/includes/class-cost-tracker.php

<?php
/**
 * Cost Tracker
 *
 * Tracks and displays API costs for user sessions.
 *
 * @package WP_Agentic_Writer
 */

if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

/**
 * Class WP_Agentic_Writer_Cost_Tracker
 *
 * @since 0.1.0
 */
class WP_Agentic_Writer_Cost_Tracker {

	/**
	 * Get singleton instance.
	 *
	 * @since 0.1.0
	 * @return WP_Agentic_Writer_Cost_Tracker
	 */
	public static function get_instance() {
		static $instance = null;

		if ( null === $instance ) {
			$instance = new self();
		}

		return $instance;
	}

	/**
	 * Constructor.
	 *
	 * @since 0.1.0
	 */
	private function __construct() {
		// Hooks for tracking costs.
		add_action( 'wp_aw_after_api_request', array( $this, 'add_request' ), 10, 6 );
	}

	/**
	 * Add API request to cost tracking.
	 *
	 * @since 0.1.0
	 * @param int    $post_id Post ID.
	 * @param string $model Model name.
	 * @param string $action Action type (planning, execution, research, image).
	 * @param int    $input_tokens Input tokens.
	 * @param int    $output_tokens Output tokens.
	 * @param float  $cost Cost in USD.
	 */
	public function add_request( $post_id, $model, $action, $input_tokens, $output_tokens, $cost ) {
		global $wpdb;

		$table_name = $wpdb->prefix . 'wpaw_cost_tracking';

		$wpdb->insert(
			$table_name,
			array(
				'post_id' => $post_id,
				'model' => $model,
				'action' => $action,
				'input_tokens' => $input_tokens,
				'output_tokens' => $output_tokens,
				'cost' => $cost,
				'created_at' => current_time( 'mysql' ),
			),
			array( '%d', '%s', '%s', '%d', '%d', '%f', '%s' )
		);
	}

	/**
	 * Get session total cost.
	 *
	 * @since 0.1.0
	 * @param int $post_id Post ID.
	 * @return float Session total cost.
	 */
	public function get_session_total( $post_id ) {
		global $wpdb;

		$table_name = $wpdb->prefix . 'wpaw_cost_tracking';

		$total = $wpdb->get_var(
			$wpdb->prepare(
				"SELECT SUM(cost) FROM {$table_name} WHERE post_id = %d",
				$post_id
			)
		);

		return floatval( $total );
	}

	/**
	 * Get monthly total cost.
	 *
	 * @since 0.1.0
	 * @return float Monthly total cost.
	 */
	public function get_monthly_total() {
		global $wpdb;

		$table_name = $wpdb->prefix . 'wpaw_cost_tracking';

		$total = $wpdb->get_var(
			$wpdb->prepare(
				"SELECT SUM(cost) FROM {$table_name} WHERE created_at >= %s",
				date( 'Y-m-01 00:00:00' )
			)
		);

		return floatval( $total );
	}

	/**
	 * Get today's usage breakdown.
	 *
	 * @since 0.1.0
	 * @return array Today's usage.
	 */
	public function get_today_usage() {
		global $wpdb;

		$table_name = $wpdb->prefix . 'wpaw_cost_tracking';

		$results = $wpdb->get_results(
			$wpdb->prepare(
				"SELECT action, SUM(input_tokens + output_tokens) as tokens, SUM(cost) as cost
				FROM {$table_name}
				WHERE created_at >= %s
				GROUP BY action",
				date( 'Y-m-d 00:00:00' )
			),
			ARRAY_A
		);

		$usage = array();
		$total_cost = 0;
		$total_tokens = 0;

		foreach ( $results as $row ) {
			$usage[ $row['action'] ] = array(
				'tokens' => intval( $row['tokens'] ),
				'cost' => floatval( $row['cost'] ),
			);
			$total_cost += floatval( $row['cost'] );
			$total_tokens += intval( $row['tokens'] );
		}

		$usage['total'] = array(
			'cost' => $total_cost,
			'tokens' => $total_tokens,
		);

		return $usage;
	}

	/**
	 * Format cost for display.
	 *
	 * @since 0.1.0
	 * @param float $cost Cost in USD.
	 * @return string Formatted cost.
	 */
	public function format_cost( $cost ) {
		return '$' . number_format( $cost, 4, '.', ',' );
	}

	/**
	 * Format tokens for display.
	 *
	 * @since 0.1.0
	 * @param int $tokens Number of tokens.
	 * @return string Formatted tokens.
	 */
	public function format_tokens( $tokens ) {
		return number_format( $tokens ) . ' tokens';
	}

	/**
	 * Get detailed cost history for a post.
	 *
	 * @since 0.1.0
	 * @param int $post_id Post ID.
	 * @param int $limit Number of records to return.
	 * @return array Cost history records.
	 */
	public function get_post_history( $post_id, $limit = 50 ) {
		global $wpdb;

		$table_name = $wpdb->prefix . 'wpaw_cost_tracking';

		$results = $wpdb->get_results(
			$wpdb->prepare(
				"SELECT * FROM {$table_name} WHERE post_id = %d ORDER BY created_at DESC LIMIT %d",
				$post_id,
				$limit
			),
			ARRAY_A
		);

		return $results;
	}

	/**
	 * Get cost tracking data for frontend.
	 *
	 * @since 0.1.0
	 * @param int $post_id Post ID.
	 * @return array Cost tracking data.
	 */
	public function get_frontend_data( $post_id = 0 ) {
		$today_usage = $this->get_today_usage();
		$monthly_total = $this->get_monthly_total();
		$session_total = $post_id > 0 ? $this->get_session_total( $post_id ) : 0;
		$post_history = $post_id > 0 ? $this->get_post_history( $post_id ) : array();

		// Get settings for budget.
		$settings = get_option( 'wp_agentic_writer_settings', array() );
		$monthly_budget = $settings['monthly_budget'] ?? 600;

		return array(
			'today' => $today_usage,
			'monthly' => array(
				'used' => $monthly_total,
				'budget' => $monthly_budget,
				'percentage' => $monthly_budget > 0 ? ( $monthly_total / $monthly_budget ) * 100 : 0,
				'remaining' => max( 0, $monthly_budget - $monthly_total ),
			),
			'session' => $session_total,
			'history' => $post_history,
		);
	}
}