<?php

declare(strict_types=1);

// phpcs:disable WordPress.Security.EscapeOutput -- No HTML output in this file; exception messages are not rendered.

namespace WordPress\CloudflareWorkersAiProvider\Models;

use WordPress\AiClient\Common\Exception\InvalidArgumentException;
use WordPress\AiClient\Common\Exception\RuntimeException;
use WordPress\AiClient\Files\DTO\File;
use WordPress\AiClient\Messages\DTO\Message;
use WordPress\AiClient\Messages\DTO\MessagePart;
use WordPress\AiClient\Messages\Enums\MessageRoleEnum;
use WordPress\AiClient\Providers\ApiBasedImplementation\AbstractApiBasedModel;
use WordPress\AiClient\Providers\Http\Contracts\HttpTransporterInterface;
use WordPress\AiClient\Providers\Http\DTO\Request;
use WordPress\AiClient\Providers\Http\DTO\Response;
use WordPress\AiClient\Providers\Http\Enums\HttpMethodEnum;
use WordPress\AiClient\Providers\Models\ImageGeneration\Contracts\ImageGenerationModelInterface;
use WordPress\AiClient\Results\DTO\Candidate;
use WordPress\AiClient\Results\DTO\GenerativeAiResult;
use WordPress\AiClient\Results\DTO\TokenUsage;
use WordPress\AiClient\Results\Enums\FinishReasonEnum;
use WordPress\CloudflareWorkersAiProvider\Gateway;
use WordPress\CloudflareWorkersAiProvider\Http\CloudflareTransporter;
use WordPress\CloudflareWorkersAiProvider\Routing\TaskContext;
use WordPress\CloudflareWorkersAiProvider\Usage\UsageTracker;

class CloudflareWorkersAiImageGenerationModel extends AbstractApiBasedModel implements ImageGenerationModelInterface
{
    private const DEFAULT_TIMEOUT = 90.0;

    // Models that return raw bytes report no usage; FLUX.1 schnell reports ~172 for a 1024px image.
    private const NEURONS_PER_IMAGE = 172.0;

    private const SIZE_CAPABLE = '#/(stable-diffusion|dreamshaper|lucid-origin|phoenix)#i';

    public function getHttpTransporter(): HttpTransporterInterface
    {
        $modelId  = $this->metadata()->getId();
        $estimate = (float) apply_filters('aipcf_image_neurons_estimate', self::NEURONS_PER_IMAGE, $modelId);
        return new CloudflareTransporter(parent::getHttpTransporter(), $modelId, max(0.0, $estimate));
    }

    public function generateImageResult(array $prompt): GenerativeAiResult
    {
        // No cheaper image model to fall back to, so both budget actions pause here.
        if (UsageTracker::overBudget()) {
            throw new RuntimeException(
                __('The daily budget is used up. Image generation resumes at midnight UTC. To allow more, raise the daily budget under AIPCF > Settings > Budget.', 'aipcf-ai-provider-for-cloudflare')
            );
        }

        $modelId = $this->metadata()->getId();
        $params  = $this->prepareParams($prompt);

        $options = $this->getRequestOptions();
        if ($options === null || $options->getTimeout() === null) {
            $options = $options !== null ? clone $options : new \WordPress\AiClient\Providers\Http\DTO\RequestOptions();
            $options->setTimeout((float) apply_filters('aipcf_request_timeout', self::DEFAULT_TIMEOUT, $modelId));
        }

        $headers = ['Content-Type' => 'application/json', 'Accept' => '*/*']
            + Gateway::headers(TaskContext::current(), isset($params['seed']), $modelId);

        $request = new Request(HttpMethodEnum::POST(), Gateway::runUrl($modelId), $headers, $params, $options);
        $request = $this->getRequestAuthentication()->authenticateRequest($request);

        $response = $this->getHttpTransporter()->send($request, $options);
        $this->throwIfNotSuccessful($response);

        return $this->parseResponse($response, $modelId);
    }

    protected function prepareParams(array $prompt): array
    {
        $text = '';
        foreach ($prompt as $message) {
            foreach ($message->getParts() as $part) {
                if ($part->getType()->isText()) {
                    $text .= ($text === '' ? '' : "\n") . (string) $part->getText();
                }
            }
        }
        $text = trim($text);
        if ($text === '') {
            throw new InvalidArgumentException(__('The image prompt must contain text.', 'aipcf-ai-provider-for-cloudflare'));
        }

        $config  = $this->getConfig();
        $modelId = $this->metadata()->getId();
        $params  = ['prompt' => mb_substr($text, 0, 2048)];

        $sys = $config->getSystemInstruction();
        if ($sys) {
            $params['prompt'] = mb_substr(trim($sys) . ' ' . $params['prompt'], 0, 2048);
        }

        if (preg_match(self::SIZE_CAPABLE, $modelId)) {
            [$w, $h] = $this->sizeFor($config->getOutputMediaOrientation(), $config->getOutputMediaAspectRatio());
            $params['width']  = $w;
            $params['height'] = $h;
        }

        foreach ($config->getCustomOptions() as $key => $value) {
            $params[$key] = $value;
        }

        return (array) apply_filters('aipcf_image_params', $params, $modelId);
    }

    private function sizeFor($orientation, ?string $aspect): array
    {
        if ($aspect !== null && preg_match('/^(\d+):(\d+)$/', $aspect, $m)) {
            $ratio = (int) $m[1] / max(1, (int) $m[2]);
            if ($ratio > 1) {
                return [1024, max(512, (int) (round(1024 / $ratio / 64) * 64))];
            }
            if ($ratio < 1) {
                return [max(512, (int) (round(1024 * $ratio / 64) * 64)), 1024];
            }
            return [1024, 1024];
        }
        if ($orientation !== null) {
            if ($orientation->isLandscape()) {
                return [1024, 768];
            }
            if ($orientation->isPortrait()) {
                return [768, 1024];
            }
        }
        return [1024, 1024];
    }

    protected function throwIfNotSuccessful(Response $response): void
    {
        if ($response->isSuccessful()) {
            return;
        }
        $data = $response->getData();
        $msg  = is_array($data) && isset($data['errors'][0]['message']) ? (string) $data['errors'][0]['message'] : '';
        if ($response->getStatusCode() === 429) {
            $msg = $msg !== '' ? $msg : 'rate or spend limit reached (AI Gateway / Cloudflare)';
        }
        throw new RuntimeException(sprintf(
            'Cloudflare Workers AI image request failed (HTTP %d)%s',
            $response->getStatusCode(),
            $msg !== '' ? ': ' . $msg : ''
        ));
    }

    protected function parseResponse(Response $response, string $modelId): GenerativeAiResult
    {
        $body = (string) $response->getBody();
        $data = $response->getData();

        if (is_array($data)) {
            if (isset($data['result']['image']) && is_string($data['result']['image'])) {
                $base64 = $data['result']['image'];
            } elseif (isset($data['result']) && is_string($data['result'])) {
                $base64 = $data['result'];
            } else {
                throw new RuntimeException('Cloudflare Workers AI image response contained no image data.');
            }
            $bytes = base64_decode($base64, true); // phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_decode -- image bytes from the API
            if ($bytes === false) {
                throw new RuntimeException('Cloudflare Workers AI image response was not valid base64.');
            }
        } else {
            $bytes  = $body;
            $base64 = base64_encode($bytes); // phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_encode -- image bytes for the SDK result
        }

        $mime = self::sniffMime($bytes);
        $file = new File('data:' . $mime . ';base64,' . $base64, $mime);

        $message   = new Message(MessageRoleEnum::model(), [new MessagePart($file)]);
        $candidate = new Candidate($message, FinishReasonEnum::stop());

        return new GenerativeAiResult(
            '',
            [$candidate],
            new TokenUsage(0, 0, 0),
            $this->providerMetadata(),
            $this->metadata(),
            ['model' => $modelId, 'bytes' => strlen($bytes)]
        );
    }

    public static function sniffMime(string $bytes): string
    {
        if (strncmp($bytes, "\x89PNG", 4) === 0) {
            return 'image/png';
        }
        if (strncmp($bytes, "\xFF\xD8\xFF", 3) === 0) {
            return 'image/jpeg';
        }
        if (strncmp($bytes, 'RIFF', 4) === 0 && substr($bytes, 8, 4) === 'WEBP') {
            return 'image/webp';
        }
        if (strncmp($bytes, 'GIF8', 4) === 0) {
            return 'image/gif';
        }
        return 'image/png';
    }
}