<?php
/**
 * Class Loco_Automatic_Translate_Addon_Pro\AI_Translate\Google\Google_AI_Model
 */

namespace Loco_Automatic_Translate_Addon_Pro\AI_Translate\Google;

use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Enums\Content_Role;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Types\Candidate;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Types\Candidates;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Types\Content;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Types\Generation_Config;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\API\Types\Parts\Text_Part;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Contracts\Generative_AI_Model;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Exception\Generative_AI_Exception;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Traits\With_Text_Generation_Trait;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Util\Formatter;
use Loco_Automatic_Translate_Addon_Pro\AI_Translate\Services\Util\Transformer;
use Generator;
use InvalidArgumentException;

/**
 * Class representing a Google AI model.
 */
class Google_AI_Model implements Generative_AI_Model {
	use With_Text_Generation_Trait;

	/**
	 * The Google AI API instance.
	 *
	 * @var Google_AI_API_Client
	 */
	private $api;

	/**
	 * The model slug.
	 *
	 * @var string
	 */
	private $model;

	/**
	 * The generation configuration.
	 *
	 * @var Generation_Config|null
	 */
	private $generation_config;

	/**
	 * The safety settings.
	 *
	 * @var Safety_Setting[]
	 */
	private $safety_settings;

	/**
	 * The system instruction.
	 *
	 * @var Content|null
	 */
	private $system_instruction;

	/**
	 * The request options.
	 *
	 * @var array<string, mixed>
	 */
	private $request_options;

	/**
	 * Constructor.
	 *
	 * @param Google_AI_API_Client $api             The Google AI API instance.
	 * @param string               $model           The model slug.
	 * @param array<string, mixed> $model_params    Optional. Additional model parameters. See
	 *                                              {@see Google_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(  $api, string $model, array $model_params = array(), array $request_options = array() ) {
		$this->api             = $api;
		$this->request_options = $request_options;

		if ( str_contains( $model, '/' ) ) {
			$this->model = $model;
		} else {
			$this->model = 'models/' . $model;
		}	
	}

	/**
	 * Gets the model slug.
	 *
	 * @return string The model slug.
	 */
	public function get_model_slug(): string {
		if ( str_starts_with( $this->model, 'models/' ) ) {
			return substr( $this->model, 7 );
		}
		return $this->model;
	}

	/**
	 * Sends a request to generate text content.
	 *
	 * @param Content[]            $contents        Prompts for the content to generate.
	 * @param array<string, mixed> $request_options The request options.
	 * @return Candidates The response candidates with generated text content - usually just one.
	 *
	 * @throws Generative_AI_Exception Thrown if the request fails or the response is invalid.
	 */
	protected function send_generate_text_request( array $contents, array $request_options ): Candidates {

		$params = $this->prepare_generate_text_params( $contents );
		$request  = $this->api->create_generate_content_request(
			$this->model,
			$params,
			array_merge(
				$this->request_options,
				$request_options
			)
		);

		$response = $this->api->make_request( $request );

		return $this->api->process_response_data(
			$response,
			function ( $response_data ) {
				return $this->get_response_candidates( $response_data );
			}
		);
	}


	/**
	 * Prepares the API request parameters for generating text content.
	 *
	 * @param Content[] $contents The contents to generate text for.
	 * @return array<string, mixed> The parameters for generating text content.
	 */
	private function prepare_generate_text_params( array $contents ): array {


		$transformers = self::get_content_transformers();




		$params = array(
			// TODO: Add support for tools and tool config, to support code generation.
			'contents' => array_map(
				static function ( Content $content ) use ( $transformers ) {
					return Transformer::transform_content( $content, $transformers );
				},
				$contents
			),
		);
		

		return array_filter( $params );
	}

	/**
	 * Extracts the candidates with content from the response.
	 *
	 * @param array<string, mixed> $response_data The response data.
	 * @param ?Candidates          $prev_chunk_candidates The candidates from the previous chunk in case of a streaming
	 *                                                    response, or null.
	 * @return Candidates The candidates with content parts.
	 *
	 * @throws Generative_AI_Exception Thrown if the response does not have any candidates with content.
	 */
	private function get_response_candidates( array $response_data, ?Candidates $prev_chunk_candidates = null ): Candidates {
		if ( ! isset( $response_data['candidates'] ) ) {
			throw $this->api->create_missing_response_key_exception( 'candidates' );
		}

		$this->check_non_empty_candidates( $response_data['candidates'] );

		if ( null === $prev_chunk_candidates ) {
			$other_data = $response_data;
			unset( $other_data['candidates'] );

			$candidates = new Candidates();
			foreach ( $response_data['candidates'] as $index => $candidate_data ) {
				$candidates->add_candidate(
					new Candidate(
						$this->prepare_candidate_content( $candidate_data, $index ),
						array_merge( $candidate_data, $other_data )
					)
				);
			}

			return $candidates;
		}

		// Subsequent chunk of a streaming response.
		$candidates_data = $this->merge_candidates_chunk(
			$prev_chunk_candidates->to_array(),
			$response_data
		);

		return Candidates::from_array( $candidates_data );
	}

	/**
	 * Merges a streaming response chunk with the previous candidates data.
	 *
	 * @param array<string, mixed> $candidates_data The candidates data from the previous chunk.
	 * @param array<string, mixed> $chunk_data      The response chunk data.
	 * @return array<string, mixed> The merged candidates data.
	 *
	 * @throws Generative_AI_Exception Thrown if the response is invalid.
	 */
	private function merge_candidates_chunk( array $candidates_data, array $chunk_data ): array {
		if ( ! isset( $chunk_data['candidates'] ) ) {
			throw $this->api->create_missing_response_key_exception( 'candidates' );
		}

		$other_data = $chunk_data;
		unset( $other_data['candidates'] );

		foreach ( $chunk_data['candidates'] as $index => $candidate_data ) {
			$candidates_data[ $index ] = array_merge( $candidates_data[ $index ], $candidate_data, $other_data );
		}

		return $candidates_data;
	}

	/**
	 * Transforms a given candidate from the API response into a Content instance.
	 *
	 * @param array<string, mixed> $candidate_data The API response candidate data.
	 * @param int                  $index          The index of the candidate in the response.
	 * @return Content The Content instance.
	 *
	 * @throws Generative_AI_Exception Thrown if the response is invalid.
	 */
	private function prepare_candidate_content( array $candidate_data, int $index ): Content {
		if ( ! isset( $candidate_data['content']['parts'] ) ) {
			// phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped
			throw $this->api->create_missing_response_key_exception( "candidates.{$index}.content.parts" );
		}

		$role = isset( $candidate_data['content']['role'] ) && 'user' === $candidate_data['content']['role']
			? Content_Role::USER
			: Content_Role::MODEL;

		return Content::from_array(
			array(
				'role'  => $role,
				'parts' => $candidate_data['content']['parts'],
			)
		);
	}

	/**
	 * Checks that the response includes candidates with content.
	 *
	 * @param array<string, mixed>[] $candidates_data The candidates data from the response.
	 *
	 * @throws Generative_AI_Exception Thrown if the response does not include any candidates with content.
	 */
	private function check_non_empty_candidates( array $candidates_data ): void {
		$errors = array();
		foreach ( $candidates_data as $candidate_data ) {
			if ( ! isset( $candidate_data['content'] ) ) {
				if ( isset( $candidate_data['finishReason'] ) ) {
					$errors[] = $candidate_data['finishReason'];
				} else {
					$errors[] = 'unknown';
				}
			}
		}

		if ( count( $errors ) === count( $candidates_data ) ) {
			$message = __( 'The response does not include any candidates with content.', 'ai-services' );

			$errors = array_unique(
				array_filter(
					$errors,
					static function ( $error ) {
						return 'unknown' !== $error;
					}
				)
			);
			if ( count( $errors ) > 0 ) {
				$message .= ' ' . sprintf(
					/* translators: %s: finish reason code */
					__( 'Finish reason: %s', 'ai-services' ),
					implode(
						wp_get_list_item_separator(),
						$errors
					)
				);
			}

			// phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped
			throw $this->api->create_response_exception( $message );
		}
	}

	/**
	 * Gets the content transformers.
	 *
	 * @return array<string, callable> The content transformers.
	 */
	private static function get_content_transformers(): array {
		return array(
			'role'  => static function ( Content $content ) {
				return $content->get_role();
			},
			'parts' => static function ( Content $content ) {

				$parts = array();
				$contentparts = (array) ($content->get_parts());

				foreach ( $contentparts as $partData ) {
					foreach ( $partData as $part ) {
						if ( $part instanceof Text_Part ) {
							$parts[] = array( 'text' => $part->get_text() );
					    } else {
						throw new Generative_AI_Exception(
							esc_html__( 'The Google AI API only supports text, image, and audio parts.', 'ai-services' )
						);
					}
				  }
			    }
				return $parts;
			},
		);
	}

	/**
	 * Gets the generation configuration transformers.
	 *
	 * @return array<string, callable> The generation configuration transformers.
	 */
	private static function get_generation_config_transformers(): array {
		return array(
			'stopSequences'    => static function ( Generation_Config $config ) {
				return $config->get_stop_sequences();
			},
			'responseMimeType' => static function ( Generation_Config $config ) {
				return $config->get_response_mime_type();
			},
			'responseSchema'   => static function ( Generation_Config $config ) {
				if ( $config->get_response_mime_type() === 'application/json' ) {
					return $config->get_response_schema();
				}
				return array();
			},
			'candidateCount'   => static function ( Generation_Config $config ) {
				return $config->get_candidate_count();
			},
			'maxOutputTokens'  => static function ( Generation_Config $config ) {
				return $config->get_max_output_tokens();
			},
			'temperature'      => static function ( Generation_Config $config ) {
				// In the Google AI API temperature ranges from 0.0 to 2.0.
				return $config->get_temperature() * 2.0;
			},
			'topP'             => static function ( Generation_Config $config ) {
				return $config->get_top_p();
			},
			'topK'             => static function ( Generation_Config $config ) {
				return $config->get_top_k();
			},
			'presencePenalty'  => static function ( Generation_Config $config ) {
				return $config->get_presence_penalty();
			},
			'frequencyPenalty' => static function ( Generation_Config $config ) {
				return $config->get_frequency_penalty();
			},
			'responseLogprobs' => static function ( Generation_Config $config ) {
				return $config->get_response_logprobs();
			},
			'logprobs'         => static function ( Generation_Config $config ) {
				return $config->get_logprobs();
			},
		);
	}
}
