<?php

// 防止直接访问
if (!defined('ABSPATH')) {
    exit; // 如果直接访问则退出
}

// Remove AJAX handlers for cloud login and disconnect
/*
add_action('wp_ajax_handle_cloud_login', 'story321_handle_cloud_login');
function story321_handle_cloud_login() { ... }

add_action('wp_ajax_handle_cloud_disconnect', 'story321_handle_cloud_disconnect');
function story321_handle_cloud_disconnect() { ... }
*/

add_action('rest_api_init', function () {
    register_rest_route('story321/v1', '/generate_text', array(
        'methods' => 'POST',
        'callback' => 'story321_handle_ai_generation',
        'permission_callback' => function () {
            // Check if API key is set before allowing access, in addition to user capability
            return current_user_can('edit_posts') && !empty(get_option('story321_api_key'));
        },
        'args' => array(
            'prompt' => array(
                'required' => true,
                'validate_callback' => function ($param) {
                    return !empty($param);
                }
            ),
            // Add stream argument definition
            'stream' => array(
                'required' => false,
                'type' => 'boolean',
                'default' => false,
            )
        )
    ));
});

function story321_handle_ai_generation($request)
{
    // Log received parameters for debugging
    error_log('[AI Generator Debug] Received generate_text request. Params: ' . print_r($request->get_params(), true));

    $api_key = get_option('story321_api_key');
    if (empty($api_key)) {
        error_log('[AI Generator Debug] API Key missing.'); // Log missing API key
        return new WP_Error('no_api_key', __('API Key is missing. Please configure it in the plugin settings.', 'ai-generator'), array('status' => 403));
    }

    // --- WordPress Response Streaming Setup ---
    // Decide if the *WordPress response* should attempt to stream based on the request param
    $should_wp_stream = $request->get_param('stream');

    if ($should_wp_stream) {
        // 设置纯文本流响应头 (给调用 WordPress API 的客户端)
        // 注意：wp_remote_post 通常会等待完整响应，
        // 真正的流式处理可能需要更底层的方法，但这会尽力模拟。
        header('Content-Type: text/plain; charset=utf-8');
        header('Cache-Control: no-cache');
        header('X-Accel-Buffering: no'); // Important for Nginx buffering
    }

    // 获取请求参数
    $prompt = $request->get_param('prompt');
    if (empty($prompt)) {
         return new WP_Error('missing_prompt', __('Prompt parameter is required.', 'ai-generator'), array('status' => 400));
    }

    // 设置请求参数 (发送给 Next.js API)
    $url = 'https://studio.story321.com/api/tm/generate-text-new'; // Ensure this is the correct endpoint
    $args = [
        'headers' => [
            'Content-Type' => 'application/json',
            'Authorization' => 'Bearer ' . $api_key, // Use API Key for Authorization
            'Accept' => 'text/plain', // Explicitly accept plain text stream
        ],
        'body' => json_encode([
            'prompt' => $prompt,
            // 移除了 'stream' 参数，因为 Next.js API 总是流式返回
            // 移除了 token, memberPoint, point - 认证由 API Key 处理
        ]),
        'timeout' => 120, // Timeout for the entire request/response cycle
        'stream' => false, // Keep this false for wp_remote_post to collect the full body
        'filename' => null // Ensure filename is null if stream is false
        // 'connect_timeout' is implicitly handled within 'timeout' in newer WP versions,
        // but you can set 'connect_timeout' => 20 explicitly if needed for older versions.
    ];

    // 发送请求
    $response = wp_remote_post($url, $args);

    // --- 处理来自 Next.js API 的响应 ---

    if (is_wp_error($response)) {
        $error_message = $response->get_error_message();
        error_log('[AI Generator Debug] wp_remote_post failed: ' . $error_message);
        if ($should_wp_stream) {
            status_header(500);
            echo "Error: Failed to connect to the AI service. " . esc_html($error_message);
            exit;
        } else {
            return new WP_Error('api_connection_error', __('Failed to connect to the AI service:', 'ai-generator') . ' ' . esc_html($error_message), array('status' => 500));
        }
    }

    $status_code = wp_remote_retrieve_response_code($response);
    $body = wp_remote_retrieve_body($response);
    $content_type = wp_remote_retrieve_header($response, 'content-type');

    // Log API response details
    error_log('[AI Generator Debug] API Response Status: ' . $status_code);
    // error_log('[AI Generator Debug] API Response Body: ' . substr($body, 0, 500) . (strlen($body) > 500 ? '...' : '')); // Log partial body

    // Handle API errors (non-2xx status codes)
    if ($status_code >= 400) {
        error_log('[AI Generator Debug] API returned error status ' . $status_code . '. Raw body: ' . $body);
        $error_message = $body; // Assume error message might be plain text or simple JSON
        try {
            $error_data = json_decode($body, true);
            if (json_last_error() === JSON_ERROR_NONE && isset($error_data['error'])) {
                $error_message = $error_data['error']; // Use 'error' field if valid JSON with 'error' key
            }
        } catch (Exception $e) {
             // Ignore JSON decode error, use raw body
        }

        if ($should_wp_stream) {
            // Instead of echoing plain text, send a JSON error response
            // consistent with how WP_Error is handled by the REST API.
            status_header($status_code);
            // Ensure correct content type for JSON response
            header('Content-Type: application/json; charset=' . get_option('blog_charset'));
            echo wp_json_encode([
                'code' => 'api_error',
                'message' => __('API Error:', 'ai-generator') . ' (' . $status_code . ') ' . esc_html($error_message),
                'data' => ['status' => $status_code]
            ]);
            exit;
        } else {
            return new WP_Error('api_error', __('API Error:', 'ai-generator') . ' (' . $status_code . ') ' . esc_html($error_message), array('status' => $status_code));
        }
    }

    // --- 处理成功响应 (2xx status code) ---
    // The Next.js API streams plain text chunks. wp_remote_post collects them into $body.

    if ($should_wp_stream) {
        // If WordPress response should stream, echo the collected body and exit.
        // Note: This isn't true streaming chunk-by-chunk from WP, but sends the full result
        // after collecting it, under the streaming headers set earlier.
        echo $body;
        exit;
    } else {
        // If WordPress response should be a standard JSON response,
        // wrap the plain text result in a JSON structure.
        // Or potentially return plain text directly if the client expects that.
        // Let's wrap it for now.
        return new WP_REST_Response(
            ['generated_text' => $body], // Wrap plain text in a JSON object
            200,
            ['Content-Type' => 'application/json'] // Set correct content type for JSON
        );
        // Alternatively, to return plain text directly:
        // return new WP_REST_Response($body, 200, ['Content-Type' => 'text/plain']);
    }
}


/* function story321_handle_ai_generation($request)
{
    // Log received parameters for debugging
    error_log('[AI Generator Debug] Received generate_text request. Params: ' . print_r($request->get_params(), true));

    $api_key = get_option('story321_api_key');
    if (empty($api_key)) {
        error_log('[AI Generator Debug] API Key missing.'); // Log missing API key
        return new WP_Error('no_api_key', __('API Key is missing. Please configure it in the plugin settings.', 'ai-generator'), array('status' => 403));
    }

    if ($request->get_param('stream')) {
        // 设置纯文本流响应头
        header('Content-Type: text/plain; charset=utf-8');
        header('Cache-Control: no-cache');
        header('X-Accel-Buffering: no');
    }

    // 获取请求参数
    $prompt = $request->get_param('prompt');
    $stream = $request->get_param('stream'); // Get the stream parameter

    // 设置请求参数
    $url = 'https://studio.story321.com/api/tm/generate-text-new'; // Ensure this is the correct endpoint
    $args = [
        'headers' => [
            'Content-Type' => 'application/json',
            'Authorization' => 'Bearer ' . $api_key, // Use API Key for Authorization
        ],
        'body' => json_encode([
            'prompt' => $prompt,
            'stream' => $stream, // Pass stream preference to API
            // Remove token, memberPoint, point as they are likely handled by API key authentication server-side
        ]),
        'timeout' => 120,
        'connect_timeout' => 20 // Explicitly set connection timeout to 20 seconds
    ];

    // 发送请求
    $response = wp_remote_post($url, $args);

    if (is_wp_error($response)) {
        if ($stream) {
            status_header(500);
            echo "Error: " . esc_html($response->get_error_message());
            exit;
        } else {
            return new WP_Error('api_error', __('Failed to connect to the AI service:', 'ai-generator') . ' ' . esc_html($response->get_error_message()), array('status' => 500));
        }
    }

    $status_code = wp_remote_retrieve_response_code($response);
    $body = wp_remote_retrieve_body($response);

    // Handle potential API errors returned in the body
    if ($status_code >= 400) {
        // Log the raw error body received from the API
        error_log('[AI Generator Debug] API returned error status ' . $status_code . '. Raw body: ' . $body);

        $error_data = json_decode($body, true);
        $error_message = isset($error_data['message']) ? $error_data['message'] : __('Unknown API error', 'ai-generator');
         if ($stream) {
            status_header($status_code);
            echo "API Error ($status_code): " . esc_html($error_message);
            exit;
        } else {
            return new WP_Error('api_error', __('API Error:', 'ai-generator') . ' ' . esc_html($error_message), array('status' => $status_code));
        }
    }

    // If streaming, output the body directly and exit
    if ($stream) {
        echo $body;
        exit;
    }

    // If not streaming, assume JSON response (though the current stream implementation doesn't seem to return JSON)
    // This part might need adjustment based on the non-stream API response format
    $decoded_body = json_decode($body, true);
     if (json_last_error() !== JSON_ERROR_NONE) {
        // If the body is not valid JSON (and not streaming), return it as plain text or handle as an error
        return new WP_REST_Response($body, 200, ['Content-Type' => 'text/plain']);
    }
    return rest_ensure_response($decoded_body);
} */


// 注册图片生成API端点
add_action('rest_api_init', function () {
    register_rest_route('story321/v1', '/generate-text2image', array(
        'methods' => 'POST',
        'callback' => 'story321_handle_generate_image',
        'permission_callback' => function () {
            // Check if API key is set before allowing access
             return current_user_can('edit_posts') && !empty(get_option('story321_api_key'));
        },
        'args' => array( // Define expected arguments
            'prompt' => array(
                'required' => true,
                'type' => 'string',
                'description' => __('The prompt for image generation.', 'ai-generator'),
                'validate_callback' => function($param) {
                    return !empty(trim($param));
                }
            ),
            // Optional: Add aspectRatio if you want to support it from WordPress side
            'aspectRatio' => array(
                'required' => false, // Make it optional
                'type' => 'string',
                'description' => __('The aspect ratio for image generation (e.g., "1:1", "16:9").', 'ai-generator'),
                'validate_callback' => function($param) {
                    // Basic validation for format like "number:number"
                    return empty($param) || (is_string($param) && preg_match('/^\d+:\d+$/', $param));
                },
                 'default' => '1:1', // Set a default value
            ),
        )
    ));
});

// 处理图片生成请求
function story321_handle_generate_image($request)
{
    $api_key = get_option('story321_api_key');
     if (empty($api_key)) {
        return new WP_Error('no_api_key', __('API Key is missing. Please configure it in the plugin settings.', 'ai-generator'), array('status' => 403, 'data' => ['success' => false]));
    }

    // 获取请求参数
    $prompt = $request->get_param('prompt');
    $aspectRatio = $request->get_param('aspectRatio'); // Get aspectRatio, defaults to '1:1' if not provided or invalid due to 'default' in args

    // Basic prompt validation (already handled by register_rest_route args, but keep for safety)
    if (empty($prompt)) {
        return new WP_Error('empty_prompt', __('Prompt cannot be empty.', 'ai-generator'), array('status' => 400, 'data' => ['success' => false]));
    }

    // 构建发送到API的请求体
    $api_request_body = ['prompt' => $prompt];
    // Add aspectRatio if provided and valid (it should be due to validation/default)
    if (!empty($aspectRatio)) {
        $api_request_body['aspectRatio'] = $aspectRatio;
    }

    // 发送请求到外部API
    $url = 'https://studio.story321.com/api/tm/generate-text2image-new'; // Ensure this is the correct endpoint
    $response = wp_remote_post($url, array(
        'headers' => array(
            'Content-Type' => 'application/json',
            'Authorization' => 'Bearer ' . $api_key, // Use API Key
            'Accept' => 'application/json'
        ),
        'body' => json_encode($api_request_body), // Use the body with prompt and aspectRatio
        'timeout' => 120, // Timeout for the entire request/response cycle
        'stream' => false, // Keep this false for wp_remote_post to collect the full body
        'filename' => null // Ensure filename is null if stream is false

    ));

    if (is_wp_error($response)) {
        $error_message = $response->get_error_message();
        error_log('[AI Generator Image Debug] wp_remote_post failed: ' . $error_message);
        return new WP_Error('api_connection_error', __('Failed to connect to the image generation service:', 'ai-generator') . ' ' . esc_html($error_message), array('status' => 500, 'data' => ['success' => false]));
    } 

    $status_code = wp_remote_retrieve_response_code($response);
    $body = wp_remote_retrieve_body($response);

    // Log API response details
    error_log('[AI Generator Image Debug] API Response Status: ' . $status_code);
    // error_log('[AI Generator Image Debug] API Response Body: ' . substr($body, 0, 500) . (strlen($body) > 500 ? '...' : '')); // Keep partial log for brevity in general logs

    // Handle API errors (non-2xx status codes)
    if ($status_code >= 400) {
         error_log('[AI Generator Image Debug] API returned error status ' . $status_code . '. Raw body: ' . $body);
         $error_message = $body; // Default to raw body
         try {
             $error_data = json_decode($body, true);
             if (json_last_error() === JSON_ERROR_NONE) {
                 if (isset($error_data['error'])) {
                     $error_message = $error_data['error']; // Prefer 'error' key
                 } elseif (isset($error_data['message'])) {
                     $error_message = $error_data['message']; // Fallback to 'message'
                 }
             }
         } catch (Exception $e) {
             // Ignore JSON errors, use raw body
         }
        // Use the status code from the API response if available and seems appropriate (4xx/5xx)
        $wp_error_status = ($status_code >= 400 && $status_code < 600) ? $status_code : 500;
        return new WP_Error('api_error', __('API Error:', 'ai-generator') . ' (' . $status_code . ') ' . esc_html($error_message), array('status' => $wp_error_status, 'data' => ['success' => false]));
    }

    // Handle successful response (2xx status code)
    // ADDED: Log the full raw body for successful status codes before decoding attempt
    error_log('[AI Generator Image Debug] Raw successful API Response Body (' . $status_code . '): ' . $body);

    $decoded_body = json_decode($body, true);

    // Check if JSON decoding failed OR if 'success' is not explicitly true OR 'image' field is missing/empty
    if (json_last_error() !== JSON_ERROR_NONE || !isset($decoded_body['success']) || $decoded_body['success'] !== true || empty($decoded_body['image'])) {
        error_log('[AI Generator Image Debug] Invalid or unsuccessful API response body: ' . $body);
        $error_message = __('Image generation failed: The API returned an unexpected or unsuccessful response.', 'ai-generator');
        // Try to get a more specific error from the response if possible
         if (isset($decoded_body['error'])) {
             $error_message = $decoded_body['error'];
         } elseif(isset($decoded_body['message'])) {
             $error_message = $decoded_body['message'];
         }
        return new WP_Error('invalid_api_response', esc_html($error_message), array('status' => 500, 'data' => ['success' => false]));
    }

    // --- Process the successful response with base64 image data ---
    $imageBase64 = $decoded_body['image']; // Get base64 data from 'image' field

    // Decode base64 image data (Handle potential Data URL format, although API likely sends raw base64)
    if (strpos($imageBase64, 'data:image/') === 0) {
        // Extract base64 part from Data URL if present
        $image_data_base64_only = explode(',', $imageBase64, 2)[1];
    } else {
        // Assume it's direct base64 data
        $image_data_base64_only = $imageBase64;
    }

    $image_data = base64_decode($image_data_base64_only, true); // Use strict mode
    if ($image_data === false) {
        error_log('[AI Generator Image Debug] Failed to decode base64 image data.');
        return new WP_Error('decode_error', __('Failed to decode image data from API response.', 'ai-generator'), array('status' => 500, 'data' => ['success' => false]));
    }

    // Validate image data (basic check)
    if (!getimagesizefromstring($image_data)) {
        error_log('[AI Generator Image Debug] Decoded data is not a valid image.');
        return new WP_Error('invalid_image_data', __('The data received from the API is not a valid image.', 'ai-generator'), array('status' => 500, 'data' => ['success' => false]));
    }

    // --- Save image to WordPress Media Library ---
    $upload_dir = wp_upload_dir();
    if (!empty($upload_dir['error'])) {
         error_log('[AI Generator Image Debug] WordPress upload directory error: ' . $upload_dir['error']);
         return new WP_Error('upload_dir_error', __('Could not get WordPress upload directory:', 'ai-generator') . ' ' . esc_html($upload_dir['error']), array('status' => 500, 'data' => ['success' => false]));
    }

    // Determine file extension (try to infer from decoded data, default to png)
    // Note: This is a basic inference. The API *should* ideally provide the format.
    $finfo = finfo_open();
    $mime_type = finfo_buffer($finfo, $image_data, FILEINFO_MIME_TYPE);
    finfo_close($finfo);
    $extension = 'png'; // Default
    $allowed_extensions = ['jpg' => 'image/jpeg', 'jpeg' => 'image/jpeg', 'png' => 'image/png', 'gif' => 'image/gif', 'webp' => 'image/webp'];
    $found_ext = array_search($mime_type, $allowed_extensions);
    if ($found_ext !== false) {
        $extension = $found_ext;
    }

    $filename = sanitize_file_name('ai-generated-' . time() . '.' . $extension);
    $filepath = $upload_dir['path'] . '/' . $filename;

    // Save the image file
    if (file_put_contents($filepath, $image_data) === false) {
        error_log('[AI Generator Image Debug] Failed to save image file to: ' . $filepath);
        return new WP_Error('save_image_error', __('Failed to save the generated image to the server.', 'ai-generator'), array('status' => 500, 'data' => ['success' => false]));
    }

    // --- Create WordPress Attachment ---
    $wp_filetype = wp_check_filetype($filename, null);

    $attachment = array(
        'guid'           => $upload_dir['url'] . '/' . basename($filename),
        'post_mime_type' => $wp_filetype['type'] ?: $mime_type, // Use inferred mime type if check fails
        'post_title'     => $prompt, // Use prompt as title
        'post_content'   => '', // Description can be added if needed
        'post_status'    => 'inherit'
    );

    $attach_id = wp_insert_attachment($attachment, $filepath);

    if (is_wp_error($attach_id)) {
        error_log('[AI Generator Image Debug] Failed to create attachment. Error: ' . $attach_id->get_error_message());
        // Clean up saved file if attachment creation fails
        unlink($filepath);
        return new WP_Error('attachment_error', __('Failed to create attachment for the generated image.', 'ai-generator'), array('status' => 500, 'data' => ['success' => false]));
    }

    // Include necessary file for image metadata generation
    require_once(ABSPATH . 'wp-admin/includes/image.php');

    // Generate attachment metadata
    $attach_data = wp_generate_attachment_metadata($attach_id, $filepath);
    if (is_wp_error($attach_data)) {
        // Log error, but don't fail the whole request as the image is saved and attached
        error_log('[AI Generator Image Debug] Failed to generate attachment metadata. Error: ' . $attach_data->get_error_message());
    } else {
        wp_update_attachment_metadata($attach_id, $attach_data);
    }

    // --- Return success response with attachment details ---
    // This matches the structure expected by the generate.js frontend
    return rest_ensure_response(array(
        'success' => true,
        'data' => array(
            'id' => $attach_id,
            'url' => wp_get_attachment_url($attach_id) // Get the URL of the full-size image
        )
    ));
}

// Add permission callback check for API key in generate_text route (already done above)

// Add permission callback check for API key in generate_text2image route (already done above)
// ... existing code ...

// Function to fetch user points from the API
function story321_get_user_points() {
    $api_key = get_option('story321_api_key');
    if (empty($api_key)) {
        error_log('[AI Generator Points Debug] API Key missing, cannot fetch points.');
        return false; // Return false specifically for missing API key
    }

    $url = 'https://studio.story321.com/api/tm/get-points'; // Assuming this is the correct endpoint
    $args = [
        'method'    => 'POST', // As per the Next.js code provided
        'headers'   => [
            'Authorization' => 'Bearer ' . $api_key,
            'Content-Type'  => 'application/json',
            'Accept'        => 'application/json'
        ],
        'timeout'   => 30, // Shorter timeout for fetching points
        'body'      => null // No body needed for this request based on the API code
    ];

    error_log('[AI Generator Points Debug] Calling points API: ' . $url);
    $response = wp_remote_post($url, $args);

    if (is_wp_error($response)) {
        $error_message = $response->get_error_message();
        error_log('[AI Generator Points Debug] wp_remote_post failed: ' . $error_message);
        return null; // Return null on connection error
    }

    $status_code = wp_remote_retrieve_response_code($response);
    $body = wp_remote_retrieve_body($response);
    error_log('[AI Generator Points Debug] API Response Status: ' . $status_code);
    // error_log('[AI Generator Points Debug] API Response Body: ' . $body);

    if ($status_code >= 400) {
        error_log('[AI Generator Points Debug] API returned error status ' . $status_code . '. Body: ' . $body);
        return null; // Return null on API error status
    }

    $decoded_body = json_decode($body, true);
error_log('[AI Generator Points Debug] Decoded Body: ' . print_r($decoded_body, true));
    if (json_last_error() !== JSON_ERROR_NONE || !isset($decoded_body['success']) || $decoded_body['success'] !== true || !isset($decoded_body['points'])) {
        error_log('[AI Generator Points Debug] Invalid or unsuccessful API response body: ' . $body);
        return null; // Return null if response is not valid JSON or indicates failure or missing points
    }

    error_log('[AI Generator Points Debug] Successfully fetched points: ' . $decoded_body['points']);
    // Return the 'points' value. Change to 'memberPoints' or handle both if needed.
    return $decoded_body;
}