<?php

namespace Limb_Chatbot\Includes\AI_Providers\Open_Router\Services;

use Exception;
use Limb_Chatbot\Includes\AI_Providers\Open_Router\Open_Router;
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 variant types for the OpenRouter provider.
 *
 * OpenRouter supports model variants:
 * - `:free` - Free variant with no cost
 * - `:extended` - Extended context window
 * - `:exacto` - Higher tool-calling accuracy
 * - `:thinking` - Extended reasoning capabilities
 * - `:online` - Real-time web search
 * - `:nitro` - High-speed inference
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Open_Router\Services
 * @since 1.0.14
 */
class AI_Model_Cost_Service {

	/**
	 * Free variant suffix.
	 *
	 * @since 1.0.14
	 * @var string
	 */
	const VARIANT_FREE = ':free';

	/**
	 * Extended context variant suffix.
	 *
	 * @since 1.0.14
	 * @var string
	 */
	const VARIANT_EXTENDED = ':extended';

	/**
	 * Exacto (tool-calling accuracy) variant suffix.
	 *
	 * @since 1.0.14
	 * @var string
	 */
	const VARIANT_EXACTO = ':exacto';

	/**
	 * Thinking (reasoning) variant suffix.
	 *
	 * @since 1.0.14
	 * @var string
	 */
	const VARIANT_THINKING = ':thinking';

	/**
	 * Online (web search) variant suffix.
	 *
	 * @since 1.0.14
	 * @var string
	 */
	const VARIANT_ONLINE = ':online';

	/**
	 * Nitro (high-speed) variant suffix.
	 *
	 * @since 1.0.14
	 * @var string
	 */
	const VARIANT_NITRO = ':nitro';

	/**
	 * All available variant suffixes.
	 *
	 * @since 1.0.14
	 * @var array
	 */
	const VARIANTS = [
		self::VARIANT_FREE,
		self::VARIANT_EXTENDED,
		self::VARIANT_EXACTO,
		self::VARIANT_THINKING,
		self::VARIANT_ONLINE,
		self::VARIANT_NITRO,
	];

	/**
	 * Checks if the model is a free variant and updates 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.14
	 */
	public static function check_model_variant( Token_Usage $usage, AI_Model $model, Chat $chat ): void {
		if ( $model->get_ai_provider_id() !== Open_Router::$id ) {
			return;
		}

		if ( self::is_free_variant( $model ) ) {
			self::apply_free_variant_costs( $model );
		} else {
			self::restore_paid_costs( $model );
		}
	}

	/**
	 * Checks if a model is using the free variant.
	 *
	 * @param AI_Model $model AI model instance.
	 *
	 * @return bool True if model uses :free variant.
	 * @since 1.0.14
	 */
	public static function is_free_variant( AI_Model $model ): bool {
		$model_name = $model->get_name();

		return self::has_variant( $model_name, self::VARIANT_FREE );
	}

	/**
	 * Checks if a model name has a specific variant suffix.
	 *
	 * @param string $model_name Model name/ID.
	 * @param string $variant    Variant suffix to check.
	 *
	 * @return bool True if model has the variant.
	 * @since 1.0.14
	 */
	public static function has_variant( string $model_name, string $variant ): bool {
		return str_ends_with( $model_name, $variant );
	}

	/**
	 * Gets the variant from a model name.
	 *
	 * @param string $model_name Model name/ID.
	 *
	 * @return string|null The variant suffix or null if no variant.
	 * @since 1.0.14
	 */
	public static function get_variant( string $model_name ): ?string {
		foreach ( self::VARIANTS as $variant ) {
			if ( self::has_variant( $model_name, $variant ) ) {
				return $variant;
			}
		}

		return null;
	}

	/**
	 * Gets the base model name without variant suffix.
	 *
	 * @param string $model_name Model name/ID with potential variant.
	 *
	 * @return string Base model name.
	 * @since 1.0.14
	 */
	public static function get_base_model_name( string $model_name ): string {
		foreach ( self::VARIANTS as $variant ) {
			if ( self::has_variant( $model_name, $variant ) ) {
				return substr( $model_name, 0, -strlen( $variant ) );
			}
		}

		return $model_name;
	}

	/**
	 * Applies the :free variant to a model name.
	 *
	 * @param string $model_name Base model name.
	 *
	 * @return string Model name with :free variant.
	 * @since 1.0.14
	 */
	public static function apply_free_variant( string $model_name ): string {
		$base_name = self::get_base_model_name( $model_name );

		return $base_name . self::VARIANT_FREE;
	}

	/**
	 * Applies zero costs to a model (for free variants).
	 *
	 * @param AI_Model|null $model AI model instance.
	 *
	 * @return void
	 * @throws Exception
	 * @since 1.0.14
	 */
	public static function apply_free_variant_costs( ?AI_Model $model ): void {
		if ( $model === null ) {
			return;
		}

		// Store original costs before zeroing (if not already stored)
		$input_cost  = $model->get_meta( 'input_token_cost' );
		$output_cost = $model->get_meta( 'output_token_cost' );

		if ( $input_cost > 0 && ! $model->get_meta( 'paid_input_token_cost' ) ) {
			$model->update_meta( 'paid_input_token_cost', $input_cost );
		}

		if ( $output_cost > 0 && ! $model->get_meta( 'paid_output_token_cost' ) ) {
			$model->update_meta( 'paid_output_token_cost', $output_cost );
		}

		// Set costs to zero for free variant
		$model->update_meta( 'input_token_cost', 0 );
		$model->update_meta( 'output_token_cost', 0 );
	}

	/**
	 * Restores paid costs to a model (for non-free variants).
	 *
	 * @param AI_Model|null $model AI model instance.
	 *
	 * @return void
	 * @throws Exception
	 * @since 1.0.14
	 */
	public static function restore_paid_costs( ?AI_Model $model ): void {
		if ( $model === null ) {
			return;
		}

		// Only restore if current cost is 0 and we have stored paid costs
		$current_input_cost = $model->get_meta( 'input_token_cost' );

		if ( $current_input_cost == 0 ) {
			if ( $paid_input = $model->get_meta( 'paid_input_token_cost' ) ) {
				$model->update_meta( 'input_token_cost', $paid_input );
			}

			if ( $paid_output = $model->get_meta( 'paid_output_token_cost' ) ) {
				$model->update_meta( 'output_token_cost', $paid_output );
			}
		}
	}

	/**
	 * Handles model cost updates before checking usage limits.
	 *
	 * @param Collection $limits Collection of usage limits.
	 *
	 * @return void
	 * @throws Exception
	 * @since 1.0.14
	 */
	public static function before_limits_check( Collection $limits ): void {
		$chatbot = $limits->first()->get_chatbot();

		if ( $chatbot->get_ai_provider_id() !== Open_Router::$id ) {
			return;
		}

		$model = $chatbot->get_ai_model();

		if ( $model === null ) {
			return;
		}

		if ( self::is_free_variant( $model ) ) {
			self::apply_free_variant_costs( $model );
		} else {
			self::restore_paid_costs( $model );
		}
	}

	/**
	 * Checks if a model has a free variant available.
	 *
	 * This can be used to offer users the option to switch to free variant.
	 *
	 * @param AI_Model $model AI model instance.
	 *
	 * @return bool True if free variant exists.
	 * @since 1.0.14
	 */
	public static function has_free_variant_available( AI_Model $model ): bool {
		$base_name       = self::get_base_model_name( $model->get_name() );
		$free_model_name = $base_name . self::VARIANT_FREE;

		return AI_Model::find_by_name( $free_model_name ) !== null;
	}

	/**
	 * Gets the free variant model for a given model.
	 *
	 * @param AI_Model $model AI model instance.
	 *
	 * @return AI_Model|null Free variant model or null if not found.
	 * @since 1.0.14
	 */
	public static function get_free_variant_model( AI_Model $model ): ?AI_Model {
		$base_name       = self::get_base_model_name( $model->get_name() );
		$free_model_name = $base_name . self::VARIANT_FREE;

		return AI_Model::find_by_name( $free_model_name );
	}

	/**
	 * Gets the paid (base) variant model for a given model.
	 *
	 * @param AI_Model $model AI model instance (potentially with variant).
	 *
	 * @return AI_Model|null Base model or null if not found.
	 * @since 1.0.14
	 */
	public static function get_paid_variant_model( AI_Model $model ): ?AI_Model {
		$base_name = self::get_base_model_name( $model->get_name() );

		if ( $base_name === $model->get_name() ) {
			return $model;
		}

		return AI_Model::find_by_name( $base_name );
	}
}
