<?php
/**
 * Google AI API client for Nano Banana image generation.
 *
 * @package 45Squared_AI_Image_Generator
 */

if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

/**
 * Client for the Google Gemini image generation and editing API.
 */
class Nano_Banana_API {

	const API_BASE = 'https://generativelanguage.googleapis.com/v1beta/models/';

	const ALLOWED_MIME_TYPES = array( 'image/png', 'image/jpeg', 'image/webp' );

	/**
	 * Resolve a model ID to a known-good value.
	 *
	 * The model is interpolated into the request URL path, so it must never
	 * be taken from user input unchecked.
	 *
	 * @param string $model Requested model ID.
	 * @return string A model ID guaranteed to be in the allowlist.
	 */
	private function resolve_model( $model ) {
		if ( ! empty( $model ) && array_key_exists( $model, Nano_Banana_Settings::MODELS ) ) {
			return $model;
		}
		$stored = get_option( 'nano_banana_model', 'gemini-3.1-flash-image-preview' );
		return array_key_exists( $stored, Nano_Banana_Settings::MODELS ) ? $stored : 'gemini-3.1-flash-image-preview';
	}

	/**
	 * Resolve an aspect ratio to a known-good value.
	 *
	 * @param string $ratio Requested aspect ratio.
	 * @return string
	 */
	private function resolve_aspect_ratio( $ratio ) {
		if ( ! empty( $ratio ) && array_key_exists( $ratio, Nano_Banana_Settings::ASPECT_RATIOS ) ) {
			return $ratio;
		}
		$stored = get_option( 'nano_banana_default_aspect_ratio', '1:1' );
		return array_key_exists( $stored, Nano_Banana_Settings::ASPECT_RATIOS ) ? $stored : '1:1';
	}

	/**
	 * Resolve an image size to a known-good value.
	 *
	 * @param string $size Requested image size.
	 * @return string
	 */
	private function resolve_image_size( $size ) {
		if ( ! empty( $size ) && array_key_exists( $size, Nano_Banana_Settings::IMAGE_SIZES ) ) {
			return $size;
		}
		$stored = get_option( 'nano_banana_default_image_size', '1K' );
		return array_key_exists( $stored, Nano_Banana_Settings::IMAGE_SIZES ) ? $stored : '1K';
	}

	/**
	 * Build request arguments for a Gemini API call.
	 *
	 * The API key travels in a header rather than the query string so it does
	 * not leak into server access logs or HTTP debug output.
	 *
	 * @param string     $api_key Decrypted API key.
	 * @param array|null $body    Request body to JSON-encode, or null for GET.
	 * @return array
	 */
	private function request_args( $api_key, $body = null ) {
		$args = array(
			'headers' => array(
				'Content-Type'   => 'application/json',
				'x-goog-api-key' => $api_key,
			),
			'timeout' => null === $body ? 30 : 120,
		);
		if ( null !== $body ) {
			$args['body'] = wp_json_encode( $body );
		}
		return $args;
	}

	/**
	 * Generate images from a text prompt.
	 *
	 * @param string $prompt      The text prompt.
	 * @param array  $options {
	 *     Optional parameters.
	 *     @type string $model        Model ID. Default from settings.
	 *     @type string $aspect_ratio Aspect ratio. Default from settings.
	 *     @type int    $count        Number of images (1-4). Default 1.
	 * }
	 * @return array|WP_Error Array of image data arrays, or WP_Error on failure.
	 */
	public function generate_image( $prompt, $options = array() ) {
		$api_key = Nano_Banana_Settings::get_api_key();
		if ( ! $api_key ) {
			return new WP_Error( 'no_api_key', __( 'API key is not configured.', '45squared-ai-powered-image-generator' ) );
		}

		$model        = $this->resolve_model( $options['model'] ?? '' );
		$aspect_ratio = $this->resolve_aspect_ratio( $options['aspect_ratio'] ?? '' );
		$image_size   = $this->resolve_image_size( $options['image_size'] ?? '' );
		$count        = ! empty( $options['count'] ) ? min( absint( $options['count'] ), 4 ) : 1;

		$image_config = array(
			'aspectRatio' => $aspect_ratio,
		);
		if ( ! empty( $image_size ) ) {
			$image_config['imageSize'] = $image_size;
		}

		$body = array(
			'contents'         => array(
				array(
					'parts' => array(
						array( 'text' => $prompt ),
					),
				),
			),
			'generationConfig' => array(
				'responseModalities' => array( 'TEXT', 'IMAGE' ),
				'imageConfig'        => $image_config,
			),
		);

		$images = array();

		for ( $i = 0; $i < $count; $i++ ) {
			$url      = self::API_BASE . $model . ':generateContent';
			$response = wp_remote_post( $url, $this->request_args( $api_key, $body ) );

			if ( is_wp_error( $response ) ) {
				return $response;
			}

			$code = wp_remote_retrieve_response_code( $response );
			$data = json_decode( wp_remote_retrieve_body( $response ), true );

			if ( 200 !== $code ) {
				$message = isset( $data['error']['message'] ) ? $data['error']['message'] : __( 'API request failed.', '45squared-ai-powered-image-generator' );
				return new WP_Error( 'api_error', $message );
			}

			$image_data = $this->extract_image_from_response( $data );
			if ( is_wp_error( $image_data ) ) {
				return $image_data;
			}

			$images[] = $image_data;
		}

		return $images;
	}

	/**
	 * Edit an existing image with a text prompt.
	 *
	 * @param string $image_path Absolute path to the source image.
	 * @param string $prompt     Edit instruction.
	 * @param string $model      Optional model override.
	 * @return array|WP_Error Image data array or WP_Error.
	 */
	public function edit_image( $image_path, $prompt, $model = '' ) {
		$api_key = Nano_Banana_Settings::get_api_key();
		if ( ! $api_key ) {
			return new WP_Error( 'no_api_key', __( 'API key is not configured.', '45squared-ai-powered-image-generator' ) );
		}

		if ( ! file_exists( $image_path ) ) {
			return new WP_Error( 'file_not_found', __( 'Source image not found.', '45squared-ai-powered-image-generator' ) );
		}

		$model = $this->resolve_model( $model );

		$image_bytes = file_get_contents( $image_path ); // phpcs:ignore WordPress.WP.AlternativeFunctions.file_get_contents_file_get_contents
		$mime_type   = wp_check_filetype( $image_path )['type'];
		if ( ! $mime_type ) {
			$mime_type = 'image/png';
		}

		$body = array(
			'contents'         => array(
				array(
					'parts' => array(
						array( 'text' => $prompt ),
						array(
							'inlineData' => array(
								'mimeType' => $mime_type,
								'data'     => base64_encode( $image_bytes ), // phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_encode
							),
						),
					),
				),
			),
			'generationConfig' => array(
				'responseModalities' => array( 'TEXT', 'IMAGE' ),
			),
		);

		$url      = self::API_BASE . $model . ':generateContent';
		$response = wp_remote_post( $url, $this->request_args( $api_key, $body ) );

		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$code = wp_remote_retrieve_response_code( $response );
		$data = json_decode( wp_remote_retrieve_body( $response ), true );

		if ( 200 !== $code ) {
			$message = isset( $data['error']['message'] ) ? $data['error']['message'] : __( 'API request failed.', '45squared-ai-powered-image-generator' );
			return new WP_Error( 'api_error', $message );
		}

		return $this->extract_image_from_response( $data );
	}

	/**
	 * Test the API connection with a minimal request.
	 *
	 * @return true|WP_Error
	 */
	public function test_connection() {
		$api_key = Nano_Banana_Settings::get_api_key();
		if ( ! $api_key ) {
			return new WP_Error( 'no_api_key', __( 'API key is not configured.', '45squared-ai-powered-image-generator' ) );
		}

		$model = $this->resolve_model( '' );
		$url   = self::API_BASE . $model;

		$response = wp_remote_get( $url, $this->request_args( $api_key ) );

		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$code = wp_remote_retrieve_response_code( $response );
		if ( 200 !== $code ) {
			$data    = json_decode( wp_remote_retrieve_body( $response ), true );
			$message = isset( $data['error']['message'] ) ? $data['error']['message'] : __( 'Connection failed.', '45squared-ai-powered-image-generator' );
			return new WP_Error( 'connection_failed', $message );
		}

		return true;
	}

	/**
	 * Extract image binary data from a Gemini API response.
	 *
	 * @param array $response_data Decoded API response.
	 * @return array|WP_Error { 'data' => binary, 'mime_type' => string }
	 */
	private function extract_image_from_response( $response_data ) {
		if ( empty( $response_data['candidates'][0]['content']['parts'] ) ) {
			return new WP_Error( 'no_image', __( 'No image returned by the API.', '45squared-ai-powered-image-generator' ) );
		}

		foreach ( $response_data['candidates'][0]['content']['parts'] as $part ) {
			if ( isset( $part['inlineData'] ) ) {
				$mime = isset( $part['inlineData']['mimeType'] ) ? $part['inlineData']['mimeType'] : '';
				return array(
					'data'      => base64_decode( $part['inlineData']['data'] ), // phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_decode
					'mime_type' => in_array( $mime, self::ALLOWED_MIME_TYPES, true ) ? $mime : 'image/png',
				);
			}
		}

		return new WP_Error( 'no_image', __( 'Response did not contain an image.', '45squared-ai-powered-image-generator' ) );
	}

	/**
	 * Save generated image data to the WordPress Media Library.
	 *
	 * @param array  $image_data { 'data' => binary, 'mime_type' => string }.
	 * @param string $prompt     The prompt used to generate the image.
	 * @param string $model      The model used.
	 * @return int|WP_Error Attachment ID or WP_Error.
	 */
	public function save_to_media_library( $image_data, $prompt, $model = '' ) {
		if ( empty( $image_data['data'] ) ) {
			return new WP_Error( 'invalid_image', __( 'No image data to save.', '45squared-ai-powered-image-generator' ) );
		}

		// Never trust a caller-supplied MIME type: derive it from the actual
		// bytes and require that it be an image format we allow.
		$sniffed = false;
		if ( function_exists( 'getimagesizefromstring' ) ) {
			$info = @getimagesizefromstring( $image_data['data'] ); // phpcs:ignore WordPress.PHP.NoSilencedErrors.Discouraged
			if ( is_array( $info ) && ! empty( $info['mime'] ) ) {
				$sniffed = $info['mime'];
			}
		}
		if ( false === $sniffed || ! in_array( $sniffed, self::ALLOWED_MIME_TYPES, true ) ) {
			return new WP_Error( 'invalid_image', __( 'The data provided is not a supported image.', '45squared-ai-powered-image-generator' ) );
		}

		$extensions = array(
			'image/png'  => '.png',
			'image/jpeg' => '.jpg',
			'image/webp' => '.webp',
		);
		$ext        = $extensions[ $sniffed ];
		$filename   = 'nano-banana-' . wp_generate_uuid4() . $ext;

		$upload = wp_upload_bits( $filename, null, $image_data['data'] );
		if ( ! empty( $upload['error'] ) ) {
			return new WP_Error( 'upload_failed', $upload['error'] );
		}

		// Ensure WP_Filesystem is initialized — required by some plugins
		// (e.g. WP Offload Media) that hook into attachment creation.
		if ( ! function_exists( 'WP_Filesystem' ) ) {
			require_once ABSPATH . 'wp-admin/includes/file.php';
		}
		WP_Filesystem();

		$attachment = array(
			'post_title'     => sanitize_text_field( wp_trim_words( $prompt, 10 ) ),
			'post_content'   => '',
			'post_excerpt'   => sanitize_textarea_field( $prompt ),
			'post_status'    => 'inherit',
			'post_mime_type' => $sniffed,
		);

		$attach_id = wp_insert_attachment( $attachment, $upload['file'] );
		if ( is_wp_error( $attach_id ) ) {
			return $attach_id;
		}

		require_once ABSPATH . 'wp-admin/includes/image.php';
		$metadata = wp_generate_attachment_metadata( $attach_id, $upload['file'] );
		wp_update_attachment_metadata( $attach_id, $metadata );

		update_post_meta( $attach_id, '_nano_banana_prompt', sanitize_textarea_field( $prompt ) );
		update_post_meta( $attach_id, '_nano_banana_model', sanitize_text_field( $model ) );

		return $attach_id;
	}
}