<?php
/**
 * Class Loco_Automatic_Translate_Addon_Pro\AI_Translate\OpenAI\OpenAI_AI_Service
 */

namespace Loco_Automatic_Translate_Addon_Pro\AI_Translate\OpenAI;

use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Enums\AI_Capability;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Types\Content;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Types\Parts;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Contracts\Authentication;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Contracts\Generative_AI_Model;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Contracts\Generative_AI_Service;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Exception\Generative_AI_Exception;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Util\AI_Capabilities;
use Felix_Arntz\WP_OOP_Plugin_Lib\HTTP\HTTP;
use InvalidArgumentException;

/**
 * Class for the OpenAI AI service.
 */
class OpenAI_AI_Service implements Generative_AI_Service {

	/**
	 * The OpenAI AI API instance.
	 */
	private $api;

	/**
	 * Constructor.
	 *
	 * @param Authentication $authentication The authentication credentials.
	 * @param HTTP           $http           Optional. The HTTP instance to use for requests. Default is a new instance.
	 */
	public function __construct( Authentication $authentication, HTTP $http = null ) {
		if ( ! $http ) {
			$http = new HTTP();
		}
		$this->api = new OpenAI_AI_API_Client( $authentication, $http );
	}

	/**
	 * Gets the service slug.
	 *
	 * @return string The service slug.
	 */
	public function get_service_slug(): string {
		return 'openai';
	}

	/**
	 * Gets the list of AI capabilities that the service and its models support.
	 *
	 * @see AI_Capabilities
	 * @return string[] The list of AI capabilities.
	 */
	public function get_capabilities(): array {
		return AI_Capabilities::get_model_class_capabilities( OpenAI_AI_Model::class );
	}

	/**
	 * Checks whether the service is connected.
	 *
	 * This is typically used to check whether the current service credentials are valid.
	 *
	 * @return bool True if the service is connected, false otherwise.
	 */
	public function is_connected(): bool {
		try {
			$this->list_models();
			return true;
		} catch ( Generative_AI_Exception $e ) {
			return false;
		}
	}

	/**
	 * Lists the available generative model slugs and their capabilities.
	 *
	 * @param array<string, mixed> $request_options Optional. The request options. Default empty array.
	 * @return array<string, string[]> Map of the available model slugs and their capabilities.
	 *
	 * @throws Generative_AI_Exception Thrown if the request fails or the response is invalid.
	 */
	public function list_models( array $request_options = array() ): array {
		$request       = $this->api->create_list_models_request();
		$response_data = $this->api->make_request( $request )->get_data();

		if ( ! isset( $response_data['data'] ) || ! $response_data['data'] ) {
			throw new Generative_AI_Exception(
				esc_html(
					sprintf(
						/* translators: %s: key name */
						__( 'The response from the OpenAI API is missing the "%s" key.', 'ai-services' ),
						'data'
					)
				)
			);
		}

		$model_slugs = array_map(
			static function ( array $model ) {
				return $model['id'];
			},
			$response_data['data']
		);

		// Unfortunately, the OpenAI API does not return model capabilities, so we have to hardcode them here.
		$gpt_capabilities = array(
			AI_Capability::TEXT_GENERATION,
		);
		$gpt_multimodal_capabilities = array(
			AI_Capability::MULTIMODAL_INPUT,
			AI_Capability::TEXT_GENERATION,
		);

		return array_reduce(
			$model_slugs,
			static function ( array $model_caps, string $model_slug ) use ( $gpt_capabilities, $gpt_multimodal_capabilities ) {
				if (
					str_starts_with( $model_slug, 'gpt-' )
					&& ! str_contains( $model_slug, '-instruct' )
					&& ! str_contains( $model_slug, '-realtime' )
					&& ! str_contains( $model_slug, '-audio' )
				) {
					if ( str_starts_with( $model_slug, 'gpt-4o' ) ) {
						$model_caps[ $model_slug ] = $gpt_multimodal_capabilities;
					} else {
						$model_caps[ $model_slug ] = $gpt_capabilities;
					}
				} else {
					$model_caps[ $model_slug ] = array();
				}
				return $model_caps;
			},
			array()
		);
	}

	/**
	 * Gets a generative model instance for the provided model parameters.
	 *
	 * @param array<string, mixed> $model_params    {
	 *     Optional. Model parameters. Default empty array.
	 *
	 *     @type string               $feature           Required. Unique identifier of the feature that the model
	 *                                                   will be used for. Must only contain lowercase letters,
	 *                                                   numbers, hyphens.
	 *     @type string               $model             The model slug. By default, the model will be determined
	 *                                                   based on heuristics such as the requested capabilities.
	 *     @type string[]             $capabilities      Capabilities requested for the model to support. It is
	 *                                                   recommended to specify this if you do not explicitly specify a
	 *                                                   model slug.
	 *     @type Generation_Config?   $generationConfig  Model generation configuration options. Default none.
	 *     @type string|Parts|Content $systemInstruction The system instruction for the model. Default none.
	 * }
	 * @param array<string, mixed> $request_options Optional. The request options. Default empty array.
	 * @return Generative_AI_Model The generative model.
	 *
	 * @throws InvalidArgumentException Thrown if the model slug or parameters are invalid.
	 * @throws Generative_AI_Exception Thrown if getting the model fails.
	 */
	public function get_model( array $model_params = array(), array $request_options = array() ): Generative_AI_Model {
		if ( isset( $model_params['model'] ) ) {
			$model = $model_params['model'];
			unset( $model_params['model'] );
		} else {
			if ( isset( $model_params['capabilities'] ) ) {
				$model_slugs = AI_Capabilities::get_model_slugs_for_capabilities(
					$this->list_models( $request_options ),
					$model_params['capabilities']
				);
			} else {
				$model_slugs = array_keys( $this->list_models( $request_options ) );
			}
			$model = $this->sort_models_by_preference( $model_slugs )[0];
		}

		return new OpenAI_AI_Model( $this->api, $model, $model_params, $request_options );
	}

	/**
	 * Sorts model slugs by preference.
	 *
	 * @param string[] $model_slugs The model slugs to sort.
	 * @return string[] The model slugs, sorted by preference.
	 */
	private function sort_models_by_preference( array $model_slugs ): array {
		$get_preference_group = static function ( $model_slug ) {
			if ( str_starts_with( $model_slug, 'gpt-3.5' ) ) {
				if ( str_ends_with( $model_slug, '-turbo' ) ) {
					return 0;
				}
				if ( str_contains( $model_slug, '-turbo' ) ) {
					return 1;
				}
				return 2;
			}
			if ( str_starts_with( $model_slug, 'gpt-' ) ) {
				if ( str_contains( $model_slug, '-turbo' ) ) {
					return 3;
				}
				return 4;
			}
			return 5;
		};

		$preference_groups = array_fill( 0, 6, array() );
		foreach ( $model_slugs as $model_slug ) {
			$group                         = $get_preference_group( $model_slug );
			$preference_groups[ $group ][] = $model_slug;
		}

		return array_merge( ...$preference_groups );
	}
}
