<?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 {

	/**
	 * Singleton instance.
	 *
	 * @since 0.1.0
	 * @var WP_Agentic_Writer_Cost_Tracker
	 */
	private static $instance = null;

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

		return self::$instance;
	}

	/**
	 * Constructor.
	 *
	 * @since 0.1.0
	 */
	private function __construct() {
		// Hooks for tracking costs - accept 9 args (provider, session_id, status added).
		add_action( 'wp_aw_after_api_request', array( $this, 'add_request' ), 10, 9 );

		// Ensure table has new columns on first access.
		$this->maybe_upgrade_table();
	}

	/**
	 * Ensure table has latest schema with provider/session/status columns.
	 *
	 * @since 0.2.0
	 */
	private function maybe_upgrade_table() {
		global $wpdb;
		static $checked = false;

		if ( $checked ) {
			return;
		}

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

		// Check if table exists first.
		$table_exists = $wpdb->get_var( $wpdb->prepare( "SHOW TABLES LIKE %s", $table_name ) );
		if ( ! $table_exists ) {
			// Table missing - trigger recreation.
			if ( function_exists( 'wp_agentic_writer_create_cost_table' ) ) {
				wp_agentic_writer_create_cost_table();
			}
			$checked = true;
			return;
		}

		// Check if new columns exist.
		$columns = $wpdb->get_col( "DESCRIBE {$table_name}", 0 );

		$needs_provider = ! in_array( 'provider', $columns, true );
		$needs_session = ! in_array( 'session_id', $columns, true );
		$needs_status = ! in_array( 'status', $columns, true );

		if ( $needs_provider || $needs_session || $needs_status ) {
			$alter_parts = array();

			if ( $needs_provider ) {
				$alter_parts[] = "ADD COLUMN provider varchar(50) DEFAULT 'openrouter' AFTER action";
			}
			if ( $needs_session ) {
				$alter_parts[] = "ADD COLUMN session_id varchar(32) DEFAULT '' AFTER post_id";
			}
			if ( $needs_status ) {
				$alter_parts[] = "ADD COLUMN status varchar(20) DEFAULT 'success' AFTER cost";
			}

			if ( ! empty( $alter_parts ) ) {
				$wpdb->query( "ALTER TABLE {$table_name} " . implode( ', ', $alter_parts ) );
			}
		}

		$checked = true;
	}

	/**
	 * Add API request to cost tracking.
	 *
	 * @since 0.2.0 Parameters changed: added provider, session_id, status.
	 * @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.
	 * @param string $provider Provider name (optional, defaults to 'unknown').
	 * @param string $session_id Session ID (optional).
	 * @param string $status Request status (optional, defaults to 'success').
	 */
	public function add_request( $post_id, $model, $action, $input_tokens, $output_tokens, $cost, $provider = 'unknown', $session_id = '', $status = 'success' ) {
		global $wpdb;

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

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

	/**
	 * Legacy add_request for backward compatibility (4 params).
	 *
	 * @deprecated 0.2.0 Use add_request with all parameters.
	 * @param int    $post_id Post ID.
	 * @param string $model Model name.
	 * @param string $action Action type.
	 * @param int    $input_tokens Input tokens.
	 * @param int    $output_tokens Output tokens.
	 * @param float  $cost Cost in USD.
	 */
	public function add_request_legacy( $post_id, $model, $action, $input_tokens, $output_tokens, $cost ) {
		$this->add_request( $post_id, $model, $action, $input_tokens, $output_tokens, $cost );
	}

	/**
	 * Record usage from WP AI Client wrapper (legacy contract).
	 *
	 * This method provides backward compatibility for the WP AI Client wrapper
	 * and other callers that use a simpler interface.
	 *
	 * @deprecated 0.1.4 Use record_usage_full() instead for accurate provider attribution.
	 * @since 0.1.3
	 * @param int    $post_id Post ID.
	 * @param string $action Action/task type (e.g., 'chat', 'planning', 'writing').
	 * @param string $model Model identifier.
	 * @param float  $cost Cost in USD.
	 * @param string $session_id Session ID (optional).
	 */
	public function record_usage( $post_id, $action, $model, $cost, $session_id = '' ) {
		$this->record_usage_full(
			$post_id,
			$model,
			$action,
			0, // input_tokens - not available in wrapper
			0, // output_tokens - not available in wrapper
			$cost,
			'unknown', // deprecated wrapper - provider unknown
			$session_id,
			'success'
		);
	}

	/**
	 * Record usage with full metadata.
	 *
	 * Use this method when you have complete information about the request.
	 *
	 * @since 0.1.4
	 * @param int    $post_id Post ID.
	 * @param string $model Model identifier.
	 * @param string $action Action/task type.
	 * @param int    $input_tokens Input token count.
	 * @param int    $output_tokens Output token count.
	 * @param float  $cost Cost in USD.
	 * @param string $provider Provider name.
	 * @param string $session_id Session ID.
	 * @param string $status Request status.
	 */
	public function record_usage_full( $post_id, $model, $action, $input_tokens, $output_tokens, $cost, $provider, $session_id = '', $status = 'success' ) {
		$this->add_request(
			$post_id,
			$model,
			$action,
			$input_tokens,
			$output_tokens,
			$cost,
			$provider,
			$session_id,
			$status
		);
	}

	/**
	 * 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 ) + $this->get_image_variants_total_for_post( $post_id );
	}

	/**
	 * 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';
		$month_start = date( 'Y-m-01 00:00:00' );

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

		return floatval( $total ) + $this->get_image_variants_total_since( $month_start );
	}

	/**
	 * 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'] );
		}

		$image_cost = $this->get_image_variants_total_since( date( 'Y-m-d 00:00:00' ) );
		if ( $image_cost > 0 ) {
			$usage['image_generation'] = array(
				'tokens' => 0,
				'cost'   => $image_cost,
			);
			$total_cost += $image_cost;
		}

		$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
		);

		$history = $results ?: array();
		$image_history = $this->get_image_variants_history_for_post( $post_id, $limit );
		$history = array_merge( $history, $image_history );

		usort(
			$history,
			function( $a, $b ) {
				return strcmp( $b['created_at'] ?? '', $a['created_at'] ?? '' );
			}
		);

		return array_slice( $history, 0, $limit );
	}

	/**
	 * Check whether a table exists.
	 *
	 * @since 0.2.1
	 * @param string $table_name Table name.
	 * @return bool
	 */
	private function table_exists( $table_name ) {
		global $wpdb;
		return (bool) $wpdb->get_var( $wpdb->prepare( 'SHOW TABLES LIKE %s', $table_name ) );
	}

	/**
	 * Get image variant generation total for a post.
	 *
	 * @since 0.2.1
	 * @param int $post_id Post ID.
	 * @return float
	 */
	private function get_image_variants_total_for_post( $post_id ) {
		global $wpdb;

		$table_name = $wpdb->prefix . 'wpaw_images_variants';
		if ( $post_id <= 0 || ! $this->table_exists( $table_name ) ) {
			return 0.0;
		}

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

		return floatval( $total );
	}

	/**
	 * Get image variant generation total since a timestamp.
	 *
	 * @since 0.2.1
	 * @param string $since MySQL datetime.
	 * @return float
	 */
	private function get_image_variants_total_since( $since ) {
		global $wpdb;

		$table_name = $wpdb->prefix . 'wpaw_images_variants';
		if ( ! $this->table_exists( $table_name ) ) {
			return 0.0;
		}

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

		return floatval( $total );
	}

	/**
	 * Get image variant generation history for a post in cost-history shape.
	 *
	 * @since 0.2.1
	 * @param int $post_id Post ID.
	 * @param int $limit Limit.
	 * @return array
	 */
	private function get_image_variants_history_for_post( $post_id, $limit = 50 ) {
		global $wpdb;

		$table_name = $wpdb->prefix . 'wpaw_images_variants';
		if ( $post_id <= 0 || ! $this->table_exists( $table_name ) ) {
			return array();
		}

		$rows = $wpdb->get_results(
			$wpdb->prepare(
				"SELECT id, post_id, image_model_used, cost, generation_time, status, created_at
				FROM {$table_name}
				WHERE post_id = %d
				ORDER BY created_at DESC
				LIMIT %d",
				$post_id,
				$limit
			),
			ARRAY_A
		);

		return array_map(
			function( $row ) {
				return array(
					'id'            => 'image_variant_' . ( $row['id'] ?? '' ),
					'post_id'       => (int) ( $row['post_id'] ?? 0 ),
					'session_id'    => '',
					'model'         => $row['image_model_used'] ?? '',
					'provider'      => 'openrouter',
					'action'        => 'image_generation',
					'input_tokens'  => 0,
					'output_tokens' => 0,
					'cost'          => (float) ( $row['cost'] ?? 0 ),
					'status'        => $row['status'] ?? '',
					'created_at'    => $row['created_at'] ?? '',
				);
			},
			$rows ?: array()
		);
	}

	/**
	 * 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,
		);
	}
}