<?php

declare(strict_types=1);

namespace WordPress\CloudflareWorkersAiProvider\Embeddings;

use WordPress\CloudflareWorkersAiProvider\Embeddings\Stores\StoreFactory;
use WordPress\CloudflareWorkersAiProvider\Usage\UsageTracker;

use const WordPress\CloudflareWorkersAiProvider\AIPCF_SETTINGS_OPTION;

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

/**
 * Embeds site content into a vector store.
 *
 * ## EXAMPLES
 *
 *     wp aipcf embed status
 *     wp aipcf embed estimate
 *     wp aipcf embed run --until-done
 *     wp aipcf embed search "how much does it cost"
 */
final class Cli
{
    public static function register(): void
    {
        if (defined('WP_CLI') && WP_CLI && class_exists('\WP_CLI')) {
            \WP_CLI::add_command('aipcf embed', self::class);
        }
    }

    /**
     * Shows what is indexed, pending and failed, and the store's state. Makes no changes.
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function status(array $args, array $assoc): void
    {
        $state = State::get();
        $c     = IndexTable::counts(Settings::signature());
        \WP_CLI::log(sprintf('Embeddings: %s. Model: %s. Vector store: %s.', Settings::enabled() ? 'on' : 'off', Models::current(), (string) (StoreFactory::current() instanceof Stores\VectorStore ? StoreFactory::current()->indexName() : '-')));
        \WP_CLI::log(sprintf(
            '%d embedded (%s) | %d waiting (%d pending, %d stale, %d deleting) | %d failed | %d no text (skipped)',
            $c['indexed'],
            self::plural($c['chunks'], 'passage', 'passages'),
            $c['pending'] + $c['stale'] + $c['deleting'],
            $c['pending'],
            $c['stale'],
            $c['deleting'],
            $c['failed'],
            $c['skipped']
        ));
        $limit = Settings::dailyNeurons();
        \WP_CLI::log(sprintf('Embedding neurons today: about %.1f%s', UsageTracker::todayEmbedNeurons(), $limit > 0 ? sprintf(' of the daily limit of %d', $limit) : ' (no daily limit of its own)'));
        if ($state['reconcile'] !== null) {
            \WP_CLI::log(($state['reconcile']['phase'] ?? '') === 'add'
                ? sprintf('Re-checking the selection (page %d).', (int) ($state['reconcile']['page'] ?? 1))
                : 'Re-checking the selection (looking for items to delete).');
        }
        if ($state['paused'] !== '') {
            \WP_CLI::warning('Paused: ' . $state['paused']);
        }
        if ($state['last_run'] > 0) {
            \WP_CLI::log(sprintf('Last run: %s ago. %s', human_time_diff($state['last_run']), $state['last_message']));
        }
        if (isset($assoc['store'])) {
            $store = StoreFactory::current();
            $s     = is_wp_error($store) ? $store : $store->status();
            \WP_CLI::log('Vector store: ' . (is_wp_error($s) ? $s->get_error_message() : self::describeStatus($s)));
        }
    }

    /**
     * Estimates chunks, tokens and neurons for the current selection. Makes no API calls.
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function estimate(array $args, array $assoc): void
    {
        \WordPress\CloudflareWorkersAiProvider\Activity\Recorder::start('embeddings', 'embed_estimate');
        $e = Estimate::get(true);
        $kinds = [];
        foreach ($e['by_kind'] as $kind => $n) {
            $kinds[] = $kind . ' ' . $n;
        }
        $fits = [
            'free'       => 'fits the Vectorize free plan',
            'paid'       => 'too big for the Vectorize free plan, but fits Workers Paid',
            'paid_extra' => 'needs Workers Paid plus extra storage',
        ];
        \WP_CLI::log(sprintf(
            "%s%s\nAbout %s passages (parts of posts), %s tokens, %.1f neurons%s\nStorage: about %s dimensions%s",
            self::plural($e['objects'], 'item', 'items'),
            $kinds !== [] ? ' (' . implode(', ', $kinds) . ')' : '',
            number_format($e['chunks']),
            number_format($e['tokens']),
            $e['neurons'],
            Settings::dailyNeurons() > 0 ? sprintf(' (about %s at the daily limit)', self::plural($e['days'], 'day', 'days')) : '',
            number_format($e['stored_dims']),
            isset($fits[$e['fits']]) ? '; ' . $fits[$e['fits']] . '.' : '.'
        ));
    }

    /**
     * Runs the worker.
     *
     * ## OPTIONS
     *
     * [--until-done]
     * : Keep running batches until nothing is left, a budget is reached or it pauses.
     *
     * [--max-neurons=<n>]
     * : Stop once this run has used about this many embedding neurons.
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function run(array $args, array $assoc): void
    {
        \WordPress\CloudflareWorkersAiProvider\Activity\Recorder::start('embeddings', 'embed_run');
        $this->requireEnabled();
        $untilDone = isset($assoc['until-done']);
        $max       = isset($assoc['max-neurons']) ? max(1, (int) $assoc['max-neurons']) : 0;
        $start     = UsageTracker::todayEmbedNeurons();
        $totals    = ['embedded' => 0, 'unchanged' => 0, 'deleted' => 0, 'failed' => 0];
        do {
            $r = Ingest::tick(true, 25);
            foreach ($totals as $k => $v) {
                $totals[$k] = $v + (int) $r[$k];
            }
            $c = IndexTable::counts(Settings::signature());
            \WP_CLI::log(sprintf(
                'embedded %d, unchanged %d, removed %d, failed %d | %d waiting (%d pending, %d stale, %d deleting)%s',
                $r['embedded'],
                $r['unchanged'],
                $r['deleted'],
                $r['failed'],
                $c['pending'] + $c['stale'] + $c['deleting'],
                $c['pending'],
                $c['stale'],
                $c['deleting'],
                $r['message'] !== '' ? ' | ' . $r['message'] : ($r['reason'] !== '' ? ' | ' . $r['reason'] : '')
            ));
            $left     = $c['pending'] + $c['stale'] + $c['deleting'];
            $progress = $r['embedded'] + $r['unchanged'] + $r['deleted'] + $r['failed'] + $r['copied'] + $r['requeued'] > 0;
            $spent    = UsageTracker::todayEmbedNeurons() - $start;
            if (!$r['ran'] && $untilDone && $r['reason'] === Ingest::busyReason()) {
                // WP-Cron started a batch (e.g. after a save); wait for it rather than give up.
                \WP_CLI::log('Another batch is running. Waiting 10 s...');
                sleep(10);
                continue;
            }
            if (!$r['ran'] || State::get()['paused'] !== '' || ($max > 0 && $spent >= $max)) {
                break;
            }
            if (!$progress && $r['message'] !== '') {
                // Waiting on the index, a budget or a temporary error.
                $store = StoreFactory::current();
                if (!$untilDone || !$store instanceof Stores\VectorStore || $r['message'] !== Ingest::preparingMessage($store->indexName())) {
                    break;
                }
                \WP_CLI::log('Waiting 30 s for the index...');
                sleep(30);
            }
        } while ($untilDone && ($left > 0 || State::get()['reconcile'] !== null));

        \WP_CLI::success(sprintf(
            'Embedded %d, unchanged %d, removed %d, failed %d. About %.1f embedding neurons used.',
            $totals['embedded'],
            $totals['unchanged'],
            $totals['deleted'],
            $totals['failed'],
            UsageTracker::todayEmbedNeurons() - $start
        ));
    }

    /**
     * Moves the vectors to another store without embedding them again (no neurons), and runs
     * the move to the end. Related posts and search keep using the old store until it is done.
     *
     * ## OPTIONS
     *
     * [<store>]
     * : vectorize, qdrant or pgvector. Its connection must already be set on Settings > Content index.
     *
     * [--index=<name>]
     * : The index, collection or table in the new store. Default: one per site.
     *
     * [--status]
     * : Show how far a move is.
     *
     * [--cancel]
     * : Cancel the move under way: what was copied goes back, also without neurons.
     *
     * [--yes]
     * : Do not ask for confirmation.
     *
     * ## EXAMPLES
     *
     *     wp aipcf embed move qdrant
     *     wp aipcf embed move --status
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function move(array $args, array $assoc): void
    {
        $move = State::get()['move'];
        if (isset($assoc['status'])) {
            if ($move === null) {
                \WP_CLI::log('No move is under way.');
                return;
            }
            $left = IndexTable::signatureCounts(Settings::signatureOf((string) $move['from'], (string) $move['from_index']))['vectors'];
            \WP_CLI::log(sprintf('%s %s -> %s %s: %d of %d vectors copied, %d posts to embed again%s.', $move['from'], $move['from_index'], $move['to'], $move['to_index'], max(0, (int) $move['vectors'] - $left), (int) $move['vectors'], (int) $move['requeued'], ($move['serve'] ?? 'from') === 'to' ? ' (going back after a cancel)' : ''));
            return;
        }
        if (isset($assoc['cancel'])) {
            \WP_CLI::success(Mover::cancel() ? 'Move cancelled. Run `wp aipcf embed run --until-done` to copy back what was copied.' : 'No move is under way.');
            return;
        }
        $this->requireEnabled();
        if ($move !== null) {
            \WP_CLI::error('A move is under way already. See `wp aipcf embed move --status`, or `--cancel` it.');
        }
        $to = (string) ($args[0] ?? '');
        if (!isset(StoreFactory::available()[$to])) {
            \WP_CLI::error('Name the store to move to: vectorize, qdrant or pgvector.');
        }
        $settings                = get_option(AIPCF_SETTINGS_OPTION, []);
        $settings                = is_array($settings) ? $settings : [];
        $settings['embed_store'] = $to;
        if (isset($assoc['index'])) {
            $settings[['vectorize' => 'vectorize_index', 'qdrant' => 'qdrant_collection', 'pgvector' => 'pg_table'][$to]] = (string) $assoc['index'];
        }
        $c = IndexTable::counts(Settings::signature());
        \WP_CLI::confirm(sprintf('Move %d posts (%d vectors) from %s to %s? They are copied, without neurons.', $c['indexed'], $c['chunks'], StoreFactory::currentId(), $to), $assoc);
        update_option(AIPCF_SETTINGS_OPTION, $settings);
        if (!Mover::active()) {
            \WP_CLI::success('Nothing to copy: the store is set, and the worker embeds what is selected as usual.');
            return;
        }
        do {
            $r = Ingest::tick(true, 25);
            \WP_CLI::log(sprintf('copied %d vectors, %d posts to embed again%s', $r['copied'], $r['requeued'], $r['message'] !== '' ? ' | ' . $r['message'] : ''));
            if (Mover::active() && $r['copied'] + $r['requeued'] === 0) {
                \WP_CLI::log('Waiting 30 s (the new store is getting ready, or Vectorize is settling)...');
                sleep(30);
            }
        } while (Mover::active());
        \WP_CLI::success('Moved. Related posts and search use the new store. Delete what is left in the old one with `wp aipcf embed cleanup-old`.');
    }

    /**
     * Deletes what a move left in the store it came from: the whole index when it was this
     * site's own, or only this site's vectors from a shared one.
     *
     * ## OPTIONS
     *
     * [--store=<store>]
     * : Another leftover than the last move's: vectorize, qdrant or pgvector.
     *
     * [--index=<name>]
     * : With --store: the index, collection or table name.
     *
     * [--drop]
     * : With --store: drop the whole index (only when it is this site's alone).
     *
     * [--yes]
     * : Do not ask for confirmation.
     *
     * @subcommand cleanup-old
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function cleanupOld(array $args, array $assoc): void
    {
        $left  = State::get()['moved_from'];
        $store = isset($assoc['store']) ? (string) $assoc['store'] : ($left['store'] ?? null);
        $index = isset($assoc['store']) ? (string) ($assoc['index'] ?? '') : (string) ($left['index'] ?? '');
        if ($store === null || $index === '') {
            \WP_CLI::error('No store is left over from a move. Name one with --store and --index.');
        }
        $drop = isset($assoc['store']) ? isset($assoc['drop']) : !empty($left['own']);
        \WP_CLI::confirm(sprintf($drop ? 'Drop the index %2$s in %1$s?' : 'Delete this site\'s vectors from %2$s in %1$s?', $store, $index), $assoc);
        $r = isset($assoc['store']) ? Mover::cleanup($store, $index, $drop) : Mover::cleanup();
        if (is_wp_error($r)) {
            \WP_CLI::error($r->get_error_message());
        }
        \WP_CLI::success('Deleted.');
    }

    /**
     * Re-checks the whole selection against the table (adds missing items, removes unselected ones).
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function reconcile(array $args, array $assoc): void
    {
        \WordPress\CloudflareWorkersAiProvider\Activity\Recorder::start('embeddings', 'embed_reconcile');
        $this->requireEnabled();
        IndexTable::install();
        Ingest::startReconcile();
        while (State::get()['reconcile'] !== null) {
            Ingest::reconcileStep(microtime(true) + 30);
        }
        $c = IndexTable::counts(Settings::signature());
        \WP_CLI::success(sprintf('Selection re-checked: %d waiting to embed, %d waiting to delete. Run `wp aipcf embed run` to apply.', $c['pending'], $c['deleting']));
    }

    /**
     * Retries failed and skipped items.
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function retry(array $args, array $assoc): void
    {
        \WordPress\CloudflareWorkersAiProvider\Activity\Recorder::start('embeddings', 'embed_retry');
        \WP_CLI::success(self::plural(IndexTable::retryFailed(), 'item', 'items') . ' queued again.');
    }

    /**
     * Starts over in a new index; the old one is dropped when the new one is complete.
     *
     * ## OPTIONS
     *
     * [--yes]
     * : Do not ask for confirmation.
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function rebuild(array $args, array $assoc): void
    {
        \WordPress\CloudflareWorkersAiProvider\Activity\Recorder::start('embeddings', 'embed_rebuild');
        $this->requireEnabled();
        \WP_CLI::confirm('Embed everything again into a new index?', $assoc);
        Ingest::rebuild();
        \WP_CLI::success(sprintf('New index %s. Run `wp aipcf embed run --until-done`.', StoreFactory::current() instanceof Stores\VectorStore ? StoreFactory::current()->indexName() : ''));
    }

    /**
     * Switches embeddings off and removes this site's vectors from the store.
     *
     * ## OPTIONS
     *
     * [--yes]
     * : Do not ask for confirmation.
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function purge(array $args, array $assoc): void
    {
        \WordPress\CloudflareWorkersAiProvider\Activity\Recorder::start('embeddings', 'embed_purge');
        \WP_CLI::confirm('Switch Embeddings off and delete this site\'s vectors from the vector store? Your content is not touched.', $assoc);
        $r = Ingest::purge();
        if (is_wp_error($r)) {
            \WP_CLI::error($r->get_error_message());
        }
        \WP_CLI::success('Vectors deleted; Embeddings are off.');
    }

    /**
     * Semantic search over this site's vectors.
     *
     * ## OPTIONS
     *
     * <query>
     * : What to look for.
     *
     * [--top=<n>]
     * : How many results.
     * ---
     * default: 5
     * ---
     *
     * @param list<string>          $args
     * @param array<string, string> $assoc
     */
    public function search(array $args, array $assoc): void
    {
        $r = Ingest::search((string) ($args[0] ?? ''), (int) ($assoc['top'] ?? 5));
        if (is_wp_error($r)) {
            \WP_CLI::error($r->get_error_message());
        }
        foreach ($r as $i => $m) {
            \WP_CLI::log(sprintf('%d. %d%% match  #%d %s%s  %s', $i + 1, (int) round($m['score'] * 100), $m['object_id'], $m['title'], $m['section'] !== '' ? ' > ' . $m['section'] : '', $m['url']));
        }
        if ($r === []) {
            \WP_CLI::log('No matches (a fresh write takes about a minute to become searchable).');
        }
    }

    private function requireEnabled(): void
    {
        if (!Settings::enabled()) {
            \WP_CLI::error('The content index is switched off on this site (AIPCF > Settings > Content index).');
        }
    }

    /**
     * "1 item", "3 items" (WP-CLI output is English only).
     */
    private static function plural(int $n, string $one, string $many): string
    {
        return number_format($n) . ' ' . ($n === 1 ? $one : $many);
    }

    /**
     * The vector store's status in words, e.g. "ready, 281 vectors of 1,024 numbers, settled."
     *
     * @param array<string, mixed> $s
     */
    private static function describeStatus(array $s): string
    {
        if (empty($s['exists'])) {
            return 'the index does not exist yet.';
        }
        $line = sprintf(
            '%s, %s of %s numbers, %s.',
            empty($s['ready']) ? 'being prepared' : 'ready',
            self::plural((int) ($s['vector_count'] ?? 0), 'vector', 'vectors'),
            number_format((int) ($s['dims'] ?? 0)),
            empty($s['settled']) ? 'recent changes still being applied' : 'settled'
        );
        if (!empty($s['missing_metadata']) && is_array($s['missing_metadata'])) {
            $line .= ' Missing metadata indexes: ' . implode(', ', array_map('strval', $s['missing_metadata'])) . '.';
        }
        return $line;
    }
}