<?php
/**
 * Replicate image provider (FLUX, SDXL, Ideogram, Recraft, etc).
 *
 * @package CatalogPilot\AiStudio
 */

namespace CatalogPilot\AiStudio\Providers;

defined( 'ABSPATH' ) || exit;

final class Image_Replicate_Provider extends Image_Provider {

	const API = 'https://api.replicate.com/v1';

	public function id() {
		return 'replicate';
	}

	public function generate( $prompt, array $opts = array() ) {
		if ( '' === $this->api_key ) {
			return new \WP_Error( 'cpais_missing_key', __( 'Replicate API token is required.', 'catalogpilot-ai-studio-for-woocommerce' ) );
		}

		$model = $this->model ?: 'black-forest-labs/flux-schnell';
		if ( ! preg_match( '#^[^/]+/[^/]+$#', $model ) ) {
			return new \WP_Error( 'cpais_invalid_model', __( 'Replicate model must be in owner/name form.', 'catalogpilot-ai-studio-for-woocommerce' ) );
		}

		$aspect = $opts['size'] ?? '1:1';
		$count  = max( 1, min( 4, (int) ( $opts['count'] ?? 1 ) ) );

		$input = array(
			'prompt'        => (string) $prompt,
			'aspect_ratio'  => $aspect,
			'num_outputs'   => $count,
			'output_format' => 'png',
		);
		if ( ! empty( $opts['negative_prompt'] ) ) {
			$input['negative_prompt'] = (string) $opts['negative_prompt'];
		}
		if ( isset( $opts['seed'] ) && '' !== $opts['seed'] ) {
			$input['seed'] = (int) $opts['seed'];
		}

		$result = $this->post_json(
			self::API . '/models/' . $model . '/predictions',
			array( 'input' => $input ),
			array(
				'Authorization' => 'Token ' . $this->api_key,
				'Prefer'        => 'wait=60',
			)
		);
		if ( is_wp_error( $result ) ) {
			return $result;
		}

		if ( 'succeeded' !== ( $result['status'] ?? '' ) ) {
			// Poll for completion up to ~90s.
			$poll_url = $result['urls']['get'] ?? '';
			if ( '' === $poll_url ) {
				return new \WP_Error( 'cpais_image_error', __( 'Replicate prediction did not complete in time.', 'catalogpilot-ai-studio-for-woocommerce' ) );
			}
			$deadline = time() + 90;
			while ( time() < $deadline ) {
				sleep( 2 );
				$resp = wp_remote_get(
					$poll_url,
					array(
						'timeout' => 15,
						'headers' => array( 'Authorization' => 'Token ' . $this->api_key ),
					)
				);
				if ( is_wp_error( $resp ) ) {
					return $resp;
				}
				$result = json_decode( wp_remote_retrieve_body( $resp ), true );
				$status = $result['status'] ?? '';
				if ( 'succeeded' === $status ) {
					break;
				}
				if ( in_array( $status, array( 'failed', 'canceled' ), true ) ) {
					return new \WP_Error( 'cpais_image_error', $result['error'] ?? __( 'Prediction failed.', 'catalogpilot-ai-studio-for-woocommerce' ) );
				}
			}
		}

		$urls = array();
		$output = $result['output'] ?? null;
		if ( is_string( $output ) ) {
			$urls = array( $output );
		} elseif ( is_array( $output ) ) {
			foreach ( $output as $o ) {
				if ( is_string( $o ) && preg_match( '#^https?://#i', $o ) ) {
					$urls[] = $o;
				}
			}
		}

		if ( empty( $urls ) ) {
			return new \WP_Error( 'cpais_image_empty', __( 'Replicate returned no image URLs.', 'catalogpilot-ai-studio-for-woocommerce' ) );
		}

		$images = array();
		foreach ( $urls as $url ) {
			$fetched = self::fetch_url_bytes( $url );
			if ( ! is_wp_error( $fetched ) ) {
				$images[] = $fetched;
			}
		}

		if ( empty( $images ) ) {
			return new \WP_Error( 'cpais_image_empty', __( 'Could not download generated images.', 'catalogpilot-ai-studio-for-woocommerce' ) );
		}

		return $images;
	}

	public function test_connection() {
		if ( '' === $this->api_key ) {
			return new \WP_Error( 'cpais_missing_key', __( 'Replicate API token is required.', 'catalogpilot-ai-studio-for-woocommerce' ) );
		}
		$response = wp_remote_get(
			self::API . '/account',
			array(
				'timeout' => 15,
				'headers' => array( 'Authorization' => 'Token ' . $this->api_key ),
			)
		);
		if ( is_wp_error( $response ) ) {
			return $response;
		}
		$code = wp_remote_retrieve_response_code( $response );
		if ( 200 !== $code ) {
			return new \WP_Error( 'cpais_auth_failed', __( 'Authentication failed.', 'catalogpilot-ai-studio-for-woocommerce' ), array( 'status' => $code ) );
		}
		return true;
	}

	/**
	 * Remove the background from an image using a Replicate background-remover model.
	 * Default model: `851-labs/background-remover` (RMBG-1.4 powered).
	 *
	 * @param string $image_url Public URL of the image to process.
	 * @param string $model     Replicate model in owner/name form.
	 * @return array|\WP_Error  { bytes, mime } on success.
	 */
	public function remove_background( $image_url, $model = '851-labs/background-remover' ) {
		if ( '' === $this->api_key ) {
			return new \WP_Error( 'cpais_missing_key', __( 'Replicate API token is required.', 'catalogpilot-ai-studio-for-woocommerce' ), array( 'status' => 400 ) );
		}
		if ( ! preg_match( '#^https?://#i', (string) $image_url ) ) {
			return new \WP_Error( 'cpais_invalid_url', __( 'Image URL is required.', 'catalogpilot-ai-studio-for-woocommerce' ), array( 'status' => 400 ) );
		}
		if ( ! preg_match( '#^[^/]+/[^/]+$#', (string) $model ) ) {
			return new \WP_Error( 'cpais_invalid_model', __( 'Replicate model must be in owner/name form.', 'catalogpilot-ai-studio-for-woocommerce' ), array( 'status' => 400 ) );
		}

		// Fetch the latest version id — Replicate deprecated /models/{owner}/{name}/predictions
		// for community models, so we always go through POST /predictions with a version.
		$meta_resp = wp_remote_get(
			self::API . '/models/' . $model,
			array( 'timeout' => 15, 'headers' => array( 'Authorization' => 'Token ' . $this->api_key ) )
		);
		if ( is_wp_error( $meta_resp ) ) {
			return $meta_resp;
		}
		$meta = json_decode( wp_remote_retrieve_body( $meta_resp ), true );
		$version = $meta['latest_version']['id'] ?? '';
		if ( '' === $version ) {
			return new \WP_Error( 'cpais_model_no_version', __( 'Replicate model has no runnable version.', 'catalogpilot-ai-studio-for-woocommerce' ) );
		}

		$result = $this->post_json(
			self::API . '/predictions',
			array(
				'version' => $version,
				'input'   => array( 'image' => $image_url ),
			),
			array(
				'Authorization' => 'Token ' . $this->api_key,
				'Prefer'        => 'wait=60',
			)
		);
		if ( is_wp_error( $result ) ) {
			return $result;
		}

		// Poll if not yet succeeded.
		if ( 'succeeded' !== ( $result['status'] ?? '' ) ) {
			$poll_url = $result['urls']['get'] ?? '';
			if ( '' === $poll_url ) {
				return new \WP_Error( 'cpais_image_error', __( 'Replicate background removal did not complete in time.', 'catalogpilot-ai-studio-for-woocommerce' ) );
			}
			$deadline = time() + 90;
			while ( time() < $deadline ) {
				sleep( 2 );
				$resp = wp_remote_get(
					$poll_url,
					array(
						'timeout' => 15,
						'headers' => array( 'Authorization' => 'Token ' . $this->api_key ),
					)
				);
				if ( is_wp_error( $resp ) ) {
					return $resp;
				}
				$result = json_decode( wp_remote_retrieve_body( $resp ), true );
				$status = $result['status'] ?? '';
				if ( 'succeeded' === $status ) {
					break;
				}
				if ( in_array( $status, array( 'failed', 'canceled' ), true ) ) {
					return new \WP_Error( 'cpais_image_error', $result['error'] ?? __( 'Background removal failed.', 'catalogpilot-ai-studio-for-woocommerce' ) );
				}
			}
		}

		$output = $result['output'] ?? null;
		$out_url = '';
		if ( is_string( $output ) ) {
			$out_url = $output;
		} elseif ( is_array( $output ) && isset( $output[0] ) && is_string( $output[0] ) ) {
			$out_url = $output[0];
		}
		if ( ! preg_match( '#^https?://#i', $out_url ) ) {
			return new \WP_Error( 'cpais_image_empty', __( 'Replicate returned no output URL.', 'catalogpilot-ai-studio-for-woocommerce' ) );
		}

		$fetched = self::fetch_url_bytes( $out_url );
		if ( is_wp_error( $fetched ) ) {
			return $fetched;
		}
		// Force PNG mime — background-remover outputs always have alpha.
		$fetched['mime'] = 'image/png';
		return $fetched;
	}
}
