<?php
/**
 * Class Felix_Arntz\AI_Services\OpenAI\OpenAI_AI_Text_Generation_Model
 *
 * @since 0.1.0
 * @package ai-services
 */

namespace Felix_Arntz\AI_Services\OpenAI;

use Felix_Arntz\AI_Services\Services\API\Helpers;
use Felix_Arntz\AI_Services\Services\API\Types\Content;
use Felix_Arntz\AI_Services\Services\API\Types\Contracts\Tool;
use Felix_Arntz\AI_Services\Services\API\Types\Model_Metadata;
use Felix_Arntz\AI_Services\Services\API\Types\Parts;
use Felix_Arntz\AI_Services\Services\API\Types\Text_Generation_Config;
use Felix_Arntz\AI_Services\Services\Base\OpenAI_Compatible_AI_Text_Generation_Model;
use Felix_Arntz\AI_Services\Services\Contracts\Generative_AI_API_Client;
use Felix_Arntz\AI_Services\Services\Contracts\With_Function_Calling;
use Felix_Arntz\AI_Services\Services\Contracts\With_Multimodal_Input;
use Felix_Arntz\AI_Services\Services\Contracts\With_Multimodal_Output;
use Felix_Arntz\AI_Services\Services\Exception\Generative_AI_Exception;
use Felix_Arntz\AI_Services\Services\Traits\OpenAI_Compatible_Text_Generation_With_Function_Calling_Trait;
use InvalidArgumentException;

/**
 * Class representing an OpenAI text generation AI model.
 *
 * @since 0.1.0
 * @since 0.5.0 Renamed from `OpenAI_AI_Model`.
 * @since 0.7.0 Now extends `OpenAI_Compatible_AI_Text_Generation_Model` instead of `Abstract_AI_Model`.
 */
class OpenAI_AI_Text_Generation_Model extends OpenAI_Compatible_AI_Text_Generation_Model implements With_Function_Calling, With_Multimodal_Input, With_Multimodal_Output {
	use OpenAI_Compatible_Text_Generation_With_Function_Calling_Trait {
		prepare_generate_text_params as prepare_generate_text_params_with_function_calling;
		prepare_tool as prepare_function_declarations_tool;
	}

	/**
	 * The expected MIME type of any audio output generated by the model.
	 *
	 * Internal temporary storage to not have to pass it around, as it should not be part of the interface.
	 *
	 * @since 0.7.0
	 * @var string
	 */
	private $expected_audio_mime_type = 'audio/mpeg';

	/**
	 * Constructor.
	 *
	 * @since 0.1.0
	 *
	 * @param Generative_AI_API_Client $api_client      The AI API client instance.
	 * @param Model_Metadata           $metadata        The model metadata.
	 * @param array<string, mixed>     $model_params    Optional. Additional model parameters. See
	 *                                                  {@see OpenAI_AI_Service::get_model()} for the list of available
	 *                                                  parameters. Default empty array.
	 * @param array<string, mixed>     $request_options Optional. The request options. Default empty array.
	 *
	 * @throws InvalidArgumentException Thrown if the model parameters are invalid.
	 */
	public function __construct( Generative_AI_API_Client $api_client, Model_Metadata $metadata, array $model_params = array(), array $request_options = array() ) {
		parent::__construct( $api_client, $metadata, $model_params, $request_options );

		$this->set_tool_config_from_model_params( $model_params );
		$this->set_tools_from_model_params( $model_params );
	}

	/**
	 * Prepares the API request parameters for generating text content.
	 *
	 * @since 0.7.0
	 *
	 * @param Content[] $contents The contents to generate text for.
	 * @return array<string, mixed> The parameters for generating text content.
	 *
	 * @throws InvalidArgumentException Thrown if an invalid tool is provided.
	 */
	protected function prepare_generate_text_params( array $contents ): array {
		$params = $this->prepare_generate_text_params_with_function_calling( $contents );

		// If 'audio' output is requested, the OpenAI API requires the 'audio' parameter to be set.
		if (
			isset( $params['modalities'] ) &&
			is_array( $params['modalities'] ) &&
			in_array( 'audio', $params['modalities'], true ) &&
			! isset( $params['audio'] )
		) {
			$params['audio'] = array(
				'voice'  => 'alloy',
				'format' => 'mp3',
			);
		}

		if ( isset( $params['audio']['format'] ) ) {
			// Hack: Store the expected MIME type for audio output, as the OpenAI API does not return it.
			$this->expected_audio_mime_type = 'mp3' === $params['audio']['format'] ? 'audio/mpeg' : 'audio/' . $params['audio']['format'];
		}

		return $params;
	}

	/**
	 * Prepares a single tool for the API request, amending the provided parameters as needed.
	 *
	 * @since 0.7.0
	 *
	 * @param array<string, mixed> $params The parameters to prepare the tools for. Passed by reference.
	 * @param Tool                 $tool   The tool to prepare.
	 * @return bool True if the tool was successfully prepared, false otherwise.
	 */
	protected function prepare_tool( array &$params, Tool $tool ): bool {
		$result = $this->prepare_function_declarations_tool( $params, $tool );
		if ( ! $result ) {
			return false;
		}

		/*
		 * The OpenAI API supports a 'strict' argument for function tools, which is not part of the standard OpenAI API
		 * specification and therefore may not be supported by other providers.
		 * Since it makes sense to always use it for OpenAI, we add it here if not set.
		 */
		if ( isset( $params['tools'] ) && is_array( $params['tools'] ) ) {
			$params['tools'] = array_map(
				function ( $openai_tool_data ) {
					if ( ! isset( $openai_tool_data['type'] ) || 'function' !== $openai_tool_data['type'] ) {
						return $openai_tool_data;
					}
					// Add the 'strict' argument to the function tool if not set.
					if ( ! isset( $openai_tool_data['function']['strict'] ) ) {
						$openai_tool_data['function']['strict'] = true;
					}
					return $openai_tool_data;
				},
				$params['tools']
			);
		}

		return true;
	}

	/**
	 * Transforms a given candidate from the API response into a Parts instance.
	 *
	 * @since 0.7.0
	 *
	 * @param array<string, mixed> $candidate_data The API response candidate data.
	 * @return Parts The Parts instance.
	 *
	 * @throws Generative_AI_Exception Thrown if the response is invalid.
	 */
	protected function prepare_response_candidate_content_parts( array $candidate_data ): Parts {
		$parts = parent::prepare_response_candidate_content_parts( $candidate_data );

		if ( isset( $candidate_data['message']['audio']['data'] ) && is_string( $candidate_data['message']['audio']['data'] ) ) {
			// Add the audio data as a part.
			$mime_type = $this->expected_audio_mime_type;
			$parts->add_inline_data_part(
				$mime_type,
				Helpers::base64_data_to_base64_data_url( $candidate_data['message']['audio']['data'], $mime_type )
			);
		}

		return $parts;
	}

	/**
	 * Gets the generation configuration transformers.
	 *
	 * @since 0.7.0
	 *
	 * @return array<string, callable> The generation configuration transformers.
	 */
	protected function get_generation_config_transformers(): array {
		$transformers = parent::get_generation_config_transformers();

		// Support multimodal output (e.g. for speech generation).
		$transformers['modalities'] = static function ( Text_Generation_Config $config ) {
			return $config->get_output_modalities();
		};

		return $transformers;
	}
}