<?php

declare(strict_types=1);

namespace WordPress\CloudflareWorkersAiProvider\Embeddings;

use WordPress\CloudflareWorkersAiProvider\Activity\Recorder;
use WordPress\CloudflareWorkersAiProvider\Embeddings\Stores\StoreFactory;
use WordPress\CloudflareWorkersAiProvider\Embeddings\Stores\VectorStore;

use const WordPress\CloudflareWorkersAiProvider\AIPCF_SETTINGS_OPTION;

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

/**
 * Moves embeddings to another vector store without embedding them again.
 *
 * The vectors are read from the old store by id; the metadata stored next to them is rebuilt
 * from the post by the same Extractor and Chunker that wrote it; both are written to the new
 * store under the same ids. A post edited since it was embedded, or a vector the old store no
 * longer has, is queued for the worker instead. Until the move is done, readers stay on the
 * old store (StoreFactory::reading()).
 *
 * Cancelling turns the move around: what was already copied is copied back (also free), and
 * readers never leave the old store.
 */
final class Mover
{
    /** Index rows per batch. A row holds up to a few dozen vectors. */
    private const BATCH = 25;

    /** The settings that pick the store and its index; Cancel puts them back. */
    private const STORE_KEYS = ['embed_store', 'vectorize_index', 'qdrant_collection', 'pg_table'];

    /** @var array{id: string, index: string, own: bool, signature: string, model: string, location: string, settings: array<string, mixed>}|null The store before a settings save. */
    private static ?array $before = null;

    public static function boot(): void
    {
        add_filter('pre_update_option_' . AIPCF_SETTINGS_OPTION, [self::class, 'rememberStore'], 10, 2);
    }

    /**
     * Notes which store the vectors are in before a settings save changes it.
     *
     * @param mixed $value
     * @param mixed $old
     * @return mixed
     */
    public static function rememberStore($value, $old = [])
    {
        $store = StoreFactory::current();
        if (!is_wp_error($store)) {
            self::$before = [
                'id'        => $store->id(),
                'index'     => $store->indexName(),
                'own'       => $store->isSiteOwn(),
                'signature' => Settings::signature(),
                'model'     => Models::current(),
                'location'  => self::location($store->id()),
                'settings'  => array_intersect_key(is_array($old) ? $old : [], array_flip(self::STORE_KEYS)),
            ];
        }
        return $value;
    }

    /**
     * After a settings save: a new store or index with the same model is a move. Returns true
     * when a move started (or a move under way was turned around).
     */
    public static function afterSave(): bool
    {
        $before       = self::$before;
        self::$before = null;
        $to           = StoreFactory::current();
        if ($before === null || is_wp_error($to)) {
            return false;
        }
        $same = $before['id'] === $to->id() && $before['index'] === $to->indexName();
        if ($before['model'] !== Models::current()) {
            return false; // new vectors are needed anyway (model change)
        }
        if ($same) {
            if ($before['location'] === self::location($to->id()) || !Settings::enabled() || State::get()['move'] !== null) {
                return false; // nothing moved
            }
            // Same index name on another server (a new Qdrant URL, PostgreSQL host or
            // database): the old server cannot be read with the new settings (its key is not
            // kept), so everything is embedded again into the new one.
            $n = IndexTable::forceReembed();
            if ($n > 0) {
                Recorder::record([
                    'tool'    => 'embeddings',
                    'action'  => 'embed_new_server',
                    'status'  => 'warning',
                    'message' => sprintf(
                        /* translators: %s: number of posts */
                        _n('The vector store now points at another server, so %s post is embedded again there.', 'The vector store now points at another server, so %s posts are embedded again there.', $n, 'aipcf-ai-provider-for-cloudflare'),
                        number_format_i18n($n)
                    ),
                ]);
            }
            return $n > 0;
        }
        $move = State::get()['move'];
        if ($move !== null) {
            // Back to the store the move started from: turn it around (cancel).
            if ($to->id() === $move['from'] && $to->indexName() === $move['from_index']) {
                self::turnAround($move);
                return true;
            }
            return false; // Settings::sanitize keeps the store while a move is under way
        }
        if (!Settings::enabled() || IndexTable::signatureCounts($before['signature'])['rows'] === 0) {
            return false; // nothing to copy: the worker fills the new store as usual
        }
        self::start($before['id'], $before['index'], $before['own'], $before['settings']);
        return true;
    }

    /**
     * @param bool                 $fromOwn      Whether the old index is this site's alone (so cleaning up may drop it).
     * @param array<string, mixed> $fromSettings The store settings that pointed at it, for Cancel.
     */
    public static function start(string $fromId, string $fromIndex, bool $fromOwn, array $fromSettings = []): void
    {
        $to = StoreFactory::current();
        if (is_wp_error($to)) {
            return;
        }
        $vectors = IndexTable::signatureCounts(Settings::signatureOf($fromId, $fromIndex))['vectors'];
        State::update([
            'move'       => [
                'from'       => $fromId,
                'from_index' => $fromIndex,
                'from_own'   => $fromOwn,
                'from_settings' => $fromSettings + ['embed_store' => $fromId],
                'to'         => $to->id(),
                'to_own'     => $to->isSiteOwn(),
                'to_index'   => $to->indexName(),
                'serve'      => 'from',
                'started'    => time(),
                'vectors'    => $vectors,
                'copied'     => 0,
                'requeued'   => 0,
                'last_error' => '',
            ],
            'moved_from' => null,
        ]);
        Recorder::record([
            'tool'    => 'embeddings',
            'action'  => 'embed_move_start',
            'status'  => 'info',
            'message' => sprintf(
                /* translators: 1: number of vectors, 2: old store, 3: new store */
                __('Moving %1$s vectors from %2$s to %3$s. No neurons: they are copied.', 'aipcf-ai-provider-for-cloudflare'),
                number_format_i18n($vectors),
                self::label($fromId),
                self::label($to->id())
            ),
        ]);
        delete_transient('aipcf_embed_store_status');
    }

    /**
     * Cancels a move: the settings go back to the old store and what was copied is copied back.
     */
    public static function cancel(): bool
    {
        $move = State::get()['move'];
        if ($move === null) {
            return false;
        }
        if (($move['serve'] ?? 'from') === 'to') {
            return true; // already being turned around
        }
        $settings = get_option(AIPCF_SETTINGS_OPTION, []);
        $settings = is_array($settings) ? $settings : [];
        foreach ((array) ($move['from_settings'] ?? ['embed_store' => $move['from']]) as $key => $value) {
            if (in_array($key, self::STORE_KEYS, true)) {
                $settings[$key] = $value;
            }
        }
        update_option(AIPCF_SETTINGS_OPTION, $settings); // afterSave() turns the move around
        return State::get()['move'] === null || (State::get()['move']['serve'] ?? '') === 'to';
    }

    /**
     * @param array<string, mixed> $move
     */
    private static function turnAround(array $move): void
    {
        $back = [
            'from'       => (string) $move['to'],
            'from_index' => (string) $move['to_index'],
            'from_own'   => (bool) ($move['to_own'] ?? false),
            'to'         => (string) $move['from'],
            'to_own'     => (bool) ($move['from_own'] ?? false),
            'to_index'   => (string) $move['from_index'],
            'serve'      => 'to', // readers stay on the store the move started from
            'started'    => time(),
            'vectors'    => (int) $move['copied'],
            'copied'     => 0,
            'requeued'   => 0,
            'last_error' => '',
        ];
        State::update(['move' => $back]);
        Recorder::record([
            'tool'    => 'embeddings',
            'action'  => 'embed_move_cancel',
            'status'  => 'info',
            'message' => sprintf(
                /* translators: %s: the store the move started from */
                __('Move cancelled. What was copied goes back to %s; nothing changes for related posts or search.', 'aipcf-ai-provider-for-cloudflare'),
                self::label((string) $move['from'])
            ),
        ]);
    }

    /**
     * @phpstan-impure
     */
    public static function active(): bool
    {
        return State::get()['move'] !== null;
    }

    /**
     * Copies batches until the deadline. Called by the worker with the new store ready.
     *
     * @return array{copied: int, requeued: int, done: bool, message: string}
     */
    public static function step(VectorStore $to, float $deadline): array
    {
        $out  = ['copied' => 0, 'requeued' => 0, 'done' => false, 'message' => ''];
        $move = State::get()['move'];
        if ($move === null) {
            return $out;
        }
        $from = StoreFactory::make((string) $move['from'], (string) $move['from_index']);
        if (is_wp_error($from)) {
            $out['message'] = $from->get_error_message();
            return $out;
        }
        $oldSig    = Settings::signatureOf((string) $move['from'], (string) $move['from_index']);
        $newSig    = Settings::signature();
        $maxTokens = (int) (Models::get(Models::current())['max_tokens'] ?? 512);

        while (microtime(true) < $deadline) {
            $rows = IndexTable::indexedWithSignature($oldSig, self::BATCH);
            if ($rows === []) {
                $out['done'] = true;
                break;
            }
            $r = self::copyBatch($rows, $from, $to, $oldSig, $newSig, $maxTokens);
            if (is_wp_error($r)) {
                $out['message'] = $r->get_error_message();
                State::update(['move' => ['last_error' => $out['message']] + $move]);
                break;
            }
            $out['copied']   += $r['copied'];
            $out['requeued'] += $r['requeued'];
        }

        $move             = State::get()['move'] ?? $move;
        $move['copied']   = (int) $move['copied'] + $out['copied'];
        $move['requeued'] = (int) $move['requeued'] + $out['requeued'];
        if ($out['copied'] + $out['requeued'] > 0) {
            $move['last_error'] = '';
        }
        State::update(['move' => $move]);
        if ($out['copied'] + $out['requeued'] > 0) {
            Recorder::record([
                'tool'         => 'embeddings',
                'action'       => 'embed_move_batch',
                'status'       => 'ok',
                'message'      => sprintf(
                    /* translators: 1: vectors copied, 2: posts to embed again */
                    __('Copied %1$s vectors; %2$s posts to embed again (edited since, or missing in the old store).', 'aipcf-ai-provider-for-cloudflare'),
                    number_format_i18n($out['copied']),
                    number_format_i18n($out['requeued'])
                ),
                'items_ok'     => $out['copied'],
                'items_failed' => 0,
            ]);
        }
        if ($out['done']) {
            self::finish($move, $to);
        }
        return $out;
    }

    /**
     * @param list<array<string, mixed>> $rows
     * @return array{copied: int, requeued: int}|\WP_Error
     */
    private static function copyBatch(array $rows, VectorStore $from, VectorStore $to, string $oldSig, string $newSig, int $maxTokens)
    {
        $plans = [];
        $ids   = [];
        $out   = ['copied' => 0, 'requeued' => 0];
        foreach ($rows as $row) {
            $type = (string) $row['object_type'];
            $id   = (int) $row['object_id'];
            $post = get_post($id);
            if (!$post instanceof \WP_Post || !Selection::isEligible($post)) {
                IndexTable::update($type, $id, ['status' => IndexTable::PENDING]); // the worker removes it
                $out['requeued']++;
                continue;
            }
            $doc    = Extractor::extract($post);
            $chunks = Chunker::chunk($doc['title'], $doc['body'], $maxTokens);
            // Same text as when it was embedded (the hash holds the old store's signature)?
            if ($chunks === [] || count($chunks) !== (int) $row['chunk_count'] || self::hash($oldSig, $doc) !== $row['content_hash']) {
                IndexTable::update($type, $id, ['status' => IndexTable::PENDING]);
                $out['requeued']++;
                continue;
            }
            $vids = [];
            foreach (array_keys($chunks) as $n) {
                $vids[] = Ingest::vectorId($type, $id, $n);
            }
            $plans[] = ['row' => $row, 'doc' => $doc, 'chunks' => $chunks, 'vids' => $vids];
            $ids     = array_merge($ids, $vids);
        }
        if ($plans === []) {
            return $out;
        }

        $values = $from->vectors($ids);
        if (is_wp_error($values)) {
            return $values;
        }

        $write = [];
        $ready = [];
        foreach ($plans as $p) {
            $missing = array_diff($p['vids'], array_keys($values));
            if ($missing !== []) {
                IndexTable::update((string) $p['row']['object_type'], (int) $p['row']['object_id'], ['status' => IndexTable::PENDING]);
                $out['requeued']++;
                continue;
            }
            foreach ($p['chunks'] as $n => $chunk) {
                $write[] = ['id' => $p['vids'][$n], 'values' => $values[$p['vids'][$n]], 'metadata' => Ingest::metadataFor($p['doc'], $n, $chunk)];
            }
            $ready[] = $p;
        }
        if ($write !== []) {
            $r = $to->upsert($write);
            if (is_wp_error($r)) {
                return $r;
            }
        }
        foreach ($ready as $p) {
            $row  = $p['row'];
            $done = IndexTable::updateIfUnchanged((string) $row['object_type'], (int) $row['object_id'], [
                'signature'    => $newSig,
                'content_hash' => self::hash($newSig, $p['doc']),
            ], (int) $row['revision']);
            if ($done) {
                $out['copied'] += count($p['vids']);
            }
        }
        return $out;
    }

    /**
     * The content hash the worker stores for a post (Ingest::prepare), under a signature.
     *
     * @param array{title: string, body: string, metadata: array<string, mixed>} $doc
     */
    public static function hash(string $signature, array $doc): string
    {
        return sha1(implode('|', [$signature, Settings::storeText() ? '1' : '0', $doc['title'], $doc['body'], (string) wp_json_encode($doc['metadata'])]));
    }

    /**
     * Everything is copied (and Vectorize can read it back): readers move to the new store.
     *
     * @param array<string, mixed> $move
     */
    private static function finish(array $move, VectorStore $to): void
    {
        $status = $to->status();
        if ($to->id() === 'vectorize' && (is_wp_error($status) || empty($status['settled']))) {
            return; // Vectorize shows new vectors after about a minute: switch on the next run
        }
        $back = ($move['serve'] ?? 'from') === 'to';
        // What is left there: the store's own count when the index is this site's alone (a
        // cancelled move may have copied nothing, yet the store holds an earlier move's vectors).
        $left = (int) $move['vectors'];
        if (!empty($move['from_own'])) {
            $from   = StoreFactory::make((string) $move['from'], (string) $move['from_index']);
            $status = is_wp_error($from) ? $from : $from->status();
            if (!is_wp_error($status) && $status['vector_count'] !== null) {
                $left = (int) $status['vector_count'];
            }
        }
        State::update([
            'move'       => null,
            // After a move, the old store still holds the vectors. After a cancelled move,
            // the store it was going to holds what was copied before the cancel.
            'moved_from' => [
                'store'   => (string) $move['from'],
                'index'   => (string) $move['from_index'],
                'own'     => (bool) ($move['from_own'] ?? false),
                'vectors' => $left,
                'at'      => time(),
            ],
        ]);
        \WordPress\CloudflareWorkersAiProvider\Features\Related::flushCache();
        delete_transient('aipcf_embed_store_status');
        Recorder::record([
            'tool'    => 'embeddings',
            'action'  => $back ? 'embed_move_back' : 'embed_move_done',
            'status'  => 'ok',
            'message' => $back
                /* translators: %s: store name */
                ? sprintf(__('Back on %s, as before the move.', 'aipcf-ai-provider-for-cloudflare'), self::label((string) $move['to']))
                : sprintf(
                    /* translators: 1: number of vectors copied, 2: new store, 3: old store */
                    __('Moved: %1$s vectors copied to %2$s. Related posts and search use it from now on. The vectors in %3$s can be deleted on Settings > Content index.', 'aipcf-ai-provider-for-cloudflare'),
                    number_format_i18n((int) $move['copied']),
                    self::label((string) $move['to']),
                    self::label((string) $move['from'])
                ),
            'items_ok' => (int) $move['copied'],
        ]);
    }

    /**
     * Deletes what a move left in the store it came from: the whole index when it is this
     * site's own, or only this site's vectors from a shared one.
     *
     * @return int|\WP_Error Vectors deleted (an estimate when a whole index is dropped).
     */
    public static function cleanup(?string $storeId = null, ?string $index = null, bool $drop = false)
    {
        $state = State::get();
        $left  = $state['moved_from'];
        if ($storeId === null) {
            if ($left === null) {
                return new \WP_Error('aipcf_nothing', __('No store is left over from a move.', 'aipcf-ai-provider-for-cloudflare'));
            }
            $storeId = (string) $left['store'];
            $index   = (string) $left['index'];
            $drop    = !empty($left['own']);
        }
        $current = StoreFactory::current();
        if (!is_wp_error($current) && $current->id() === $storeId && $current->indexName() === (string) $index) {
            return new \WP_Error('aipcf_in_use', __('That is the store Embeddings use now; it is not deleted.', 'aipcf-ai-provider-for-cloudflare'));
        }
        $store = StoreFactory::make($storeId, $index);
        if (is_wp_error($store)) {
            return $store;
        }
        $vectors = (int) ($left['vectors'] ?? 0);
        if ($drop) {
            $r = $store->dropIndex();
        } else {
            // A shared index: only this site's vector ids, as far as the index table knows them.
            $r    = true;
            $page = 0;
            $n    = 0;
            while ($r === true && ($rows = IndexTable::page($page, 500)) !== []) {
                $ids = [];
                foreach ($rows as $row) {
                    $page = (int) $row['id'];
                    for ($c = 0; $c < max(1, (int) $row['chunk_count']); $c++) {
                        $ids[] = Ingest::vectorId((string) $row['object_type'], (int) $row['object_id'], $c);
                    }
                }
                $n += count($ids);
                $r  = $store->delete($ids);
            }
            $vectors = $vectors > 0 ? $vectors : $n;
        }
        if (is_wp_error($r)) {
            return $r;
        }
        State::update(['moved_from' => null]);
        Recorder::record([
            'tool'    => 'embeddings',
            'action'  => 'embed_cleanup_old',
            'status'  => 'ok',
            'message' => sprintf(
                /* translators: 1: store name, 2: index name */
                __('Deleted what the move left in %1$s (%2$s).', 'aipcf-ai-provider-for-cloudflare'),
                self::label($storeId),
                (string) $index
            ),
        ]);
        return $vectors;
    }

    /**
     * Keeps what a move left in the old store and stops offering to delete it.
     */
    public static function keepOld(): void
    {
        State::update(['moved_from' => null]);
    }

    /**
     * Which server a store type points at in the saved settings: the Qdrant URL, the PostgreSQL
     * host, port and database, or the Cloudflare account.
     */
    private static function location(string $storeId): string
    {
        $s = \WordPress\CloudflareWorkersAiProvider\get_plugin_settings();
        if ($storeId === 'qdrant') {
            return rtrim((string) ($s['qdrant_url'] ?? ''), '/');
        }
        if ($storeId === 'pgvector') {
            return (string) ($s['pg_host'] ?? '') . ':' . (int) ($s['pg_port'] ?? 5432) . '/' . (string) ($s['pg_database'] ?? '');
        }
        return \WordPress\CloudflareWorkersAiProvider\get_account_id();
    }

    public static function label(string $storeId): string
    {
        return StoreFactory::available()[$storeId] ?? $storeId;
    }
}