<?php
/**
 * Replicate API handler class
 *
 * @package ReplicateWP_Virtual_TryOn
 */

// Exit if accessed directly
if ( ! defined( 'ABSPATH' ) ) {
    exit;
}

/**
 * Replicate API handler class
 */
class RWPVTO_Replicate_API {
    
    /**
     * API endpoint
     *
     * @var string
     */
    private $api_endpoint = 'https://api.replicate.com/v1/predictions';
    
    /**
     * API key
     *
     * @var string
     */
    private $api_key;
    
    /**
     * Model version
     *
     * @var string
     */
    private $model_version;
    
    /**
     * Constructor
     */
    public function __construct() {
        $settings = get_option( 'rwpvto_settings', array() );
        $this->api_key = isset( $settings['api_key'] ) ? $settings['api_key'] : '';
        $this->model_version = isset( $settings['model_version'] ) ? $settings['model_version'] : 'cuuupid/idm-vton:0513734a452173b8173e907e3a59d19a36266e55b48528559432bd21c7d7e985';
    }
    
    /**
     * Generate virtual try-on image
     *
     * @param array $params Generation parameters
     * @return array|WP_Error
     */
    public function generate_tryon( $params ) {
        // Validate API key
        if ( empty( $this->api_key ) ) {
            return new WP_Error(
                'no_api_key',
                __( 'Replicate API key is not configured. Please add it in settings.', 'ai-virtual-try-on-for-woocommerce' )
            );
        }
        
        // Prepare default parameters
        $defaults = array(
            'garm_img' => '',
            'human_img' => '',
            'garment_des' => '',
            'category' => 'upper_body',
            'crop' => false,
            'seed' => 42,
            'steps' => 30,
            'force_dc' => false,
            'mask_only' => false,
        );
        
        $input_params = wp_parse_args( $params, $defaults );
        
        // Validate required parameters
        if ( empty( $input_params['garm_img'] ) ) {
            return new WP_Error(
                'missing_garment_image',
                __( 'Garment image URL is required.', 'ai-virtual-try-on-for-woocommerce' )
            );
        }
        
        if ( empty( $input_params['human_img'] ) ) {
            return new WP_Error(
                'missing_model_image',
                __( 'Model image URL is required.', 'ai-virtual-try-on-for-woocommerce' )
            );
        }
        
        // Prepare request body
        $body = array(
            'version' => $this->model_version,
            'input' => $input_params,
        );
        
        // Make API request
        $response = wp_remote_post(
            $this->api_endpoint,
            array(
                'headers' => array(
                    'Authorization' => 'Bearer ' . $this->api_key,
                    'Content-Type' => 'application/json',
                ),
                'body' => wp_json_encode( $body ),
                'timeout' => 30, // Initial request timeout
            )
        );
        
        // Check for errors
        if ( is_wp_error( $response ) ) {
            return $response;
        }
        
        $response_code = wp_remote_retrieve_response_code( $response );
        $response_body = wp_remote_retrieve_body( $response );
        $data = json_decode( $response_body, true );
        
        // Handle API errors (but 202 is valid - it means processing)
        if ( $response_code !== 200 && $response_code !== 201 && $response_code !== 202 ) {
            $error_message = isset( $data['detail'] ) ? $data['detail'] : __( 'Unknown API error occurred.', 'ai-virtual-try-on-for-woocommerce' );
            
            return new WP_Error(
                'api_error',
                sprintf(
                    /* translators: %1$d: HTTP status code, %2$s: Error message */
                    __( 'Replicate API error (HTTP %1$d): %2$s', 'ai-virtual-try-on-for-woocommerce' ),
                    $response_code,
                    $error_message
                )
            );
        }
        
        // If 202, the API is processing asynchronously - poll for results
        if ( $response_code === 202 && isset( $data['id'] ) ) {
            $prediction_id = $data['id'];
            $max_attempts = 150; // 150 attempts = 5 minutes max
            $attempt = 0;
            
            while ( $attempt < $max_attempts ) {
                sleep( 2 ); // Wait 2 seconds between polls
                
                $status_result = $this->get_prediction_status( $prediction_id );
                
                if ( is_wp_error( $status_result ) ) {
                    return $status_result;
                }
                
                // Check if completed
                if ( isset( $status_result['status'] ) && $status_result['status'] === 'succeeded' ) {
                    if ( isset( $status_result['output'] ) && ! empty( $status_result['output'] ) ) {
                        return array(
                            'success' => true,
                            'prediction_id' => $prediction_id,
                            'status' => 'succeeded',
                            'output_url' => is_array( $status_result['output'] ) ? $status_result['output'][0] : $status_result['output'],
                            'raw_response' => $status_result,
                        );
                    } else {
                        // Status is succeeded but no output - something went wrong
                        return new WP_Error(
                            'no_output',
                            __( 'Generation completed but no output was returned. Response: ', 'ai-virtual-try-on-for-woocommerce' ) . wp_json_encode( $status_result )
                        );
                    }
                }
                
                // Check if failed
                if ( isset( $status_result['status'] ) && $status_result['status'] === 'failed' ) {
                    $error_msg = isset( $status_result['error'] ) ? $status_result['error'] : __( 'Generation failed', 'ai-virtual-try-on-for-woocommerce' );
                    return new WP_Error( 'generation_failed', $error_msg );
                }
                
                // Check if canceled
                if ( isset( $status_result['status'] ) && $status_result['status'] === 'canceled' ) {
                    return new WP_Error( 'canceled', __( 'Generation was canceled.', 'ai-virtual-try-on-for-woocommerce' ) );
                }
                
                $attempt++;
            }
            
            // Timeout
            return new WP_Error(
                'timeout',
                __( 'Generation timed out after 5 minutes. The image may still be processing. Check your Replicate dashboard.', 'ai-virtual-try-on-for-woocommerce' )
            );
        }
        
        // For 200/201 responses, check if it's actually still processing
        if ( isset( $data['status'] ) && in_array( $data['status'], array( 'starting', 'processing' ) ) && isset( $data['id'] ) ) {
            // The API returned 200/201 but it's still processing - poll for results
            $prediction_id = $data['id'];
            $max_attempts = 150; // 150 attempts = 5 minutes max
            $attempt = 0;
            
            while ( $attempt < $max_attempts ) {
                sleep( 2 );
                
                $status_result = $this->get_prediction_status( $prediction_id );
                
                if ( is_wp_error( $status_result ) ) {
                    return $status_result;
                }
                
                // Check if completed
                if ( isset( $status_result['status'] ) && $status_result['status'] === 'succeeded' ) {
                    if ( isset( $status_result['output'] ) && ! empty( $status_result['output'] ) ) {
                        return array(
                            'success' => true,
                            'prediction_id' => $prediction_id,
                            'status' => 'succeeded',
                            'output_url' => is_array( $status_result['output'] ) ? $status_result['output'][0] : $status_result['output'],
                            'raw_response' => $status_result,
                        );
                    }
                }
                
                // Check if failed
                if ( isset( $status_result['status'] ) && $status_result['status'] === 'failed' ) {
                    $error_msg = isset( $status_result['error'] ) ? $status_result['error'] : __( 'Generation failed', 'ai-virtual-try-on-for-woocommerce' );
                    return new WP_Error( 'generation_failed', $error_msg );
                }
                
                // Check if canceled
                if ( isset( $status_result['status'] ) && $status_result['status'] === 'canceled' ) {
                    return new WP_Error( 'canceled', __( 'Generation was canceled.', 'ai-virtual-try-on-for-woocommerce' ) );
                }
                
                $attempt++;
            }
            
            // Timeout
            return new WP_Error(
                'timeout',
                __( 'Generation timed out after 5 minutes. The image may still be processing. Check your Replicate dashboard.', 'ai-virtual-try-on-for-woocommerce' )
            );
        }
        
        // Output should be available immediately for completed status
        if ( ! isset( $data['output'] ) || empty( $data['output'] ) ) {
            return new WP_Error(
                'invalid_response',
                sprintf(
                    /* translators: %s: JSON encoded response */
                    __( 'Invalid response from Replicate API. No output generated. Status: %s', 'ai-virtual-try-on-for-woocommerce' ),
                    isset( $data['status'] ) ? $data['status'] : 'unknown'
                )
            );
        }
        
        return array(
            'success' => true,
            'prediction_id' => isset( $data['id'] ) ? $data['id'] : null,
            'status' => isset( $data['status'] ) ? $data['status'] : 'unknown',
            'output_url' => is_array( $data['output'] ) ? $data['output'][0] : $data['output'],
            'raw_response' => $data,
        );
    }
    
    /**
     * Get prediction status
     *
     * @param string $prediction_id Prediction ID
     * @return array|WP_Error
     */
    public function get_prediction_status( $prediction_id ) {
        if ( empty( $this->api_key ) ) {
            return new WP_Error(
                'no_api_key',
                __( 'Replicate API key is not configured.', 'ai-virtual-try-on-for-woocommerce' )
            );
        }
        
        $response = wp_remote_get(
            $this->api_endpoint . '/' . $prediction_id,
            array(
                'headers' => array(
                    'Authorization' => 'Bearer ' . $this->api_key,
                    'Content-Type' => 'application/json',
                ),
                'timeout' => 30,
            )
        );
        
        if ( is_wp_error( $response ) ) {
            return $response;
        }
        
        $response_code = wp_remote_retrieve_response_code( $response );
        $response_body = wp_remote_retrieve_body( $response );
        $data = json_decode( $response_body, true );
        
        if ( $response_code !== 200 ) {
            return new WP_Error(
                'api_error',
                __( 'Failed to retrieve prediction status.', 'ai-virtual-try-on-for-woocommerce' )
            );
        }
        
        return $data;
    }
    
    /**
     * Test API connection
     *
     * @return bool|WP_Error
     */
    public function test_connection() {
        if ( empty( $this->api_key ) ) {
            return new WP_Error(
                'no_api_key',
                __( 'Replicate API key is not configured.', 'ai-virtual-try-on-for-woocommerce' )
            );
        }
        
        // Make a simple request to test the API key
        $response = wp_remote_get(
            'https://api.replicate.com/v1/predictions',
            array(
                'headers' => array(
                    'Authorization' => 'Bearer ' . $this->api_key,
                    'Content-Type' => 'application/json',
                ),
                'timeout' => 15,
            )
        );
        
        if ( is_wp_error( $response ) ) {
            return new WP_Error(
                'connection_error',
                sprintf(
                    /* translators: %s: Error message */
                    __( 'Connection failed: %s', 'ai-virtual-try-on-for-woocommerce' ),
                    $response->get_error_message()
                )
            );
        }
        
        $response_code = wp_remote_retrieve_response_code( $response );
        $response_body = wp_remote_retrieve_body( $response );
        
        if ( $response_code === 200 ) {
            return true;
        }
        
        // Try to get more specific error message
        $data = json_decode( $response_body, true );
        $error_detail = '';
        
        if ( isset( $data['detail'] ) ) {
            $error_detail = ' - ' . $data['detail'];
        }
        
        return new WP_Error(
            'api_error',
            sprintf(
                /* translators: %1$d: HTTP status code, %2$s: Error details */
                __( 'API connection failed (HTTP %1$d)%2$s. Please check your API key.', 'ai-virtual-try-on-for-woocommerce' ),
                $response_code,
                $error_detail
            )
        );
    }
}

