<?php

declare(strict_types=1);

namespace WordPress\CloudflareWorkersAiProvider\Embeddings;

use WordPress\CloudflareWorkersAiProvider\Usage\UsageTracker;

use function WordPress\CloudflareWorkersAiProvider\get_plugin_settings;

if (!defined('ABSPATH')) {
    exit;
}

/**
 * Workers AI text-embedding models the plugin can index with.
 *
 * The catalog API lists these under "Text Embeddings" with a price but, for the
 * multilingual ones, no dimensions or input limit, so the table is kept here.
 * Embedding responses report no usage, so the price is what neurons are estimated from.
 */
final class Models
{
    public const DEFAULT = '@cf/baai/bge-m3';

    /**
     * @return array<string, array{label: string, dims: int, max_tokens: int, price: float, multilingual: bool}>
     */
    public static function all(): array
    {
        $models = [
            '@cf/baai/bge-m3' => [
                'label'        => 'BGE-M3',
                'dims'         => 1024,
                'max_tokens'   => 8192,
                'price'        => 0.0118,
                'multilingual' => true,
            ],
            '@cf/qwen/qwen3-embedding-0.6b' => [
                'label'        => 'Qwen3 Embedding 0.6B',
                'dims'         => 1024,
                'max_tokens'   => 8192,
                'price'        => 0.0118,
                'multilingual' => true,
            ],
            '@cf/baai/bge-small-en-v1.5' => [
                'label'        => 'BGE small (English)',
                'dims'         => 384,
                'max_tokens'   => 512,
                'price'        => 0.0202,
                'multilingual' => false,
            ],
            '@cf/baai/bge-base-en-v1.5' => [
                'label'        => 'BGE base (English)',
                'dims'         => 768,
                'max_tokens'   => 512,
                'price'        => 0.0666,
                'multilingual' => false,
            ],
            '@cf/baai/bge-large-en-v1.5' => [
                'label'        => 'BGE large (English)',
                'dims'         => 1024,
                'max_tokens'   => 512,
                'price'        => 0.204,
                'multilingual' => false,
            ],
        ];

        /**
         * Filters the embedding models offered for indexing.
         *
         * @param array $models Model id => label, dims, max_tokens, price (USD per M input tokens), multilingual.
         */
        $models = (array) apply_filters('aipcf_embedding_models', $models);

        $clean = [];
        foreach ($models as $id => $m) {
            if (!is_string($id) || !is_array($m) || (int) ($m['dims'] ?? 0) <= 0) {
                continue;
            }
            $clean[$id] = [
                'label'        => (string) ($m['label'] ?? $id),
                'dims'         => (int) $m['dims'],
                'max_tokens'   => max(1, (int) ($m['max_tokens'] ?? 512)),
                'price'        => max(0.0, (float) ($m['price'] ?? 0)),
                'multilingual' => !empty($m['multilingual']),
            ];
        }
        return $clean;
    }

    /**
     * @return array{label: string, dims: int, max_tokens: int, price: float, multilingual: bool}|null
     */
    public static function get(string $modelId): ?array
    {
        return self::all()[$modelId] ?? null;
    }

    /**
     * The model chosen in the settings, or the default when that one is no longer offered.
     */
    public static function current(): string
    {
        $id = (string) (get_plugin_settings()['embed_model'] ?? '');
        return self::get($id) !== null ? $id : self::DEFAULT;
    }

    /**
     * Rough token count: ~4 characters per token for Latin script (M0 measured 4.4 on
     * English, so this errs high), ~2 when most of the text is another script.
     */
    public static function estimateTokens(string $text): int
    {
        $chars = mb_strlen($text);
        if ($chars === 0) {
            return 0;
        }
        $ascii = (int) preg_match_all('/[\x00-\x7F]/', $text);
        $perToken = $ascii / $chars >= 0.7 ? 4 : 2;
        return (int) ceil($chars / $perToken);
    }

    /**
     * Estimated neurons for embedding this many tokens with the model.
     */
    public static function neuronsFor(string $modelId, int $tokens): float
    {
        $model = self::get($modelId);
        if ($model === null || $tokens <= 0) {
            return 0.0;
        }
        return round($tokens / 1000 * UsageTracker::neuronsPer1k($model['price']), 3);
    }
}