<?php
declare(strict_types=1);
// ==================================================
// Nuclear Fancy Pants Edition™ -- NUCLEAR+++
// File: /home/ploi/automarketer-ai.com/public/wp-content/plugins/automarketer-ai/modules/img-sd.php
// Description: Stable Diffusion (Stability AI) image provider module -- NUCLEAR+++
// Last Edited: 2026-05-25
// ==================================================

namespace AutoMarketerAI\Modules;

use AutoMarketerAI\ImageProviderInterface;
use AutoMarketerAI\ImageProvider;

defined('ABSPATH') || exit;

final class ImgSd implements ImageProviderInterface
{
    private const API_BASE      = 'https://api.stability.ai/v2beta/stable-image/generate';
    private const PREFIX        = 'sk-';
    private const DEFAULT_MODEL = 'core';

    // Model slugs map to API endpoint suffixes.
    private const MODEL_ENDPOINTS = [
        'core'  => 'core',
        'ultra' => 'ultra',
        'sd3'   => 'sd3',
    ];

    // Per-model cost strings shown on the card chip.
    private const MODEL_COSTS = [
        'core'  => '$0.03 / image',
        'ultra' => '$0.08 / image',
        'sd3'   => '$0.065 / image',
    ];

    public function get_id(): string
    {
        return 'img-sd';
    }

    public function get_name(): string
    {
        return 'Stable Diffusion';
    }

    public function get_description(): string
    {
        return 'Stability AI\'s Stable Diffusion generates diverse, high-quality marketing visuals with strong creative range and competitive pricing.';
    }

    public function get_logo_url(): string
    {
        return '';
    }

    public function get_category(): string
    {
        return 'Image';
    }

    public function get_model_count(): int
    {
        return count(self::MODEL_ENDPOINTS);
    }

    public function get_available_models(): array
    {
        return [
            'core'  => 'Stable Image Core -- $0.03 / image',
            'ultra' => 'Stable Image Ultra -- $0.08 / image',
            'sd3'   => 'Stable Diffusion 3.5 -- $0.065 / image',
        ];
    }

    public function get_single_model_notice(): string
    {
        return '';
    }

    public function validate_api_key(string $key): bool
    {
        if (!str_starts_with($key, self::PREFIX)) {
            return false;
        }

        // Validate by hitting the user account balance endpoint.
        $response = wp_remote_get('https://api.stability.ai/v1/user/balance', [
            'headers' => [
                'Authorization' => 'Bearer ' . $key,
                'Accept'        => 'application/json',
            ],
            'timeout' => 10,
        ]);

        if (is_wp_error($response)) {
            return false;
        }

        $code = (int) wp_remote_retrieve_response_code($response);
        return $code === 200;
    }

    public function get_api_key_prefix(): string
    {
        return self::PREFIX;
    }

    public function get_default_model(): string
    {
        return self::DEFAULT_MODEL;
    }

    public function render_activation_popup(): void
    {
        // Not used -- activation modal rendered by shared template.
    }

    public function register(): void
    {
        // No additional hooks needed.
    }

    public function generate_raw(string $prompt, string $api_key = '', string $model = ''): string
    {
        // Image providers do not implement generate_raw.
        return '';
    }

    // --------------------------------------------------
    // Generate an image using the Stability AI REST API.
    // Sends a multipart/form-data request and receives raw
    // image bytes directly in the response body.
    // Returns a local attachment URL on success, null on failure.
    // --------------------------------------------------
    public function generate_image(string $prompt, string $api_key): ?string
    {
        if ($prompt === '' || $api_key === '') {
            return null;
        }

        $active_model = (string) get_option('amai_model_img-sd', self::DEFAULT_MODEL);
        $endpoint     = self::MODEL_ENDPOINTS[$active_model] ?? self::MODEL_ENDPOINTS[self::DEFAULT_MODEL];
        $safe_prompt  = mb_substr($prompt, 0, 10000);

        $boundary = wp_generate_password(24, false);

        $body  = '--' . $boundary . "\r\n";
        $body .= 'Content-Disposition: form-data; name="prompt"' . "\r\n\r\n";
        $body .= $safe_prompt . "\r\n";
        $body .= '--' . $boundary . "\r\n";
        $body .= 'Content-Disposition: form-data; name="output_format"' . "\r\n\r\n";
        $body .= 'png' . "\r\n";

        // SD3 requires mode field.
        if ($endpoint === 'sd3') {
            $body .= '--' . $boundary . "\r\n";
            $body .= 'Content-Disposition: form-data; name="mode"' . "\r\n\r\n";
            $body .= 'text-to-image' . "\r\n";
            $body .= '--' . $boundary . "\r\n";
            $body .= 'Content-Disposition: form-data; name="model"' . "\r\n\r\n";
            $body .= 'sd3.5-large' . "\r\n";
        }

        $body .= '--' . $boundary . '--' . "\r\n";

        $response = wp_remote_post(self::API_BASE . '/' . $endpoint, [
            'headers' => [
                'Authorization' => 'Bearer ' . $api_key,
                'Accept'        => 'image/*',
                'Content-Type'  => 'multipart/form-data; boundary=' . $boundary,
            ],
            'body'    => $body,
            'timeout' => 60,
        ]);

        if (is_wp_error($response)) {
            return null;
        }

        $code = (int) wp_remote_retrieve_response_code($response);

        if ($code !== 200) {
            return null;
        }

        // Check for content filter rejection.
        $finish_reason = wp_remote_retrieve_header($response, 'finish-reason');
        if ($finish_reason === 'CONTENT_FILTERED') {
            return null;
        }

        $bytes = wp_remote_retrieve_body($response);
        if ($bytes === '') {
            return null;
        }

        return ImageProvider::saveBytesToMediaLibrary($bytes, 'png', 'amai-sd');
    }
}