<?php

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

/**
 * The training data as stored: the sources table (one row per trained post,
 * PDF or note) and the chunks table (each source's text in chunks, with their
 * embeddings). Everything that writes training data, or asks what's trained,
 * goes through here.
 *
 * Posts and PDFs are found by their post ID; notes, which aren't posts, by
 * their row's ID (the source ID).
 */
class RSVL_Index
{
    /* ---------------------------------------------------------------------
     * Writing
     * ------------------------------------------------------------------- */

    /**
     * Store a post's embedded chunks in place of what it had, with a fresh
     * source row. Nothing is replaced unless there's something to store, so a
     * failed embed never wipes good data.
     *
     * The delete and inserts are one transaction, and the post's source row
     * is locked first, so another job storing the same post at the same moment
     * (auto-retrain beside a training run, say) waits for this one, then
     * replaces it cleanly. Otherwise it could delete the new source row between
     * the two inserts, leaving chunks that point at nothing, which search still
     * finds and nothing ever deletes.
     *
     * @param array  $item    A prepared item (see RSVL_Indexer::prepare_post()).
     * @param array  $vectors The chunks' embeddings, keyed by chunk position.
     * @param string $model   The embedding model.
     */
    public static function replace_post(array $item, array $vectors, string $model): void
    {
        if (! array_filter($vectors)) {
            return;
        }

        // Two jobs that both lock the post before either stores it deadlock,
        // and MySQL cancels one: that one tries again, once the other is done.
        for ($attempt = 1; $attempt <= 2; $attempt++) {
            $error = self::store_post($item, $vectors, $model);
            if (null === $error) {
                return;
            }
        }

        resolver_log('Could not store the training for post ' . $item['post_id'] . ' (' . $error . ').');
    }

    /**
     * Store a source's chunks, with their embeddings, in place of the ones it
     * had, in a single INSERT. A chunk without an embedding is left out.
     *
     * @param string[] $chunks  The chunks' text.
     * @param array    $vectors Their embeddings, keyed by chunk position.
     */
    public static function replace_chunks(int $source_id, string $source_url, string $content_type, array $chunks, array $vectors, string $model): void
    {
        global $wpdb;

        $wpdb->delete(RSVL_Schema::chunks_table(), ['source_id' => $source_id], ['%d']);

        $rows   = [];
        $values = [];
        foreach (array_values($chunks) as $index => $text) {
            if (empty($vectors[ $index ])) {
                continue;
            }
            $rows[] = '(%d, %s, %d, %d, %s, %s, %s, %s)';
            array_push($values, $source_id, $source_url, $index, count($chunks), $text, RSVL_Embedder::pack_vector($vectors[ $index ]), $model, $content_type);
        }

        if (empty($rows)) {
            return;
        }

        $sql = 'INSERT INTO %i (source_id, source_url, chunk_index, total_chunks, content, embedding, embedding_model, content_type) VALUES '
             . implode(', ', $rows);

        $wpdb->query($wpdb->prepare($sql, RSVL_Schema::chunks_table(), ...$values)); // phpcs:ignore WordPress.DB.PreparedSQL.NotPrepared, PluginCheck.Security.DirectDB.UnescapedDBParameter -- one placeholder group per chunk, built above.
    }

    /**
     * Remove a post or PDF from the training data: its source row and chunks.
     * Its row is locked as it's read (FOR UPDATE), so inside replace_post()'s
     * transaction, a job storing the same post waits here for the other.
     */
    public static function delete_post(int $post_id): void
    {
        global $wpdb;

        $source_id = $wpdb->get_var($wpdb->prepare('SELECT id FROM %i WHERE post_id = %d FOR UPDATE', RSVL_Schema::sources_table(), $post_id));
        if ($source_id) {
            self::delete_source((int) $source_id);
        }
    }

    /** Remove a source, such as a note, and its chunks. */
    public static function delete_source(int $source_id): void
    {
        global $wpdb;

        $wpdb->delete(RSVL_Schema::chunks_table(), ['source_id' => $source_id], ['%d']);
        $wpdb->delete(RSVL_Schema::sources_table(), ['id' => $source_id], ['%d']);
    }

    /**
     * Remove the trained posts (not PDFs or notes) a callback picks, such as
     * those that stopped being trainable when a setting changed.
     *
     * @param callable $should_delete function( WP_Post $post ): bool
     */
    public static function delete_posts_where(callable $should_delete): void
    {
        global $wpdb;

        $post_ids = $wpdb->get_col($wpdb->prepare("SELECT post_id FROM %i WHERE post_id IS NOT NULL AND post_type != 'rsvl-pdf'", RSVL_Schema::sources_table()));

        foreach ($post_ids as $post_id) {
            $post = get_post((int) $post_id);
            if ($post && $should_delete($post)) {
                self::delete_post((int) $post_id);
            }
        }
    }

    /**
     * Delete chunks whose source row is gone, which search would otherwise
     * keep finding and nothing else would ever delete. replace_post() doesn't
     * leave them any more; this clears any an interrupted write left. Runs
     * with the daily cleanup.
     *
     * @return int How many were deleted.
     */
    public static function delete_orphaned_chunks(): int
    {
        global $wpdb;

        return (int) $wpdb->query($wpdb->prepare(
            'DELETE c FROM %i c
             LEFT JOIN %i s ON s.id = c.source_id
             WHERE s.id IS NULL',
            RSVL_Schema::chunks_table(),
            RSVL_Schema::sources_table()
        ));
    }

    /** Record why a trained post failed to retrain. Posts never trained have no row, and get none. */
    public static function set_training_error(int $post_id, string $error): void
    {
        global $wpdb;

        $wpdb->update(
            RSVL_Schema::sources_table(),
            ['training_error' => resolver_sanitize_plain_text($error)],
            ['post_id'        => $post_id],
            ['%s'],
            ['%d']
        );
    }

    /** Unlink notes from chats that are being deleted. */
    public static function unlink_chats(array $session_ids): void
    {
        global $wpdb;

        if (empty($session_ids)) {
            return;
        }

        $wpdb->query($wpdb->prepare(
            'UPDATE %i SET chat_session_id = NULL WHERE chat_session_id IN (' . implode(',', array_fill(0, count($session_ids), '%s')) . ')',
            array_merge([RSVL_Schema::sources_table()], array_values($session_ids))
        ));
    }

    /* ---------------------------------------------------------------------
     * Reading
     * ------------------------------------------------------------------- */

    public static function is_indexed(int $post_id): bool
    {
        global $wpdb;

        return (bool) $wpdb->get_var($wpdb->prepare('SELECT 1 FROM %i WHERE post_id = %d LIMIT 1', RSVL_Schema::sources_table(), $post_id));
    }

    /** Whether anything is stored at all, trained or failed. */
    public static function has_sources(): bool
    {
        global $wpdb;

        return (bool) $wpdb->get_var($wpdb->prepare('SELECT 1 FROM %i LIMIT 1', RSVL_Schema::sources_table()));
    }

    /** Whether anything trained successfully. */
    public static function has_trained_content(): bool
    {
        global $wpdb;

        return (bool) $wpdb->get_var($wpdb->prepare('SELECT 1 FROM %i WHERE training_error IS NULL LIMIT 1', RSVL_Schema::sources_table()));
    }

    /**
     * Every trained post's and PDF's ID.
     *
     * @return int[]
     */
    public static function get_post_ids(): array
    {
        global $wpdb;

        return array_map('intval', $wpdb->get_col($wpdb->prepare('SELECT post_id FROM %i WHERE post_id IS NOT NULL', RSVL_Schema::sources_table())));
    }

    /**
     * Trained posts that have any of these meta keys, e.g. for retraining the
     * pages a custom field appears on.
     *
     * @return int[]
     */
    public static function get_post_ids_with_meta(array $meta_keys): array
    {
        global $wpdb;

        if (empty($meta_keys)) {
            return [];
        }

        return array_map('intval', $wpdb->get_col($wpdb->prepare(
            'SELECT DISTINCT s.post_id FROM %i s
             INNER JOIN %i pm ON s.post_id = pm.post_id
             WHERE pm.meta_key IN (' . implode(',', array_fill(0, count($meta_keys), '%s')) . ') AND s.post_id IS NOT NULL',
            array_merge([RSVL_Schema::sources_table(), $wpdb->postmeta], array_values($meta_keys))
        )));
    }

    /**
     * Whether these posts are trained, and how.
     *
     * @param int[] $post_ids
     * @return array Keyed by post ID: ['indexed', 'indexed_at', 'char_count', 'found_fields', 'training_error'].
     */
    public static function get_post_data(array $post_ids): array
    {
        global $wpdb;

        $post_ids = array_values(array_filter(array_map('absint', $post_ids)));
        if (empty($post_ids)) {
            return [];
        }

        $rows = $wpdb->get_results($wpdb->prepare(
            'SELECT post_id, indexed_at, char_count, found_fields, training_error FROM %i WHERE post_id IN (' . implode(',', array_fill(0, count($post_ids), '%d')) . ')',
            array_merge([RSVL_Schema::sources_table()], $post_ids)
        ));

        $result = [];
        foreach ($rows as $row) {
            $result[ (int) $row->post_id ] = [
                'indexed'        => true,
                'indexed_at'     => $row->indexed_at,
                'char_count'     => (int) $row->char_count,
                'found_fields'   => $row->found_fields ? json_decode($row->found_fields, true) : [],
                'training_error' => $row->training_error ?? null,
            ];
        }

        foreach ($post_ids as $post_id) {
            $result[ $post_id ] ??= [
                'indexed'        => false,
                'indexed_at'     => null,
                'char_count'     => 0,
                'found_fields'   => [],
                'training_error' => null,
            ];
        }

        return $result;
    }

    /**
     * Every trained post (not PDFs or notes, which have their own panels),
     * grouped by post type, for the Pages table. Posts since deleted are left
     * out.
     *
     * @return array Keyed by post type: ['label' => …, 'posts' => [[post_id, title, url, type_slug, type_label, indexed_at, char_count, found_fields, training_error], …]].
     */
    public static function get_posts_by_type(): array
    {
        global $wpdb;

        $rows = $wpdb->get_results($wpdb->prepare(
            // title is deliberately not selected: it is only populated for notes
            // (see the column comment in the schema). Every row here has a real
            // WordPress post, so the title comes from $post->post_title below.
            "SELECT id, post_id, post_type AS content_type, indexed_at, char_count, found_fields, training_error
             FROM %i
             WHERE post_type NOT IN ( 'rsvl-note', 'rsvl-pdf' )
             ORDER BY post_type",
            RSVL_Schema::sources_table()
        ));

        $grouped = [];
        foreach ($rows ?: [] as $row) {
            $type = $row->content_type;
            $post = get_post((int) $row->post_id);
            if (! $post) {
                continue;
            }

            $pt_obj = get_post_type_object($type);
            if (! isset($grouped[ $type ])) {
                $grouped[ $type ] = [
                    'label' => $pt_obj ? $pt_obj->labels->name : ucfirst($type),
                    'posts' => [],
                ];
            }

            $grouped[ $type ]['posts'][] = [
                'post_id' => $post->ID,
                'title'   => $post->post_title,
                // The editor, not the permalink: from this table the useful next
                // step is fixing the content that was trained. Falls back to the
                // permalink for anyone without edit rights on the post.
                'url'            => get_edit_post_link($post->ID) ?: (string) get_permalink($post->ID),
                'type_slug'      => $type,
                'type_label'     => $pt_obj ? $pt_obj->labels->singular_name : ucfirst($type),
                'indexed_at'     => $row->indexed_at,
                'char_count'     => (int) $row->char_count,
                'found_fields'   => $row->found_fields ? json_decode($row->found_fields, true) : [],
                'training_error' => $row->training_error ?? null,
            ];
        }

        foreach ($grouped as &$group) {
            usort($group['posts'], fn ($a, $b) => strcmp($a['title'], $b['title']));
        }
        unset($group);

        uasort($grouped, fn ($a, $b) => strcmp($a['label'], $b['label']));

        return $grouped;
    }

    /**
     * The most recently trained posts (not PDFs or notes), newest first.
     * Posts since deleted are left out.
     *
     * @return array[] Each: ['post' => WP_Post, 'indexed_at' => …].
     */
    public static function get_recent_posts(int $limit): array
    {
        global $wpdb;

        // Over-fetch rather than LIMIT to $limit directly: rows whose post has
        // since been deleted still have to be dropped in PHP below, and doing
        // that after a tight LIMIT silently returns a short list.
        $rows = $wpdb->get_results($wpdb->prepare(
            "SELECT post_id, indexed_at FROM %i
             WHERE post_id IS NOT NULL AND post_type NOT IN ( 'rsvl-note', 'rsvl-pdf' )
             ORDER BY indexed_at DESC, id DESC
             LIMIT %d",
            RSVL_Schema::sources_table(),
            $limit * 3
        ), ARRAY_A) ?: [];

        $posts = [];
        foreach ($rows as $row) {
            $post = get_post((int) $row['post_id']);
            if ($post) {
                $posts[] = ['post' => $post, 'indexed_at' => $row['indexed_at']];
            }
            if (count($posts) >= $limit) {
                break;
            }
        }

        return $posts;
    }

    /** Counts of what's stored. */
    public static function get_stats(): array
    {
        global $wpdb;

        $table = RSVL_Schema::chunks_table();

        return [
            'total_chunks' => (int) $wpdb->get_var($wpdb->prepare('SELECT COUNT(*) FROM %i', $table)),
            'total_posts'  => (int) $wpdb->get_var($wpdb->prepare('SELECT COUNT(DISTINCT source_id) FROM %i WHERE source_id IS NOT NULL', $table)),
            'by_type'      => $wpdb->get_results($wpdb->prepare('SELECT content_type, COUNT(*) as chunk_count, COUNT(DISTINCT source_id) as post_count FROM %i GROUP BY content_type', $table), ARRAY_A) ?: [],
        ];
    }

    /**
     * One try at replace_post(), as a transaction.
     *
     * @return string|null Why it failed, or null when it was stored.
     */
    private static function store_post(array $item, array $vectors, string $model): ?string
    {
        global $wpdb;

        $wpdb->query('START TRANSACTION');

        self::delete_post($item['post_id']);

        $found    = $item['found_fields'] ?? [];
        $inserted = $wpdb->insert(
            RSVL_Schema::sources_table(),
            [
                'post_id'      => $item['post_id'],
                'post_type'    => $item['post_type'],
                'indexed_at'   => gmdate('Y-m-d H:i:s'),
                'char_count'   => array_sum(array_map('strlen', $item['chunks'])),
                'found_fields' => ! empty($found) ? wp_json_encode($found) : null,
            ],
            ['%d', '%s', '%s', '%d', '%s']
        );

        // A failed insert leaves $wpdb->insert_id pointing at an unrelated
        // earlier insert, so its chunks must never be stored against it.
        if (false === $inserted) {
            $error = $wpdb->last_error;
            $wpdb->query('ROLLBACK');

            return $error;
        }

        self::replace_chunks((int) $wpdb->insert_id, $item['source_url'], $item['post_type'], $item['chunks'], $vectors, $model);

        $wpdb->query('COMMIT');

        return null;
    }
}