<?php

namespace Limb_Chatbot\Includes\AI_Providers\Gemini\Services;

use Exception;
use Limb_Chatbot\Includes\AI_Providers\Gemini\Gemini;
use Limb_Chatbot\Includes\Data_Objects\AI_Model;
use Limb_Chatbot\Includes\Data_Objects\Chat;
use Limb_Chatbot\Includes\Data_Objects\Token_Usage;
use Limb_Chatbot\Includes\Services\Collection;

/**
 * Class AI_Model_Cost_Service
 *
 * Handles cost management logic for AI models based on user tiers for the Gemini provider.
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Gemini\Services
 * @since 1.0.0
 */
class AI_Model_Cost_Service {

	/**
	 * Checks user tier and updates model costs accordingly.
	 *
	 * @param  Token_Usage  $usage  Token usage data.
	 * @param  AI_Model  $model  AI model instance.
	 * @param  Chat  $chat  Chat instance.
	 *
	 * @return void
	 * @throws Exception
	 * @since 1.0.0
	 */
	public static function check_user_tier( Token_Usage $usage, AI_Model $model, Chat $chat ): void {
		$chatbot = $chat->get_chatbot();
		$params  = $chatbot->get_config()->get_params();
		if ( $model->get_ai_provider_id() == Gemini::$id ) {
			if ( empty( $params['paid_tier'] ) ) {
				self::maybe_remove_model_costs( $model );
			} else {
				self::maybe_add_model_costs( $model );
			}
		}
	}

	/**
	 * Removes token costs from the model.
	 *
	 * @param  AI_Model|null  $model  AI model instance.
	 *
	 * @return void
	 * @throws Exception
	 * @since 1.0.0
	 */
	public static function maybe_remove_model_costs( ?AI_Model $model ): void {
		if ( $model == null ) {
			return;
		}
		if ( ! $model->get_meta( 'input_token_cost' ) ) {
			return;
		}
		$model->update_meta( 'input_token_cost', 0 );
		$model->update_meta( 'output_token_cost', 0 );
	}

	/**
	 * Adds token costs to the model if not already set.
	 *
	 * @param  AI_Model|null  $model  AI model instance.
	 *
	 * @return void
	 * @throws Exception
	 * @since 1.0.0
	 */
	public static function maybe_add_model_costs( ?AI_Model $model ): void {
		if ( $model == null ) {
			return;
		}
		if ( ( $cost = $model->get_meta( 'input_token_cost' ) ) && $cost > 0 ) {
			return;
		}
		if ( $input_token_cost = $model->get_meta( 'paid_input_token_cost' ) ) {
			$model->update_meta( 'input_token_cost', $input_token_cost );
		}
		if ( $output_token_cost = $model->get_meta( 'paid_output_token_cost' ) ) {
			$model->update_meta( 'output_token_cost', $output_token_cost );
		}
	}

	/**
	 * Handles model cost updates before checking usage limits.
	 *
	 * @param  Collection  $limits  Collection of usage limits.
	 *
	 * @return void
	 * @throws Exception
	 * @since 1.0.0
	 */
	public static function before_limits_check( Collection $limits ): void {
		$chatbot = $limits->first()->get_chatbot();
		if ( $limits->first()->get_chatbot()->get_ai_provider_id() == Gemini::$id ) {
			if ( $config = $chatbot->get_config() ) {
				if ( empty( $config->get_params()['paid_tier'] ) ) {
					self::maybe_remove_model_costs( $limits->first()->get_chatbot()->get_ai_model() );
				} else {
					self::maybe_add_model_costs( $limits->first()->get_chatbot()->get_ai_model() );
				}
			}
		}
	}
}