<?php
/**
 * Non-blocking fal queue client (submit / status / result).
 *
 * Lets consumers enqueue many fal jobs without holding a PHP worker for the
 * full inference window. Uses the same queue.fal.run flow as FalImageGenerationModel.
 *
 * @package KraftySprouts\AiProviderForFalAi\Queue
 * @since   1.1.8
 */

declare( strict_types=1 );

namespace KraftySprouts\AiProviderForFalAi\Queue;

use KraftySprouts\AiProviderForFalAi\Provider\FalProvider;
use WordPress\AiClient\AiClient;
use WP_Error;

/**
 * Class FalQueueClient
 *
 * @since 1.1.8
 */
class FalQueueClient {

	/**
	 * Per-call HTTP timeout (seconds).
	 *
	 * @since 1.1.8
	 * @var int
	 */
	private const HTTP_TIMEOUT = 30;

	/**
	 * Whether an API key is available for queue calls.
	 *
	 * @since 1.1.8
	 * @return bool
	 */
	public static function isAvailable(): bool {
		return '' !== self::apiKey();
	}

	/**
	 * Submit a text-to-image job; returns tracking URLs immediately.
	 *
	 * @since 1.1.8
	 * @param string      $model_id     fal model id.
	 * @param string      $prompt       Image prompt.
	 * @param string|null $aspect_ratio Optional aspect ratio (e.g. 16:9).
	 * @return array{request_id:string,status_url:string,response_url:string,model_id:string}|WP_Error
	 */
	public static function submit( string $model_id, string $prompt, ?string $aspect_ratio = null ) {
		$key = self::apiKey();
		if ( '' === $key ) {
			return new WP_Error(
				'aipffa_fal_no_key',
				__( 'fal.ai API key is not configured.', 'ai-provider-for-fal-ai' )
			);
		}

		$model_id = ltrim( $model_id, '/' );
		if ( '' === $model_id || '' === trim( $prompt ) ) {
			return new WP_Error(
				'aipffa_fal_invalid_submit',
				__( 'Model id and prompt are required.', 'ai-provider-for-fal-ai' )
			);
		}

		$body = self::buildRequestBody( $model_id, $prompt, $aspect_ratio );
		$url  = FalProvider::queueUrl( $model_id );

		$response = self::request( 'POST', $url, $key, $body );
		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$request_id   = isset( $response['request_id'] ) && is_string( $response['request_id'] ) ? $response['request_id'] : '';
		$status_url   = isset( $response['status_url'] ) && is_string( $response['status_url'] ) ? $response['status_url'] : '';
		$response_url = isset( $response['response_url'] ) && is_string( $response['response_url'] ) ? $response['response_url'] : '';

		if ( '' === $request_id || '' === $status_url || '' === $response_url ) {
			return new WP_Error(
				'aipffa_fal_bad_submit',
				__( 'fal.ai queue submit did not return tracking URLs.', 'ai-provider-for-fal-ai' )
			);
		}

		if ( is_wp_error( self::assertQueueUrl( $status_url ) ) ) {
			return self::assertQueueUrl( $status_url );
		}
		if ( is_wp_error( self::assertQueueUrl( $response_url ) ) ) {
			return self::assertQueueUrl( $response_url );
		}

		return array(
			'request_id'   => $request_id,
			'status_url'   => $status_url,
			'response_url' => $response_url,
			'model_id'     => $model_id,
		);
	}

	/**
	 * One status poll. Returns uppercase fal status (COMPLETED, IN_PROGRESS, …).
	 *
	 * @since 1.1.8
	 * @param string $status_url Queue status URL.
	 * @param string $request_id Request id (for errors).
	 * @return string|WP_Error
	 */
	public static function status( string $status_url, string $request_id = '' ) {
		$key = self::apiKey();
		if ( '' === $key ) {
			return new WP_Error(
				'aipffa_fal_no_key',
				__( 'fal.ai API key is not configured.', 'ai-provider-for-fal-ai' )
			);
		}

		$guard = self::assertQueueUrl( $status_url );
		if ( is_wp_error( $guard ) ) {
			return $guard;
		}

		$response = self::request( 'GET', $status_url, $key, null );
		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$status = isset( $response['status'] ) && is_string( $response['status'] )
			? strtoupper( $response['status'] )
			: '';

		if ( 'COMPLETED' === $status
			&& isset( $response['error'] )
			&& is_string( $response['error'] )
			&& '' !== $response['error']
		) {
			return new WP_Error(
				'aipffa_fal_failed',
				sprintf(
					/* translators: 1: request id, 2: fal error. */
					__( 'fal.ai request %1$s failed: %2$s', 'ai-provider-for-fal-ai' ),
					$request_id,
					$response['error']
				)
			);
		}

		if ( in_array( $status, array( 'FAILED', 'CANCELLED', 'CANCELED' ), true ) ) {
			$message = isset( $response['error'] ) && is_string( $response['error'] ) && '' !== $response['error']
				? $response['error']
				: $status;
			return new WP_Error(
				'aipffa_fal_failed',
				sprintf(
					/* translators: 1: request id, 2: status or error. */
					__( 'fal.ai request %1$s ended with status: %2$s', 'ai-provider-for-fal-ai' ),
					$request_id,
					$message
				)
			);
		}

		return $status;
	}

	/**
	 * Fetch completed result and extract the first image URL (+ mime).
	 *
	 * @since 1.1.8
	 * @param string $response_url Queue response URL.
	 * @return array{url:string,mime:string}|WP_Error
	 */
	public static function fetchImage( string $response_url ) {
		$key = self::apiKey();
		if ( '' === $key ) {
			return new WP_Error(
				'aipffa_fal_no_key',
				__( 'fal.ai API key is not configured.', 'ai-provider-for-fal-ai' )
			);
		}

		$guard = self::assertQueueUrl( $response_url );
		if ( is_wp_error( $guard ) ) {
			return $guard;
		}

		$response = self::request( 'GET', $response_url, $key, null );
		if ( is_wp_error( $response ) ) {
			return $response;
		}

		return self::extractImageFromResult( $response );
	}

	/**
	 * Fetch completed result by model id + request id (recovery path).
	 *
	 * @since 1.1.9
	 * @param string $model_id   fal model id.
	 * @param string $request_id Queue request id.
	 * @return array{url:string,mime:string}|WP_Error
	 */
	public static function fetchByRequestId( string $model_id, string $request_id ) {
		$model_id   = ltrim( $model_id, '/' );
		$request_id = trim( $request_id );
		if ( '' === $model_id || '' === $request_id ) {
			return new WP_Error(
				'aipffa_fal_invalid_ids',
				__( 'Model id and request id are required.', 'ai-provider-for-fal-ai' )
			);
		}

		$key = self::apiKey();
		if ( '' === $key ) {
			return new WP_Error(
				'aipffa_fal_no_key',
				__( 'fal.ai API key is not configured.', 'ai-provider-for-fal-ai' )
			);
		}

		$url = add_query_arg(
			array(
				'endpoint_id' => $model_id,
				'request_id' => $request_id,
				'expand'     => 'payloads',
				'limit'      => 1,
			),
			'https://api.fal.ai/v1/models/requests/by-endpoint'
		);
		$response = self::request( 'GET', $url, $key, null );
		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$items = isset( $response['items'] ) && is_array( $response['items'] )
			? array_values( $response['items'] )
			: array();
		if ( empty( $items[0]['json_output'] ) || ! is_array( $items[0]['json_output'] ) ) {
			return new WP_Error(
				'aipffa_fal_request_not_found',
				__( 'No completed fal.ai result was found for that request id.', 'ai-provider-for-fal-ai' )
			);
		}

		return self::extractImageFromResult( $items[0]['json_output'] );
	}

	/**
	 * Fetch raw fal result JSON (for prompt matching / diagnostics).
	 *
	 * @since 1.1.9
	 * @param string $response_url Absolute queue.fal.run result URL.
	 * @return array<string, mixed>|WP_Error
	 */
	public static function fetchRawResult( string $response_url ) {
		$key = self::apiKey();
		if ( '' === $key ) {
			return new WP_Error(
				'aipffa_fal_no_key',
				__( 'fal.ai API key is not configured.', 'ai-provider-for-fal-ai' )
			);
		}

		$guard = self::assertQueueUrl( $response_url );
		if ( is_wp_error( $guard ) ) {
			return $guard;
		}

		return self::request( 'GET', $response_url, $key, null );
	}

	/**
	 * List recent successful requests, including input/output payloads.
	 *
	 * Uses fal's authenticated platform history API so consumers can recover
	 * completed requests even when their local queue handle was lost.
	 *
	 * @since 1.1.10
	 * @param string[] $model_ids Endpoint/model IDs.
	 * @param int      $hours     Lookback window in hours.
	 * @param int      $limit     Maximum requests to return.
	 * @return list<array<string, mixed>>|WP_Error
	 */
	public static function listRecentRequests( array $model_ids, int $hours = 48, int $limit = 500 ) {
		$key = self::apiKey();
		if ( '' === $key ) {
			return new WP_Error(
				'aipffa_fal_no_key',
				__( 'fal.ai API key is not configured.', 'ai-provider-for-fal-ai' )
			);
		}

		$models = array();
		foreach ( $model_ids as $model_id ) {
			$model_id = is_string( $model_id ) ? ltrim( trim( $model_id ), '/' ) : '';
			if ( '' !== $model_id && preg_match( '#^[a-zA-Z0-9._/-]+$#', $model_id ) ) {
				$models[] = $model_id;
			}
		}
		$models = array_slice( array_values( array_unique( $models ) ), 0, 50 );
		if ( array() === $models ) {
			return new WP_Error(
				'aipffa_fal_no_models',
				__( 'At least one valid fal.ai model id is required.', 'ai-provider-for-fal-ai' )
			);
		}

		$hours = max( 1, min( 2160, $hours ) );
		$limit = max( 1, min( 1000, $limit ) );
		$items = array();
		$cursor = '';

		do {
			$page_size = min( 100, $limit - count( $items ) );
			$query     = array(
				'endpoint_id' => implode( ',', $models ),
				'start'       => gmdate( 'Y-m-d\TH:i:s\Z', time() - ( $hours * HOUR_IN_SECONDS ) ),
				'expand'      => 'payloads',
				'limit'       => $page_size,
			);
			if ( '' !== $cursor ) {
				$query['cursor'] = $cursor;
			}

			$url      = add_query_arg( $query, 'https://api.fal.ai/v1/models/requests/by-endpoint' );
			$response = self::request( 'GET', $url, $key, null );
			if ( is_wp_error( $response ) ) {
				return $response;
			}

			$page_items = isset( $response['items'] ) && is_array( $response['items'] )
				? $response['items']
				: array();
			foreach ( $page_items as $item ) {
				if ( is_array( $item ) ) {
					// Drop heavy base64 / unused fields immediately — full payloads
					// have OOM'd WordPress admin recover requests on large batches.
					$items[] = self::slimHistoryItem( $item );
				}
				if ( count( $items ) >= $limit ) {
					break;
				}
			}

			$cursor = isset( $response['next_cursor'] ) && is_string( $response['next_cursor'] )
				? $response['next_cursor']
				: '';
		} while ( '' !== $cursor && count( $items ) < $limit );

		return $items;
	}

	/**
	 * Keep only the fields recovery needs from a fal history row.
	 *
	 * @since 1.1.11
	 * @param array<string, mixed> $item Raw history item.
	 * @return array<string, mixed>
	 */
	private static function slimHistoryItem( array $item ): array {
		$input  = isset( $item['json_input'] ) && is_array( $item['json_input'] ) ? $item['json_input'] : array();
		$output = isset( $item['json_output'] ) && is_array( $item['json_output'] ) ? $item['json_output'] : array();

		$prompt = '';
		foreach ( array( 'prompt', 'image_prompt' ) as $key ) {
			if ( isset( $input[ $key ] ) && is_string( $input[ $key ] ) && '' !== trim( $input[ $key ] ) ) {
				$prompt = trim( $input[ $key ] );
				break;
			}
		}

		$image = array();
		$raw_images = array();
		if ( isset( $output['images'] ) && is_array( $output['images'] ) ) {
			$raw_images = array_values( $output['images'] );
		} elseif ( isset( $output['image'] ) && is_array( $output['image'] ) ) {
			$raw_images = array( $output['image'] );
		}
		if ( isset( $raw_images[0] ) && is_array( $raw_images[0] ) ) {
			$first = $raw_images[0];
			$url   = isset( $first['url'] ) && is_string( $first['url'] ) ? $first['url'] : '';
			if ( '' !== $url ) {
				$mime = 'image/jpeg';
				if ( isset( $first['content_type'] ) && is_string( $first['content_type'] ) ) {
					$mime = $first['content_type'];
				} elseif ( isset( $first['contentType'] ) && is_string( $first['contentType'] ) ) {
					$mime = $first['contentType'];
				}
				$image = array(
					'url'          => $url,
					'content_type' => $mime,
				);
			}
		}

		return array(
			'request_id'  => isset( $item['request_id'] ) ? (string) $item['request_id'] : '',
			'endpoint_id' => isset( $item['endpoint_id'] ) ? (string) $item['endpoint_id'] : '',
			'status_code' => isset( $item['status_code'] ) ? (int) $item['status_code'] : 0,
			'sent_at'     => isset( $item['sent_at'] ) ? (string) $item['sent_at'] : '',
			'json_input'  => array(
				'prompt' => $prompt,
			),
			'json_output' => array(
				'images' => $image ? array( $image ) : array(),
			),
		);
	}

	/**
	 * Pull first image URL/mime from a fal result body.
	 *
	 * @since 1.1.9
	 * @param array<string, mixed> $response Decoded fal body.
	 * @return array{url:string,mime:string}|WP_Error
	 */
	private static function extractImageFromResult( array $response ) {
		$images = array();
		if ( isset( $response['images'] ) && is_array( $response['images'] ) ) {
			$images = array_values( $response['images'] );
		} elseif ( isset( $response['image'] ) && is_array( $response['image'] ) ) {
			$images = array( $response['image'] );
		}

		if ( empty( $images ) || ! is_array( $images[0] ) ) {
			return new WP_Error(
				'aipffa_fal_no_image',
				__( 'fal.ai result contained no image.', 'ai-provider-for-fal-ai' )
			);
		}

		$image = $images[0];
		$mime  = 'image/jpeg';
		if ( isset( $image['content_type'] ) && is_string( $image['content_type'] ) ) {
			$mime = $image['content_type'];
		} elseif ( isset( $image['contentType'] ) && is_string( $image['contentType'] ) ) {
			$mime = $image['contentType'];
		}

		if ( isset( $image['url'] ) && is_string( $image['url'] ) && '' !== $image['url'] ) {
			return array(
				'url'  => $image['url'],
				'mime' => $mime,
			);
		}

		return new WP_Error(
			'aipffa_fal_no_image_url',
			__( 'fal.ai result image had no URL.', 'ai-provider-for-fal-ai' )
		);
	}

	/**
	 * Build fal text-to-image body (mirrors FalImageGenerationModel sizing rules).
	 *
	 * @since 1.1.8
	 * @param string      $model_id     Model id.
	 * @param string      $prompt       Prompt.
	 * @param string|null $aspect_ratio Aspect ratio.
	 * @return array<string, mixed>
	 */
	public static function buildRequestBody( string $model_id, string $prompt, ?string $aspect_ratio = null ): array {
		$params = array(
			'prompt'     => $prompt,
			'num_images' => 1,
		);

		if ( self::supportsSharedImageSizeParam( $model_id ) ) {
			$size = self::prepareImageSizeParam( $model_id, $aspect_ratio );
			if ( null !== $size ) {
				$params['image_size'] = $size;
			}
		} elseif ( self::supportsAspectRatioParam( $model_id ) && is_string( $aspect_ratio ) && '' !== $aspect_ratio ) {
			$params['aspect_ratio'] = $aspect_ratio;
		}

		return $params;
	}

	/**
	 * Resolve fal API key from AI Client registry or FAL_KEY.
	 *
	 * @since 1.1.8
	 * @return string
	 */
	private static function apiKey(): string {
		if ( class_exists( AiClient::class ) ) {
			try {
				$registry = AiClient::defaultRegistry();
				if ( method_exists( $registry, 'getProviderRequestAuthentication' ) ) {
					$auth = $registry->getProviderRequestAuthentication( FalProvider::class );
					if ( is_object( $auth ) && method_exists( $auth, 'getApiKey' ) ) {
						$key = (string) $auth->getApiKey();
						if ( '' !== trim( $key ) ) {
							return trim( $key );
						}
					}
				}
			} catch ( \Throwable $e ) { // phpcs:ignore Generic.CodeAnalysis.EmptyStatement.DetectedCatch
				// Fall through to env/constant.
			}
		}

		$key = getenv( 'FAL_KEY' );
		if ( false === $key && defined( 'FAL_KEY' ) ) {
			$constant = constant( 'FAL_KEY' );
			$key      = is_scalar( $constant ) ? (string) $constant : '';
		}

		return is_string( $key ) ? trim( $key ) : '';
	}

	/**
	 * Authenticated JSON request to fal queue.
	 *
	 * @since 1.1.8
	 * @param string                    $method GET|POST.
	 * @param string                    $url    Absolute URL.
	 * @param string                    $key    API key.
	 * @param array<string, mixed>|null $body   JSON body or null.
	 * @return array<string, mixed>|WP_Error
	 */
	private static function request( string $method, string $url, string $key, ?array $body ) {
		$args = array(
			'method'  => $method,
			'timeout' => self::HTTP_TIMEOUT,
			'headers' => array(
				'Authorization' => 'Key ' . $key,
				'Accept'        => 'application/json',
			),
		);

		if ( null !== $body ) {
			$args['headers']['Content-Type'] = 'application/json';
			$args['body']                    = wp_json_encode( $body );
		}

		$response = wp_remote_request( $url, $args );
		if ( is_wp_error( $response ) ) {
			return $response;
		}

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

		if ( $code < 200 || $code >= 300 ) {
			$message = '';
			if ( is_array( $data ) && isset( $data['detail'] ) && is_string( $data['detail'] ) ) {
				$message = $data['detail'];
			} elseif ( is_array( $data ) && isset( $data['error']['message'] ) && is_string( $data['error']['message'] ) ) {
				$message = $data['error']['message'];
			}
			if ( '' === $message ) {
				$message = sprintf(
					/* translators: %d: HTTP status code. */
					__( 'fal.ai HTTP error %d', 'ai-provider-for-fal-ai' ),
					$code
				);
			}
			return new WP_Error(
				'aipffa_fal_http',
				$message,
				array(
					'status'   => $code,
					'response' => $data,
				)
			);
		}

		if ( ! is_array( $data ) ) {
			return new WP_Error(
				'aipffa_fal_bad_json',
				__( 'fal.ai returned a non-JSON body.', 'ai-provider-for-fal-ai' )
			);
		}

		return $data;
	}

	/**
	 * SSRF guard: only queue.fal.run hosts.
	 *
	 * @since 1.1.8
	 * @param string $url URL from fal.
	 * @return true|WP_Error
	 */
	private static function assertQueueUrl( string $url ) {
		$host = wp_parse_url( $url, PHP_URL_HOST );
		if ( ! is_string( $host ) || 'queue.fal.run' !== strtolower( $host ) ) {
			return new WP_Error(
				'aipffa_fal_bad_url',
				__( 'Queue tracking URL must be on queue.fal.run.', 'ai-provider-for-fal-ai' )
			);
		}
		return true;
	}

	/**
	 * @since 1.1.8
	 * @param string $model_id Model id.
	 * @return bool
	 */
	private static function supportsSharedImageSizeParam( string $model_id ): bool {
		$with_image_size = array(
			'openai/gpt-image-2',
			'bytedance/seedream/v5/pro/text-to-image',
		);

		return in_array( $model_id, $with_image_size, true )
			|| 0 === strpos( $model_id, 'fal-ai/flux' )
			|| 0 === strpos( $model_id, 'fal-ai/recraft' );
	}

	/**
	 * @since 1.1.8
	 * @param string $model_id Model id.
	 * @return bool
	 */
	private static function supportsAspectRatioParam( string $model_id ): bool {
		return 0 === strpos( $model_id, 'fal-ai/nano-banana' );
	}

	/**
	 * @since 1.1.8
	 * @param string      $model_id     Model id.
	 * @param string|null $aspect_ratio Aspect ratio.
	 * @return string|array{width:int,height:int}|null
	 */
	private static function prepareImageSizeParam( string $model_id, ?string $aspect_ratio ) {
		$wants_16_9 = ( '16:9' === $aspect_ratio );

		if ( $wants_16_9 && 'bytedance/seedream/v5/pro/text-to-image' === $model_id ) {
			return array(
				'width'  => 2048,
				'height' => 1152,
			);
		}

		if ( null !== $aspect_ratio ) {
			$map = array(
				'1:1'  => 'square_hd',
				'16:9' => 'landscape_16_9',
				'9:16' => 'portrait_16_9',
				'4:3'  => 'landscape_4_3',
				'3:4'  => 'portrait_4_3',
			);
			if ( isset( $map[ $aspect_ratio ] ) ) {
				return $map[ $aspect_ratio ];
			}
		}

		return null;
	}
}