<?php
/**
 * Yandex GPT & ART Client
 *
 * Unified client for YandexGPT (text generation) and YandexART (image generation).
 * Uses WordPress HTTP API instead of Guzzle.
 *
 * Based on tigusigalpa/yandexgpt-php, adapted for WordPress.
 *
 * @package SmartyPress_AI_Engine
 * @since 1.1.0
 */

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

/**
 * Yandex Client class
 */
class SmartyPress_AI_Yandex_Client {

	/**
	 * API endpoints
	 */
	const COMPLETION_ENDPOINT       = 'https://llm.api.cloud.yandex.net/foundationModels/v1/completion';
	const IMAGE_GENERATION_ENDPOINT = 'https://llm.api.cloud.yandex.net/foundationModels/v1/imageGenerationAsync';
	const OPERATIONS_ENDPOINT       = 'https://operation.api.cloud.yandex.net/operations';

	/**
	 * Auth manager instance
	 *
	 * @var SmartyPress_AI_Yandex_Auth_Manager
	 */
	private $auth_manager;

	/**
	 * Yandex Cloud folder ID
	 *
	 * @var string
	 */
	private $folder_id;

	/**
	 * Constructor
	 *
	 * @param string $api_key   Yandex Cloud service account API key.
	 * @param string $folder_id Yandex Cloud folder ID.
	 * @throws Exception If folder ID is empty.
	 */
	public function __construct( $api_key, $folder_id ) {
		if ( empty( $folder_id ) ) {
			throw new Exception(
				esc_html__( 'Yandex Cloud Folder ID cannot be empty.', 'smartypress-ai-engine' )
			);
		}

		$this->auth_manager = new SmartyPress_AI_Yandex_Auth_Manager( $api_key );
		$this->folder_id    = $folder_id;
	}

	/**
	 * Generate text from a single prompt
	 *
	 * @param string $prompt User prompt.
	 * @param string $model  Model name (default: yandexgpt-lite).
	 * @param array  $options Options: temperature, max_tokens, messages.
	 * @return array {success, data} or {success, error}.
	 */
	public function generate_text( $prompt, $model = 'yandexgpt-lite', $options = array() ) {
		$model_uri = SmartyPress_AI_Yandex_Models::get_model_uri( $model, $this->folder_id );

		// Yandex API accepts temperature only in range [0, 1], normalize it
		$temperature = min( 1.0, max( 0.0, floatval( $options['temperature'] ?? 0.6 ) ) );

		$body = array(
			'modelUri'          => $model_uri,
			'completionOptions' => array(
				'stream'      => false,
				'temperature' => $temperature,
				'maxTokens'   => intval( $options['max_tokens'] ?? 2000 ),
			),
			'messages'          => $options['messages'] ?? array(
				array(
					'role' => 'user',
					'text' => $prompt,
				),
			),
		);

		$response = wp_remote_post(
			self::COMPLETION_ENDPOINT,
			array(
				'headers' => array(
					'Content-Type'  => 'application/json',
					'Authorization' => $this->auth_manager->get_authorization_header(),
				),
				'body'    => wp_json_encode( $body ),
				'timeout' => 60,
			)
		);

		return $this->handle_response( $response );
	}

	/**
	 * Generate text from a dialog (multiple messages)
	 *
	 * @param array  $messages Array of messages [{role, text}, ...].
	 * @param string $model    Model name.
	 * @param array  $options  Options: temperature, max_tokens.
	 * @return array {success, data} or {success, error}.
	 */
	public function generate_from_messages( $messages, $model = 'yandexgpt-lite', $options = array() ) {
		$options['messages'] = $messages;
		return $this->generate_text( '', $model, $options );
	}

	/**
	 * Generate image synchronously (with polling)
	 *
	 * Sends an async request to YandexART, then polls for the result.
	 *
	 * @param array|string $prompt  Prompt string or array of [{text, weight}, ...].
	 * @param array        $options Options: max_wait (seconds), poll_interval (seconds).
	 * @return array {success, b64_json, operation_id} or {success, error}.
	 */
	public function generate_image( $prompt, $options = array() ) {
		// Start async generation.
		$operation = $this->generate_image_async( $prompt, $options );

		if ( ! $operation['success'] ) {
			return $operation;
		}

		$operation_id = $operation['data']['id'] ?? null;

		if ( ! $operation_id ) {
			return array(
				'success' => false,
				'error'   => __( 'Operation ID not found in YandexART response.', 'smartypress-ai-engine' ),
			);
		}

		// Poll for result.
		$max_wait      = intval( $options['max_wait'] ?? 600 );
		$poll_interval = intval( $options['poll_interval'] ?? 2 );
		$elapsed       = 0;

		while ( $elapsed < $max_wait ) {
			sleep( $poll_interval );
			$elapsed += $poll_interval;

			$status = $this->get_operation( $operation_id );

			if ( ! $status['success'] ) {
				return $status;
			}

			$done = $status['data']['done'] ?? false;

			if ( $done ) {
				// Check for errors in the operation.
				if ( ! empty( $status['data']['error'] ) ) {
					$error_message = $status['data']['error']['message']
						?? __( 'Unknown YandexART error', 'smartypress-ai-engine' );

					return array(
						'success' => false,
						'error'   => $error_message,
					);
				}

				// Extract base64 image.
				$image_base64 = $status['data']['response']['image'] ?? null;

				if ( ! $image_base64 ) {
					return array(
						'success' => false,
						'error'   => __( 'Image data not found in YandexART response.', 'smartypress-ai-engine' ),
					);
				}

				return array(
					'success'      => true,
					'b64_json'     => $image_base64,
					'operation_id' => $operation_id,
				);
			}
		}

		return array(
			'success' => false,
			'error'   => __( 'YandexART image generation timed out.', 'smartypress-ai-engine' ),
		);
	}

	/**
	 * Start async image generation (YandexART)
	 *
	 * @param array|string $prompt  Prompt string or array of [{text, weight}, ...].
	 * @param array        $options Generation options.
	 * @return array {success, data} or {success, error}.
	 */
	public function generate_image_async( $prompt, $options = array() ) {
		$model_uri = SmartyPress_AI_Yandex_Models::get_art_model_uri( $this->folder_id );

		// Normalize messages.
		$messages = $this->normalize_art_messages( $prompt );

		// Ensure generationOptions is an empty object, not an array.
		$generation_options = ! empty( $options['generation_options'] )
			? $options['generation_options']
			: new stdClass();

		$body = array(
			'modelUri'          => $model_uri,
			'generationOptions' => $generation_options,
			'messages'          => $messages,
		);

		$response = wp_remote_post(
			self::IMAGE_GENERATION_ENDPOINT,
			array(
				'headers' => array(
					'Content-Type'  => 'application/json',
					'Authorization' => $this->auth_manager->get_authorization_header(),
				),
				'body'    => wp_json_encode( $body ),
				'timeout' => 30,
			)
		);

		return $this->handle_response( $response );
	}

	/**
	 * Get operation status
	 *
	 * @param string $operation_id Operation ID.
	 * @return array {success, data} or {success, error}.
	 */
	public function get_operation( $operation_id ) {
		$response = wp_remote_get(
			self::OPERATIONS_ENDPOINT . '/' . $operation_id,
			array(
				'headers' => array(
					'Authorization' => $this->auth_manager->get_authorization_header(),
					'Content-Type'  => 'application/json',
				),
				'timeout' => 30,
			)
		);

		return $this->handle_response( $response );
	}

	/**
	 * Normalize messages for YandexART
	 *
	 * Accepts a string (single prompt) or an array of strings/arrays.
	 * Converts each element to {text, weight}.
	 *
	 * @param array|string $messages Input messages.
	 * @return array Normalized messages.
	 */
	private function normalize_art_messages( $messages ) {
		if ( is_string( $messages ) ) {
			return array(
				array(
					'text'   => $messages,
					'weight' => 1,
				),
			);
		}

		$normalized = array();

		foreach ( $messages as $msg ) {
			if ( is_string( $msg ) ) {
				$normalized[] = array(
					'text'   => $msg,
					'weight' => 1,
				);
			} elseif ( is_array( $msg ) ) {
				$normalized[] = array(
					'text'   => $msg['text'] ?? '',
					'weight' => $msg['weight'] ?? 1,
				);
			}
		}

		return $normalized;
	}

	/**
	 * Handle HTTP response
	 *
	 * @param array|WP_Error $response WordPress HTTP response.
	 * @return array {success, data} or {success, error, code}.
	 */
	private function handle_response( $response ) {
		if ( is_wp_error( $response ) ) {
			return array(
				'success' => false,
				'error'   => $response->get_error_message(),
			);
		}

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

		if ( json_last_error() !== JSON_ERROR_NONE ) {
			return array(
				'success' => false,
				'error'   => __( 'Failed to decode Yandex API response.', 'smartypress-ai-engine' ),
			);
		}

		if ( 200 !== $code ) {
			// Yandex API can return errors in different formats.
			$error = $data['error']['message']
				?? $data['message']
				?? __( 'Unknown Yandex API error', 'smartypress-ai-engine' );

			return array(
				'success' => false,
				'error'   => $error,
				'code'    => $code,
			);
		}

		return array(
			'success' => true,
			'data'    => $data,
		);
	}

	/**
	 * Get the auth manager instance
	 *
	 * @return SmartyPress_AI_Yandex_Auth_Manager
	 */
	public function get_auth_manager() {
		return $this->auth_manager;
	}

	/**
	 * Get the folder ID
	 *
	 * @return string
	 */
	public function get_folder_id() {
		return $this->folder_id;
	}
}
