<?php

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

/**
 * Embeddings for training and search, from the embedding provider chosen on
 * the Model tab. The provider (a RSVL_Provider) makes the requests; this picks
 * the provider and model, splits large batches to fit its limits, checks the
 * vectors that come back, and logs the tokens used for training.
 */
class RSVL_Embedder
{
    private ?RSVL_Provider $provider = null;
    private string $model             = '';

    public function __construct()
    {
        $provider = RSVL_Provider::get_provider((string) resolver_get_setting('embedding_provider'));
        if (! $provider || ! $provider->supports_embeddings() || ! $provider->is_connected()) {
            return;
        }

        $this->provider = $provider;
        $stored         = (string) resolver_get_setting('embedding_model_' . $provider->get_id());
        $this->model    = '' !== $stored ? $stored : $provider->get_latest_embedding_model();
    }

    /** Whether embeddings can be made: a provider is chosen and its API key entered. */
    public function is_available(): bool
    {
        return null !== $this->provider;
    }

    /** The embedding provider's ID, or '' when none is available. */
    public function get_provider(): string
    {
        return $this->provider ? $this->provider->get_id() : '';
    }

    /** The embedding provider's name, e.g. "OpenAI", or '' when none is available. */
    public function get_provider_name(): string
    {
        return $this->provider ? $this->provider->get_name() : '';
    }

    public function get_model(): string
    {
        return $this->model;
    }

    /** Whether the provider has a newer embedding model than the one in use. */
    public function is_upgrade_available(): bool
    {
        return $this->provider && $this->model !== $this->provider->get_latest_embedding_model();
    }

    public function get_latest_model(): string
    {
        return $this->provider ? $this->provider->get_latest_embedding_model() : '';
    }

    /**
     * The date the model in use is deprecated, or false when it isn't.
     *
     * @return string|false Y-m-d.
     */
    public function get_deprecation_date()
    {
        return self::models()[ $this->model ]['deprecated'] ?? false;
    }

    /**
     * The model the one in use should be upgraded to, or false when none.
     *
     * @return string|false
     */
    public function get_upgrade_to()
    {
        return self::models()[ $this->model ]['upgrade_to'] ?? false;
    }

    public function get_dimensions(): int
    {
        return self::models()[ $this->model ]['dimensions'] ?? 1536;
    }

    /**
     * How similar a chunk must be to a visitor's question to be given to the
     * chat model. Each model has its own: scores only compare within one
     * model, whose unrelated texts can score 0.1 or 0.6.
     */
    public function get_similarity_threshold(): float
    {
        return (float) (self::models()[ $this->model ]['similarity_threshold'] ?? 0.35);
    }

    /**
     * Embed one search query: a visitor's message, at chat time. Its tokens
     * aren't logged as training tokens.
     *
     * @return array|WP_Error The vector.
     */
    public function embed(string $text)
    {
        $vectors = $this->embed_all([$this->input($text, null)], 30, false);

        return is_wp_error($vectors) ? $vectors : $vectors[0];
    }

    /**
     * Embed several texts as search queries, the way embed() does a visitor's
     * message, such as the phrases a visitor's message is compared with.
     *
     * @return array|WP_Error Vectors in the same order as $texts.
     */
    public function embed_queries(array $texts)
    {
        if (empty($texts)) {
            return $this->is_available() ? [] : $this->unavailable();
        }

        return $this->embed_all(array_map(fn (string $text) => $this->input($text, null), array_values($texts)), 90, true);
    }

    /**
     * Embed one page's chunks for training, in as few requests as the
     * provider allows.
     *
     * @param string $title The page's title, which each chunk starts with, or '' for none.
     * @return array|WP_Error Vectors in the same order as $texts.
     */
    public function batch_embed(array $texts, string $title = '')
    {
        if (empty($texts)) {
            return $this->is_available() ? [] : $this->unavailable();
        }

        return $this->embed_all(array_map(fn (string $text) => $this->input($text, $title), array_values($texts)), 90, true);
    }

    /**
     * Embed several items' chunks for training, together, in as few requests
     * as the provider allows, then hand each item back its own vectors.
     *
     * @param array $chunk_lists Each item's chunks, e.g. [0 => ['chunk', …], 1 => […]].
     * @param array $titles      Each item's title, keyed as $chunk_lists; '' or missing for none.
     * @return array|WP_Error Each item's vectors, keyed as $chunk_lists, each
     *                        list keyed by chunk position.
     */
    public function embed_chunk_lists(array $chunk_lists, array $titles = [])
    {
        $flat = [];
        $map  = [];
        foreach ($chunk_lists as $item => $chunks) {
            foreach (array_values($chunks) as $position => $chunk) {
                $flat[] = $this->input($chunk, (string) ($titles[ $item ] ?? ''));
                $map[]  = [$item, $position];
            }
        }

        if (empty($flat)) {
            return $this->is_available() ? array_fill_keys(array_keys($chunk_lists), []) : $this->unavailable();
        }

        $vectors = $this->embed_all($flat, 90, true);
        if (is_wp_error($vectors)) {
            return $vectors;
        }

        $by_item = array_fill_keys(array_keys($chunk_lists), []);
        foreach ($map as $flat_index => [$item, $position]) {
            $by_item[ $item ][ $position ] = $vectors[ $flat_index ];
        }

        return $by_item;
    }

    /**
     * @param bool $log_tokens Log the tokens used as training ('index') tokens.
     * @return array|WP_Error
     */
    private function embed_all(array $texts, int $timeout, bool $log_tokens)
    {
        if (! $this->is_available()) {
            return $this->unavailable();
        }

        [$max_inputs, $max_tokens] = $this->provider->get_embedding_batch_limits();

        $results = [];
        foreach ($this->split_into_sub_batches($texts, $max_inputs, $max_tokens) as $batch) {
            $response = $this->provider->embed_texts($batch, $this->model, $timeout);
            if (is_wp_error($response)) {
                return $response;
            }

            $valid = $this->validate_vectors($response['vectors'], count($batch), $this->provider->get_name());
            if (is_wp_error($valid)) {
                return $valid;
            }

            foreach ($response['vectors'] as $vector) {
                $results[] = $vector;
            }

            if ($log_tokens && $response['tokens'] > 0) {
                RSVL_Conversations::log_token_usage('index', $this->provider->get_id(), $this->model, $response['tokens'], 0, $response['tokens']);
            }
        }

        return $results;
    }

    /**
     * A text as the provider's model is sent it: see RSVL_Provider::embedding_input().
     *
     * @param string|null $title Null for a search query; a chunk's title, or '' for none.
     */
    private function input(string $text, ?string $title): string
    {
        return $this->provider ? $this->provider->embedding_input($text, $title, $this->model) : $text;
    }

    private function unavailable(): WP_Error
    {
        return new WP_Error('resolver_no_embedder', __('No embedding provider is configured. Please add an OpenAI or Gemini API key.', 'resolver-ai-chatbot'));
    }

    /**
     * Split texts into sub-batches that fit within input count and token limits.
     * Token count is estimated as ceil(strlen / 4), which matches OpenAI's documented approximation.
     *
     * @param array $texts      Flat array of strings.
     * @param int   $max_inputs Maximum number of texts per sub-batch.
     * @param int   $max_tokens Maximum estimated tokens per sub-batch.
     * @return array Array of sub-arrays, each safe to send in one API call.
     */
    private function split_into_sub_batches(array $texts, int $max_inputs, int $max_tokens): array
    {
        $batches = [];
        $current = [];
        $tokens  = 0;

        foreach ($texts as $text) {
            $text_tokens = (int) ceil(strlen($text) / 4);

            if (! empty($current) &&
                (count($current) >= $max_inputs || $tokens + $text_tokens > $max_tokens)) {
                $batches[] = $current;
                $current   = [];
                $tokens    = 0;
            }

            $current[] = $text;
            $tokens += $text_tokens;
        }

        if (! empty($current)) {
            $batches[] = $current;
        }

        return $batches;
    }

    /**
     * Check the vectors parsed from one API response before anything uses them.
     *
     * pack() silently turns junk values into zeros; a missing vector would
     * shift every later one onto the wrong chunk, because the indexer matches
     * them by position; and a vector of the wrong length breaks
     * cosine_similarity(). Rejecting the response here sends all three
     * through the normal embedding-error handling instead.
     *
     * @param array  $vectors  Vectors keyed by input position.
     * @param int    $expected Number of inputs that were sent.
     * @param string $label    Provider name for the error message.
     * @return true|WP_Error
     */
    private function validate_vectors(array $vectors, int $expected, string $label)
    {
        if (count($vectors) !== $expected) {
            /* translators: 1: embedding provider, e.g. OpenAI. 2: number of embeddings returned. 3: number of texts sent. */
            return new WP_Error('resolver_embed_error', sprintf(__('%1$s returned %2$s embeddings for %3$s inputs.', 'resolver-ai-chatbot'), $label, number_format_i18n(count($vectors)), number_format_i18n($expected)));
        }
        if (array_keys($vectors) !== range(0, $expected - 1)) {
            /* translators: %s: embedding provider, e.g. OpenAI. */
            return new WP_Error('resolver_embed_error', sprintf(__('%s returned embeddings that do not match the inputs sent.', 'resolver-ai-chatbot'), $label));
        }

        // A model missing from models() still has to be self-consistent.
        $first = reset($vectors);
        $dims  = self::models()[ $this->model ]['dimensions'] ?? (is_array($first) ? count($first) : 0);

        foreach ($vectors as $vector) {
            if (! is_array($vector) || 0 === $dims || count($vector) !== $dims) {
                /* translators: 1: embedding provider, e.g. OpenAI. 2: number of values an embedding should have. */
                return new WP_Error('resolver_embed_error', sprintf(__('%1$s returned an embedding of the wrong size (expected %2$s values).', 'resolver-ai-chatbot'), $label, number_format_i18n($dims)));
            }

            // Keys must run 0..n-1, and every value must be a finite number;
            // json_decode() turns 1e999 into INF.
            $invalid = ! array_is_list($vector) || array_filter($vector, fn ($value) => (! is_float($value) && ! is_int($value)) || ! is_finite((float) $value));
            if ($invalid) {
                /* translators: %s: embedding provider, e.g. OpenAI. */
                return new WP_Error('resolver_embed_error', sprintf(__('%s returned an embedding containing invalid values.', 'resolver-ai-chatbot'), $label));
            }
        }

        return true;
    }

    /**
     * Every embedding model, from every provider: model ID => the provider's
     * details for it (see RSVL_Provider::get_embedding_models()), plus
     * 'provider', the provider's ID.
     */
    public static function models(): array
    {
        $models = [];
        foreach (RSVL_Provider::get_all_providers() as $id => $provider) {
            foreach ($provider->get_embedding_models() as $model => $details) {
                $models[ $model ] = $details + ['provider' => $id];
            }
        }

        return $models;
    }

    /**
     * Pack a float array into binary for storage.
     *
     * @param array $vector Array of floats.
     * @return string Binary packed data.
     */
    public static function pack_vector(array $vector): string
    {
        return pack('f*', ...$vector);
    }

    /**
     * Unpack binary data back into a float array.
     *
     * @param string $binary The packed binary data.
     * @return array Array of floats.
     */
    public static function unpack_vector(string $binary): array
    {
        return array_values(unpack('f*', $binary));
    }

    /**
     * Compute cosine similarity between two vectors.
     *
     * @param array $a First vector.
     * @param array $b Second vector.
     * @return float Similarity score between -1 and 1.
     */
    public static function cosine_similarity(array $a, array $b): float
    {
        $dot    = 0.0;
        $norm_a = 0.0;
        $norm_b = 0.0;
        $len    = count($a);

        for ($i = 0; $i < $len; $i++) {
            $dot += $a[ $i ] * $b[ $i ];
            $norm_a += $a[ $i ] * $a[ $i ];
            $norm_b += $b[ $i ] * $b[ $i ];
        }

        if ($norm_a == 0 || $norm_b == 0) {
            return 0.0;
        }

        return $dot / (sqrt($norm_a) * sqrt($norm_b));
    }
}