<?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\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\RequestOptions;
use WordPress\AiClient\Providers\Http\DTO\Response;
use WordPress\AiClient\Providers\Http\Enums\HttpMethodEnum;
use WordPress\AiClient\Providers\Models\TextGeneration\Contracts\TextGenerationModelInterface;
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;

use function WordPress\CloudflareWorkersAiProvider\fetch_catalog;

class CloudflareWorkersAiTranslationModel extends AbstractApiBasedModel implements TextGenerationModelInterface
{
    public const MODELS = [
        '@cf/meta/m2m100-1.2b',
        '@cf/ai4bharat/indictrans2-en-indic-1B',
    ];

    // M2M100's list price per 1K tokens, used when a response has tokens but no neurons (IndicTrans2).
    private const DEFAULT_NEURONS_PER_1K = 31.1;

    private const LANGUAGES = [
        'en' => ['en', 'eng_Latn'], 'english' => ['en', 'eng_Latn'],
        'hi' => ['hi', 'hin_Deva'], 'hindi' => ['hi', 'hin_Deva'],
        'mr' => ['mr', 'mar_Deva'], 'marathi' => ['mr', 'mar_Deva'],
        'gu' => ['gu', 'guj_Gujr'], 'gujarati' => ['gu', 'guj_Gujr'],
        'bn' => ['bn', 'ben_Beng'], 'bengali' => ['bn', 'ben_Beng'], 'bangla' => ['bn', 'ben_Beng'],
        'ta' => ['ta', 'tam_Taml'], 'tamil' => ['ta', 'tam_Taml'],
        'te' => ['te', 'tel_Telu'], 'telugu' => ['te', 'tel_Telu'],
        'kn' => ['kn', 'kan_Knda'], 'kannada' => ['kn', 'kan_Knda'],
        'ml' => ['ml', 'mal_Mlym'], 'malayalam' => ['ml', 'mal_Mlym'],
        'pa' => ['pa', 'pan_Guru'], 'punjabi' => ['pa', 'pan_Guru'],
        'or' => ['or', 'ory_Orya'], 'odia' => ['or', 'ory_Orya'], 'oriya' => ['or', 'ory_Orya'],
        'as' => ['as', 'asm_Beng'], 'assamese' => ['as', 'asm_Beng'],
        'ur' => ['ur', 'urd_Arab'], 'urdu' => ['ur', 'urd_Arab'],
        'ne' => ['ne', 'npi_Deva'], 'nepali' => ['ne', 'npi_Deva'],
        'sa' => ['sa', 'san_Deva'], 'sanskrit' => ['sa', 'san_Deva'],
        'sd' => ['sd', 'snd_Deva'], 'sindhi' => ['sd', 'snd_Deva'],
        'ks' => ['ks', 'kas_Deva'], 'kashmiri' => ['ks', 'kas_Deva'],
        'mai' => ['mai', 'mai_Deva'], 'maithili' => ['mai', 'mai_Deva'],
        'kok' => ['kok', 'gom_Deva'], 'konkani' => ['kok', 'gom_Deva'],
        'mni' => ['mni', 'mni_Beng'], 'manipuri' => ['mni', 'mni_Beng'],
        'sat' => ['sat', 'sat_Olck'], 'santali' => ['sat', 'sat_Olck'],
        'doi' => ['doi', 'doi_Deva'], 'dogri' => ['doi', 'doi_Deva'],
        'brx' => ['brx', 'brx_Deva'], 'bodo' => ['brx', 'brx_Deva'],
        'fr' => ['fr', null], 'french' => ['fr', null],
        'de' => ['de', null], 'german' => ['de', null],
        'es' => ['es', null], 'spanish' => ['es', null],
        'pt' => ['pt', null], 'portuguese' => ['pt', null],
        'it' => ['it', null], 'italian' => ['it', null],
        'ru' => ['ru', null], 'russian' => ['ru', null],
        'ja' => ['ja', null], 'japanese' => ['ja', null],
        'zh' => ['zh', null], 'chinese' => ['zh', null],
        'ar' => ['ar', null], 'arabic' => ['ar', null],
        'ko' => ['ko', null], 'korean' => ['ko', null],
        'tr' => ['tr', null], 'turkish' => ['tr', null],
        'nl' => ['nl', null], 'dutch' => ['nl', null],
        'id' => ['id', null], 'indonesian' => ['id', null],
        'vi' => ['vi', null], 'vietnamese' => ['vi', null],
        'th' => ['th', null], 'thai' => ['th', null],
        'sw' => ['sw', null], 'swahili' => ['sw', null],
    ];

    public static function supports(string $modelId): bool
    {
        return in_array($modelId, self::MODELS, true);
    }

    public static function languageCode(string $lang, string $modelId): ?string
    {
        $key = strtolower(trim($lang));
        if ($key === '') {
            return null;
        }
        $indic = strpos($modelId, 'indictrans2') !== false;

        // Already a model-native code.
        if ($indic && preg_match('/^[a-z]{3}_[A-Z][a-z]{3}$/', trim($lang))) {
            return trim($lang);
        }
        if (isset(self::LANGUAGES[$key])) {
            return $indic ? self::LANGUAGES[$key][1] : self::LANGUAGES[$key][0];
        }
        // Reverse lookup of an IndicTrans2 code when M2M100 is targeted.
        foreach (self::LANGUAGES as $entry) {
            if ($entry[1] !== null && strcasecmp($entry[1], $key) === 0) {
                return $indic ? $entry[1] : $entry[0];
            }
        }
        return $indic ? null : $key; // M2M100 accepts many more ISO codes than we list.
    }

    public static function indicLanguages(): array
    {
        $out = [];
        foreach (self::LANGUAGES as $alias => $entry) {
            if ($entry[1] !== null && strlen($alias) > 3) {
                $out[] = $alias;
            }
        }
        return $out;
    }

    public function getHttpTransporter(): HttpTransporterInterface
    {
        $modelId = $this->metadata()->getId();
        $rate    = self::DEFAULT_NEURONS_PER_1K;
        foreach (fetch_catalog() as $m) {
            if (($m['id'] ?? '') === $modelId && (float) ($m['price_in'] ?? 0) > 0) {
                $rate = UsageTracker::neuronsPer1k((float) $m['price_in']);
                break;
            }
        }
        $rate = (float) apply_filters('aipcf_translation_neurons_per_1k', $rate, $modelId);
        return new CloudflareTransporter(parent::getHttpTransporter(), $modelId, 0.0, max(0.0, $rate));
    }

    public function generateTextResult(array $prompt): GenerativeAiResult
    {
        // No cheaper translation model to fall back to, so both budget actions pause here.
        if (UsageTracker::overBudget()) {
            throw new RuntimeException(
                __('The daily budget is used up. Translation 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 RequestOptions();
            $options->setTimeout((float) apply_filters('aipcf_request_timeout', 60.0, $modelId));
        }

        $headers = ['Content-Type' => 'application/json']
            + Gateway::headers(TaskContext::current(), true, $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);
    }

    protected function prepareParams(array $prompt): array
    {
        $text = '';
        for ($i = count($prompt) - 1; $i >= 0; $i--) {
            if (!$prompt[$i]->getRole()->isUser()) {
                continue;
            }
            foreach ($prompt[$i]->getParts() as $part) {
                if ($part->getType()->isText()) {
                    $text .= (string) $part->getText();
                }
            }
            break;
        }
        $text = trim($text);
        if ($text === '') {
            throw new InvalidArgumentException(__('Nothing to translate: the prompt has no user text.', 'aipcf-ai-provider-for-cloudflare'));
        }

        $modelId = $this->metadata()->getId();
        $custom  = $this->getConfig()->getCustomOptions();
        $target  = (string) ($custom['target_lang'] ?? $custom['target_language'] ?? $custom['to'] ?? '');
        $source  = (string) ($custom['source_lang'] ?? $custom['source_language'] ?? $custom['from'] ?? 'en');

        if ($target === '') {
            throw new InvalidArgumentException(
                __('Translation needs a target language: set the custom option target_lang (for example "mr") on the model config, or use aipcf_translate().', 'aipcf-ai-provider-for-cloudflare')
            );
        }

        $targetCode = self::languageCode($target, $modelId);
        if ($targetCode === null) {
            throw new InvalidArgumentException(sprintf('Model %s cannot translate into "%s".', $modelId, $target));
        }

        if (strpos($modelId, 'indictrans2') !== false) {
            // English > Indic only; the source is implicit.
            return ['text' => $text, 'target_language' => $targetCode];
        }

        $sourceCode = self::languageCode($source, $modelId) ?? 'en';
        return ['text' => $text, 'source_lang' => $sourceCode, 'target_lang' => $targetCode];
    }

    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 (preg_match('/"([^"]+)" is not one of/', $msg, $m)) {
            $msg = sprintf('language "%s" is not supported by this model', $m[1]);
        } elseif (strlen($msg) > 240) {
            $msg = substr($msg, 0, 240) . '...';
        }
        throw new RuntimeException(sprintf(
            'Cloudflare Workers AI translation failed (HTTP %d)%s',
            $response->getStatusCode(),
            $msg !== '' ? ': ' . $msg : ''
        ));
    }

    protected function parseResponse(Response $response): GenerativeAiResult
    {
        $data   = $response->getData();
        $result = is_array($data) && isset($data['result']) && is_array($data['result']) ? $data['result'] : [];

        $text = null;
        if (isset($result['translated_text']) && is_string($result['translated_text'])) {
            $text = $result['translated_text'];
        } elseif (isset($result['translations'][0]) && is_string($result['translations'][0])) {
            $text = $result['translations'][0];
        }
        if ($text === null) {
            throw new RuntimeException('Cloudflare Workers AI translation response contained no text.');
        }

        $usage = isset($result['usage']) && is_array($result['usage']) ? $result['usage'] : [];
        $in    = (int) ($usage['prompt_tokens'] ?? 0);
        $out   = (int) ($usage['completion_tokens'] ?? 0);

        $message = new Message(MessageRoleEnum::model(), [new MessagePart($text)]);

        return new GenerativeAiResult(
            '',
            [new Candidate($message, FinishReasonEnum::stop())],
            new TokenUsage($in, $out, (int) ($usage['total_tokens'] ?? $in + $out)),
            $this->providerMetadata(),
            $this->metadata(),
            []
        );
    }
}