<?php

declare(strict_types=1);

namespace WordPress\CloudflareWorkersAiProvider\Embeddings;

use function WordPress\CloudflareWorkersAiProvider\get_plugin_settings;

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

/**
 * What embedding the current selection would take, before anything is sent:
 * counts from the database, sizes from extracting and chunking a random sample
 * (no API calls), scaled up to the whole selection.
 */
final class Estimate
{
    private const TRANSIENT = 'aipcf_embed_estimate';

    private const SAMPLE = 150;

    /** Vectorize storage included per plan, in stored dimensions (vectors x dims). */
    public const VECTORIZE_FREE_DIMS = 5000000;

    public const VECTORIZE_PAID_DIMS = 10000000;

    /**
     * @return array{objects: int, by_kind: array<string, int>, sampled: int, skipped: int, short: int, chunks: int, tokens: int, words: int, neurons: float, days: int, dims: int, stored_dims: int, fits: string, model: string, computed_at: int}
     */
    public static function get(bool $refresh = false): array
    {
        $key    = self::cacheKey();
        $cached = $refresh ? false : get_transient(self::TRANSIENT);
        if (is_array($cached) && ($cached['key'] ?? '') === $key && isset($cached['data'])) {
            return $cached['data'];
        }
        $data = self::compute();
        set_transient(self::TRANSIENT, ['key' => $key, 'data' => $data], 15 * MINUTE_IN_SECONDS);
        return $data;
    }

    public static function forget(): void
    {
        delete_transient(self::TRANSIENT);
    }

    /**
     * @return array{objects: int, by_kind: array<string, int>, sampled: int, skipped: int, short: int, chunks: int, tokens: int, words: int, neurons: float, days: int, dims: int, stored_dims: int, fits: string, model: string, computed_at: int}
     */
    private static function compute(): array
    {
        $model   = Models::current();
        $info    = Models::get($model) ?? ['dims' => 0, 'max_tokens' => 512];
        $counts  = Selection::counts();
        $objects = array_sum($counts);

        $sampled  = 0;
        $skipped  = 0;
        $short    = 0;
        $chunks   = 0;
        $tokens   = 0;
        $words    = 0;
        $minWords = Settings::minWords();
        foreach ($counts as $kind => $count) {
            if ($count === 0) {
                continue;
            }
            $take = max(1, (int) round(self::SAMPLE * $count / max(1, $objects)));
            $q    = Selection::query(['posts_per_page' => $take, 'orderby' => 'rand', 'no_found_rows' => true, 'fields' => 'ids'] + Selection::queryArgs($kind));
            foreach ($q->posts as $id) {
                $post = get_post((int) $id);
                if (!$post instanceof \WP_Post) {
                    continue;
                }
                $sampled++;
                if (!Selection::isEligible($post)) {
                    $skipped++;
                    continue;
                }
                $doc = Extractor::extract($post);
                $n   = Extractor::words($doc['body']);
                if ($post->post_type !== 'attachment' && $minWords > 0 && $n < $minWords) {
                    $skipped++;
                    $short++;
                    continue;
                }
                $pieces = Chunker::chunk($doc['title'], $doc['body'], $info['max_tokens']);
                if ($pieces === []) {
                    $skipped++;
                    continue;
                }
                $chunks += count($pieces);
                $tokens += (int) array_sum(array_column($pieces, 'tokens'));
                $words  += $n;
            }
        }

        $scale       = $sampled > 0 ? $objects / $sampled : 0.0;
        $totalChunks = (int) round($chunks * $scale);
        $totalTokens = (int) round($tokens * $scale);
        $neurons     = Models::neuronsFor($model, $totalTokens);
        $cap         = Settings::dailyNeurons();
        $storedDims  = $totalChunks * (int) $info['dims'];

        $fits = 'unknown';
        if ((string) (get_plugin_settings()['embed_store'] ?? 'vectorize') === 'vectorize') {
            $fits = $storedDims <= self::VECTORIZE_FREE_DIMS ? 'free' : ($storedDims <= self::VECTORIZE_PAID_DIMS ? 'paid' : 'paid_extra');
        }

        return [
            'objects'     => $objects,
            'by_kind'     => $counts,
            'sampled'     => $sampled,
            'skipped'     => (int) round($skipped * $scale),
            'short'       => (int) round($short * $scale),
            'chunks'      => $totalChunks,
            'tokens'      => $totalTokens,
            'words'       => (int) round($words * $scale),
            'neurons'     => round($neurons, 1),
            'days'        => $cap > 0 ? (int) max(1, ceil($neurons / $cap)) : 1,
            'dims'        => (int) $info['dims'],
            'stored_dims' => $storedDims,
            'fits'        => $fits,
            'model'       => $model,
            'computed_at' => time(),
        ];
    }

    /**
     * Changes whenever a setting that affects the estimate changes.
     */
    private static function cacheKey(): string
    {
        $s = get_plugin_settings();
        $relevant = [];
        foreach (Settings::KEYS as $k) {
            $relevant[$k] = $s[$k] ?? null;
        }
        return md5((string) wp_json_encode($relevant) . '|' . Chunker::VERSION);
    }
}